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

add source prefix

Browse Source
This commit is contained in:
hiyouga
2023-06-16 16:32:17 +08:00
parent 0574b590ef
commit fc4d8155b3
3 changed files with 9 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/cli_demo.py
View File

@@ -20,9 +20,10 @@ def main():
model_name = "BLOOM" if "bloom" in model_args.model_name_or_path else "LLaMA"
prompt_template = Template(data_args.prompt_template)
source_prefix = data_args.source_prefix if data_args.source_prefix else ""
def predict_and_print(query, history: list) -> list:
input_ids = tokenizer([prompt_template.get_prompt(query, history)], return_tensors="pt")["input_ids"]
input_ids = tokenizer([prompt_template.get_prompt(query, history, source_prefix)], 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)