mirror of
https://github.com/hiyouga/LLaMA-Factory.git
synced 2025-10-15 08:08:09 +08:00
76 lines
2.6 KiB
Python
76 lines
2.6 KiB
Python
# coding=utf-8
|
||
# Implements stream chat in command line for fine-tuned models.
|
||
# Usage: python cli_demo.py --checkpoint_dir path_to_checkpoint
|
||
|
||
|
||
import torch
|
||
from utils import ModelArguments, FinetuningArguments, load_pretrained, get_logits_processor
|
||
from transformers import HfArgumentParser
|
||
|
||
|
||
def main():
|
||
|
||
parser = HfArgumentParser((ModelArguments, FinetuningArguments))
|
||
model_args, finetuning_args = parser.parse_args_into_dataclasses()
|
||
model_name = "BLOOM" if "bloom" in model_args.model_name_or_path else "LLaMA"
|
||
model, tokenizer = load_pretrained(model_args, finetuning_args)
|
||
|
||
if torch.cuda.device_count() > 1:
|
||
from accelerate import dispatch_model, infer_auto_device_map
|
||
device_map = infer_auto_device_map(model)
|
||
model = dispatch_model(model, device_map)
|
||
else:
|
||
model = model.cuda()
|
||
|
||
model.eval()
|
||
|
||
def format_example(query):
|
||
prompt = "Below is an instruction that describes a task. "
|
||
prompt += "Write a response that appropriately completes the request.\n"
|
||
prompt += "Instruction:\nHuman: {}\nAssistant: ".format(query)
|
||
return prompt
|
||
|
||
def predict(query, history: list):
|
||
input_ids = tokenizer([format_example(query)], return_tensors="pt")["input_ids"]
|
||
input_ids = input_ids.to(model.device)
|
||
gen_kwargs = {
|
||
"do_sample": True,
|
||
"top_p": 0.7,
|
||
"temperature": 0.95,
|
||
"num_beams": 1,
|
||
"max_new_tokens": 256,
|
||
"repetition_penalty": 1.5,
|
||
"logits_processor": get_logits_processor()
|
||
}
|
||
with torch.no_grad():
|
||
generation_output = model.generate(input_ids=input_ids, **gen_kwargs)
|
||
outputs = generation_output.tolist()[0][len(input_ids[0]):]
|
||
response = tokenizer.decode(outputs, skip_special_tokens=True)
|
||
history = history + [(query, response)]
|
||
return response, history
|
||
|
||
history = []
|
||
print("欢迎使用 {} 模型，输入内容即可对话，clear清空对话历史，stop终止程序".format(model_name))
|
||
while True:
|
||
try:
|
||
query = input("\nInput: ")
|
||
except UnicodeDecodeError:
|
||
print("Detected decoding error at the inputs, please set the terminal encoding to utf-8.")
|
||
continue
|
||
except Exception:
|
||
raise
|
||
|
||
if query.strip() == "stop":
|
||
break
|
||
|
||
if query.strip() == "clear":
|
||
history = []
|
||
continue
|
||
|
||
response, history = predict(query, history)
|
||
print("{}:".format(model_name), response)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|