423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-08-23 14:22:51 +08:00
Code Issues Packages Projects Releases Wiki Activity

add source prefix

Browse Source 
Former-commit-id: fc4d8155b35dcc453a64a50b21ce59050a15be99
This commit is contained in:
hiyouga 2023-06-16 16:32:17 +08:00
parent c6d56e7109
commit 36ea46e85c
3 changed files with 9 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
src/api_demo.py
View File

@ -42,7 +42,7 @@ app = FastAPI()
@app.post("/") @app.post("/")
async def create_item(request: Request): async def create_item(request: Request):
global model, tokenizer, prompt_template, generating_args global model, tokenizer, prompt_template, source_prefix, generating_args
# Parse the request JSON # Parse the request JSON
json_post_raw = await request.json() json_post_raw = await request.json()
@ -55,7 +55,7 @@ async def create_item(request: Request):
temperature = json_post_list.get("temperature", None) temperature = json_post_list.get("temperature", None)
# Tokenize the input prompt # Tokenize the input prompt
input_ids = tokenizer([prompt_template.get_prompt(prompt, history)], return_tensors="pt")["input_ids"] input_ids = tokenizer([prompt_template.get_prompt(prompt, history, source_prefix)], return_tensors="pt")["input_ids"]
input_ids = input_ids.to(model.device) input_ids = input_ids.to(model.device)
# Generation arguments # Generation arguments
@ -94,8 +94,11 @@ async def create_item(request: Request):
if __name__ == "__main__": if __name__ == "__main__":
model_args, data_args, finetuning_args, generating_args = prepare_infer_args() model_args, data_args, finetuning_args, generating_args = prepare_infer_args()
model, tokenizer = load_pretrained(model_args, finetuning_args) model, tokenizer = load_pretrained(model_args, finetuning_args)
prompt_template = Template(data_args.prompt_template) prompt_template = Template(data_args.prompt_template)
source_prefix = data_args.source_prefix if data_args.source_prefix else ""
uvicorn.run(app, host='0.0.0.0', port=8000, workers=1) uvicorn.run(app, host='0.0.0.0', port=8000, workers=1)

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" model_name = "BLOOM" if "bloom" in model_args.model_name_or_path else "LLaMA"
prompt_template = Template(data_args.prompt_template) 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: 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) input_ids = input_ids.to(model.device)
streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True) streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)

3
src/web_demo.py
View File

@ -25,6 +25,7 @@ model_args, data_args, finetuning_args, generating_args = prepare_infer_args()
model, tokenizer = load_pretrained(model_args, finetuning_args) model, tokenizer = load_pretrained(model_args, finetuning_args)
prompt_template = Template(data_args.prompt_template) prompt_template = Template(data_args.prompt_template)
source_prefix = data_args.source_prefix if data_args.source_prefix else ""
def postprocess(self, y): def postprocess(self, y):
@ -79,7 +80,7 @@ def parse_text(text): # copy from https://github.com/GaiZhenbiao/ChuanhuChatGPT
def predict(query, chatbot, max_length, top_p, temperature, history): def predict(query, chatbot, max_length, top_p, temperature, history):
chatbot.append((parse_text(query), "")) chatbot.append((parse_text(query), ""))
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) input_ids = input_ids.to(model.device)
streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True) streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)