Added Abacus.ai Smaug and Snowflake instruct to the list of ChatML-ba…

…sed models.
fairydreaming · May 3, 2024 · 8b3e09d · 8b3e09d
1 parent a9a0013
commit 8b3e09d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/run_model.py b/run_model.py
@@ -44,7 +44,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"]):
+elif any(model_name in model_file_basename.lower() for model_name in ["qwen", "yi", "dbrx-instruct", "theprofessor", "smaug", "arctic"]):
     prompt_template = CHATML_PROMPT_TEMPLATE
 elif any(model_name in model_file_basename.lower() for model_name in ["command-r"]):
     prompt_template = COMMANDR_PROMPT_TEMPLATE