Skip to content

Commit

Permalink
Added OLMo-2 model template (same as Tulu-3 model).
Browse files Browse the repository at this point in the history
Added BOS token to the Tulu-3 model template.
  • Loading branch information
sszymczy committed Nov 27, 2024
1 parent 5e5afdf commit ce06b03
Showing 1 changed file with 3 additions and 3 deletions.
6 changes: 3 additions & 3 deletions run_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@
DEEPSEEK2_PROMPT_TEMPLATE="{SYS}\n\nUser: {USER}\n\nAssistant:"
NEMOTRON4_PROMPT_TEMPLATE="<extra_id_0>System\n{SYS}\n<extra_id_1>User\n{USER}\n<extra_id_1>Assistant\n"
DEEPSEEK2_0628_PROMPT_TEMPLATE="<|begin▁of▁sentence|>{SYS}\n\n<|User|>{USER}<|Assistant|>"
TULU3_PROMPT_TEMPLATE="<|system|>\n{SYS}\n<|user|>\n{USER}\n<|assistant|>\n"
TULU3_PROMPT_TEMPLATE="<|endoftext|><|system|>\n{SYS}\n<|user|>\n{USER}\n<|assistant|>\n"
else:
LLAMA_PROMPT_TEMPLATE="<s>[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"
Expand All @@ -44,7 +44,7 @@
DEEPSEEK2_PROMPT_TEMPLATE="User: {USER}\n\nAssistant:"
NEMOTRON4_PROMPT_TEMPLATE="<extra_id_0>System\n\n<extra_id_1>User\n{USER}<extra_id_1>Assistant\n"
DEEPSEEK2_0628_PROMPT_TEMPLATE="<|begin▁of▁sentence|><|User|>{USER}<|Assistant|>"
TULU3_PROMPT_TEMPLATE="<|user|>\n{USER}\n<|assistant|>\n"
TULU3_PROMPT_TEMPLATE="<|endoftext|><|user|>\n{USER}\n<|assistant|>\n"

model_file_basename = os.path.basename(model_file)

Expand All @@ -67,7 +67,7 @@
prompt_template = DEEPSEEK2_PROMPT_TEMPLATE
elif any(model_name in model_file_basename.lower() for model_name in ["nemotron"]):
prompt_template = NEMOTRON4_PROMPT_TEMPLATE
elif any(model_name in model_file_basename.lower() for model_name in ["tulu-3"]):
elif any(model_name in model_file_basename.lower() for model_name in ["tulu-3", "olmo-2"]):
prompt_template = TULU3_PROMPT_TEMPLATE
else:
raise RuntimeError("Could not detect model prompt template!")
Expand Down

0 comments on commit ce06b03

Please sign in to comment.