Added InternLM model to the list of models using ChatML template.

fairydreaming · Aug 14, 2024 · 2d2f829 · 2d2f829
1 parent f125445
commit 2d2f829
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/run_model.py b/run_model.py
@@ -50,7 +50,7 @@
     prompt_template = LLAMA3_PROMPT_TEMPLATE
 elif any(model_name in model_file_basename.lower() for model_name in ["llama", "gemma", "mistral", "mixtral", "miqu"]):
     prompt_template = LLAMA_PROMPT_TEMPLATE
-elif any(model_name in model_file_basename.lower() for model_name in ["qwen", "yi", "dbrx-instruct", "theprofessor", "smaug", "arctic"]):
+elif any(model_name in model_file_basename.lower() for model_name in ["qwen", "yi", "dbrx-instruct", "theprofessor", "smaug", "arctic", "internlm"]):
     prompt_template = CHATML_PROMPT_TEMPLATE
 elif any(model_name in model_file_basename.lower() for model_name in ["command-r", "aya-23"]):
     prompt_template = COMMANDR_PROMPT_TEMPLATE