423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-12-16 11:50:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

support multiturn training like FastChat

Browse Source
This commit is contained in:
hiyouga
2023-06-14 22:27:39 +08:00
parent 875e8e2349
commit b6faf0207d
5 changed files with 166 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/cli_demo.py
View File

@@ -16,23 +16,27 @@ from transformers import TextIteratorStreamer
def main():
model_args, data_args, finetuning_args, generating_args = prepare_infer_args()
model_name = "BLOOM" if "bloom" in model_args.model_name_or_path else "LLaMA"
model, tokenizer = load_pretrained(model_args, finetuning_args)
model_name = "BLOOM" if "bloom" in model_args.model_name_or_path else "LLaMA"
prompt_template = Template(data_args.prompt_template)
streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
def predict_and_print(query, history: list):
def predict_and_print(query, history: list) -> list:
input_ids = tokenizer([prompt_template.get_prompt(query, history)], return_tensors="pt")["input_ids"]
input_ids = input_ids.to(model.device)
streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
gen_kwargs = generating_args.to_dict()
gen_kwargs["input_ids"] = input_ids
gen_kwargs["logits_processor"] = get_logits_processor()
gen_kwargs["streamer"] = streamer
thread = Thread(target=model.generate, kwargs=gen_kwargs)
thread.start()
print("{}: ".format(model_name), end="", flush=True)
response = ""
print("{}: ".format(model_name), end="")
for new_text in streamer:
print(new_text, end="", flush=True)
response += new_text