Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2device = "cuda" # the device to load the model onto
3
4model = AutoModelForCausalLM.from_pretrained(
5 "ShengbinYue/LawLLM-7B",
6 torch_dtype="auto",
7 device_map="auto"
8)
9tokenizer = AutoTokenizer.from_pretrained("ShengbinYue/LawLLM-7B")
10
11prompt = "生产销售假冒伪劣商品罪如何判刑?"
12messages = [
13 {"role": "system", "content": "你是LawLLM,一个由复旦大学DISC实验室创造的法律助手。"},
14 {"role": "user", "content": prompt}
15]
16text = tokenizer.apply_chat_template(
17 messages,
18 tokenize=False,
19 add_generation_prompt=True
20)
21model_inputs = tokenizer([text], return_tensors="pt").to(device)
22
23generated_ids = model.generate(
24 model_inputs.input_ids,
25 max_new_tokens=512
26)
27generated_ids = [
28 output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
29]
30
31response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]1from transformers import AutoModelForCausalLM, AutoTokenizer
2from vllm import LLM, SamplingParams
3
4model_name ='ShengbinYue/LawLLM-7B'
5
6sampling_params = SamplingParams(
7 temperature=0.1,
8 top_p=0.9,
9 top_k=50,
10 max_tokens=4096
11)
12llm = LLM(model=model_name)
13
14tokenizer = AutoTokenizer.from_pretrained(model_name)
15prompt = "生产销售假冒伪劣商品罪如何判刑?"
16
17# prompt = "戴罪立功是什么意思"
18messages = [
19 {"role": "system", "content": "你是LawLLM,一个由复旦大学DISC实验室创造的法律助手。"},
20 {"role": "user", "content": prompt}
21]
22text = tokenizer.apply_chat_template(
23 messages,
24 tokenize=False,
25 add_generation_prompt=True
26)
27
28outputs = llm.generate([text], sampling_params)
29for output in outputs:
30 prompt = output.prompt
31 generated_text = output.outputs[0].text
32 print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")@misc{yue2023disclawllm,
title={DISC-LawLLM: Fine-tuning Large Language Models for Intelligent Legal Services},
author={Shengbin Yue and Wei Chen and Siyuan Wang and Bingxuan Li and Chenchen Shen and Shujun Liu and Yuxuan Zhou and Yao Xiao and Song Yun and Wei Lin and Xuanjing Huang and Zhongyu Wei},
year={2023},
eprint={2309.11325},
archivePrefix={arXiv},
primaryClass={cs.CL}
}
@inproceedings{yue2024lawllm,
title={LawLLM: Intelligent Legal System with Legal Reasoning and Verifiable Retrieval},
author={Yue, Shengbin and Liu, Shujun and Zhou, Yuxuan and Shen, Chenchen and Wang, Siyuan and Xiao, Yao and Li, Bingxuan and Song, Yun and Shen, Xiaoyu and Chen, Wei and others},
booktitle={International Conference on Database Systems for Advanced Applications},
pages={304--321},
year={2024},
organization={Springer}
}