diff --git a/run_model.py b/run_model.py index 6103514..765c100 100755 --- a/run_model.py +++ b/run_model.py @@ -29,12 +29,14 @@ CHATML_PROMPT_TEMPLATE="<|im_start|>system\n{SYS}<|im_end|>\n<|im_start|>user\n{USER}<|im_end|>\n<|im_start|>assistant" COMMANDR_PROMPT_TEMPLATE="<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>{SYS}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|USER_TOKEN|>{USER}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>" VICUNA_PROMPT_TEMPLATE="{SYS}\nUSER: {USER}\nASSISTANT:" + PHI3_PROMPT_TEMPLATE=None else: LLAMA_PROMPT_TEMPLATE="[INST] {USER}[/INST]\n" LLAMA3_PROMPT_TEMPLATE="<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{USER}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" CHATML_PROMPT_TEMPLATE="<|im_start|>user\n{USER}<|im_end|>\n<|im_start|>assistant\n" COMMANDR_PROMPT_TEMPLATE="<|START_OF_TURN_TOKEN|><|USER_TOKEN|>{USER}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>" VICUNA_PROMPT_TEMPLATE="USER: {USER}\nASSISTANT: " + PHI3_PROMPT_TEMPLATE="<|user|>\n{USER}<|end|>\n<|assistant|>" model_file_basename = os.path.basename(model_file) @@ -48,6 +50,8 @@ prompt_template = COMMANDR_PROMPT_TEMPLATE elif any(model_name in model_file_basename.lower() for model_name in ["wizardlm"]): prompt_template = VICUNA_PROMPT_TEMPLATE +elif any(model_name in model_file_basename.lower() for model_name in ["phi-3"]): + prompt_template = PHI3_PROMPT_TEMPLATE else: raise RuntimeError("Could not detect model prompt template!")