From 32f565a4b113ee191dc410bd135a16141038591f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= <sszymczy@gmail.com>
Date: Mon, 15 Jul 2024 09:20:20 +0200
Subject: [PATCH] Added Nemotron-4 prompt template.

---
 run_model.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/run_model.py b/run_model.py
index ed4dbb7..a19963f 100755
--- a/run_model.py
+++ b/run_model.py
@@ -31,6 +31,7 @@
     VICUNA_PROMPT_TEMPLATE="{SYS}\nUSER: {USER}\nASSISTANT:"
     PHI3_PROMPT_TEMPLATE=None
     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"
 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"
@@ -39,6 +40,7 @@
     VICUNA_PROMPT_TEMPLATE="USER: {USER}\nASSISTANT: "
     PHI3_PROMPT_TEMPLATE="<|user|>\n{USER}<|end|>\n<|assistant|>"
     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"
 
 model_file_basename = os.path.basename(model_file)
 
@@ -56,6 +58,8 @@
     prompt_template = PHI3_PROMPT_TEMPLATE
 elif any(model_name in model_file_basename.lower() for model_name in ["deepseek-v2"]):
     prompt_template = DEEPSEEK2_PROMPT_TEMPLATE
+elif any(model_name in model_file_basename.lower() for model_name in ["nemotron"]):
+    prompt_template = NEMOTRON4_PROMPT_TEMPLATE
 else:
     raise RuntimeError("Could not detect model prompt template!")