You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
67 lines
2.1 KiB
67 lines
2.1 KiB
2 years ago
|
# coding=utf-8
|
||
|
|
# Implements stream chat in command line for LLaMA fine-tuned with PEFT.
|
||
|
|
# Usage: python cli_demo.py --checkpoint_dir path_to_checkpoint
|
||
|
|
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from utils import ModelArguments, auto_configure_device_map, load_pretrained
|
||
|
|
from transformers import HfArgumentParser
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
|
||
|
|
parser = HfArgumentParser(ModelArguments)
|
||
|
|
model_args, = parser.parse_args_into_dataclasses()
|
||
|
|
model, tokenizer = load_pretrained(model_args)
|
||
|
|
if torch.cuda.device_count() > 1:
|
||
|
|
from accelerate import dispatch_model
|
||
|
|
device_map = auto_configure_device_map(torch.cuda.device_count())
|
||
|
|
model = dispatch_model(model, device_map)
|
||
|
|
else:
|
||
|
|
model = model.cuda()
|
||
|
|
model.eval()
|
||
|
|
|
||
|
|
def predict(query, history: list):
|
||
|
|
inputs = tokenizer([query], return_tensors="pt")
|
||
|
|
inputs = inputs.to(model.device)
|
||
|
|
gen_kwargs = {
|
||
|
|
"do_sample": True,
|
||
|
|
"top_p": 0.9,
|
||
|
|
"top_k": 40,
|
||
|
|
"temperature": 0.7,
|
||
|
|
"num_beams": 1,
|
||
|
|
"max_new_tokens": 256,
|
||
|
|
"repetition_penalty": 1.5
|
||
|
|
}
|
||
|
|
with torch.no_grad():
|
||
|
|
generation_output = model.generate(**inputs, **gen_kwargs)
|
||
|
|
outputs = generation_output.tolist()[0][len(inputs["input_ids"][0]):]
|
||
|
|
response = tokenizer.decode(outputs, skip_special_tokens=True)
|
||
|
|
history = history + [(query, response)]
|
||
|
|
return response, history
|
||
|
|
|
||
|
|
history = []
|
||
|
|
print("欢迎使用 LLaMA-7B 模型，输入内容即可对话，clear清空对话历史，stop终止程序")
|
||
|
|
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("LLaMA-7B:", response)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|