Fix prompts

Signed-off-by: Igor Gitman <[email protected]>

nemo_skills/inference/prompt/llama3/base.yaml

@@ -1,16 +1,14 @@
 few_shot_examples:
-  template: "Question:\n{question}\n\n{context}My solution:\n{generation}\n\n\n\n\n\n"
+  template: "Question:\n{question}{context}\n\nMy solution:\n{generation}\n\n\n\n\n\n"
 
 system: ""
 user: |-
     Here are some examples of questions and solutions followed by a new question that you need to solve.
     Make sure to put the answer (and only answer) inside \boxed{{}}.
 
     {examples}Question:
-    {question}
+    {question}{context}
 
-    {context}
-
-prompt_template: "<|begin_of_text|>{user}My solution:\n{generation}"
+prompt_template: "<|begin_of_text|>{user}\n\nMy solution:\n{generation}"
 
 stop_phrases: ["\n\n\n\n\n\n"]

nemo_skills/inference/prompt/llama3/instruct.yaml

@@ -1,23 +1,22 @@
 few_shot_examples:
-  template: "Question:\n{question}\n\n{context}My solution:\n{generation}\n\n\n\n\n\n"
+  template: "Question:\n{question}{context}\n\nMy solution:\n{generation}\n\n\n\n\n\n"
 
-system: |-
-  You are Meta AI, a sophisticated and energetic AI Assistant. You excel at solving mathematical problems.
+system: ""
 
 user: |-
     Here are some examples of questions and solutions followed by a new question that you need to solve.
     Make sure to put the answer (and only answer) inside \boxed{{}}.
 
     {examples}Question:
-    {question}
-
-    {context}
+    {question}{context}
 
 prompt_template: |-
   <|begin_of_text|><|start_header_id|>system<|end_header_id|>
 
   {system}<|eot_id|><|start_header_id|>user<|end_header_id|>
 
-  {user}<|eot_id|><|start_header_id|>assistant<|end_header_id|>{generation}
+  {user}
+
+  <|eot_id|><|start_header_id|>assistant<|end_header_id|>{generation}
 
 stop_phrases: ["<|eot_id|>"]

nemo_skills/inference/prompt/llama3/sft.yaml

@@ -1,7 +1,3 @@
-few_shot_examples:
-  template: ""
-  num_few_shots: 0
-
 system: ""
 
 user: "{question}{context}"

nemo_skills/inference/prompt/nemotron/fewshot.yaml

@@ -1,5 +1,5 @@
 few_shot_examples:
-  template: "Example question:\n{question}\n\n{context}Example solution:\n{generation}\n\n\n\n\n\n"
+  template: "Example question:\n{question}{context}\n\nExample solution:\n{generation}\n\n\n\n\n\n"
 
 system: ""
 
@@ -8,9 +8,7 @@ user: |-
     Make sure to put the answer (and only answer) inside \boxed{{}}.
 
     {examples}Question:
-    {question}
-
-    {context}
+    {question}{context}
 
     Don't forget that your final answer should be inside \boxed{{}}!
 

nemo_skills/inference/prompt/nemotron/zeroshot.yaml

@@ -1,15 +1,10 @@
-few_shot_examples:
-  template: ""
-  num_few_shots: 0
-
 system: ""
 
 user: |-
     You are an expert in math. Given the math question below, I want you to reason through the steps and then give a final answer.
 
     Your final answer should be inside \boxed{{}}.
 
-    {examples}Question:
     {question}{context}
 
 prompt_template: |-

nemo_skills/inference/prompt/openmathinstruct/base.yaml

@@ -1,14 +1,12 @@
 few_shot_examples:
-  template: "Question:\n{question}\n\n{context}My solution:\n{generation}\n\n\n\n\n\n"
+  template: "Question:\n{question}{context}\n\nMy solution:\n{generation}\n\n\n\n\n\n"
 system: ""
 user: |-
     Here are some examples of questions and solutions followed by a new question that you need to solve.
     Make sure to put the answer (and only answer) inside \boxed{{}}.
 
     {examples}Question:
-    {question}
+    {question}{context}
 
-    {context}
-
-prompt_template: "{system}{user}My solution:\n{generation}"
+prompt_template: "{system}{user}\n\nMy solution:\n{generation}"
 stop_phrases: ["\n\n\n\n\n\n"]

nemo_skills/inference/prompt/openmathinstruct/reference.yaml

@@ -1,14 +1,12 @@
 few_shot_examples:
-  template: "Question:\n{question}\n\n{context}My solution:\n{reference_solution}\n\n\n\n\n\n"
+  template: "Question:\n{question}{context}\n\nMy solution:\n{reference_solution}\n\n\n\n\n\n"
 system: ""
 user: |-
     Here are some examples of questions and solutions followed by a new question that you need to solve.
     Make sure to put the answer (and only answer) inside \boxed{{}}.
 
     {examples}Question:
-    {question}
+    {question}{context}
 
-    {context}
-
-prompt_template: "{system}{user}My solution:\n{generation}"
+prompt_template: "{system}{user}\n\nMy solution:\n{generation}"
 stop_phrases: ["\n\n\n\n\n\n"]

nemo_skills/inference/prompt/openmathinstruct/sft.yaml

@@ -1,13 +1,9 @@
-few_shot_examples:
-  template: ""
-  num_few_shots: 0
-
 system: >-  # no line breaks!
   You're an expert Python programmer and mathematician.
   Help the user to solve this problem using code when necessary.
   Make sure to put the answer (and only answer) inside \boxed{}.
 
-user: "{question}\n\n{context}"
+user: "{question}{context}\n\n"
 
 prompt_template: "System:\n{system}\n\nUser:\n{user}Assistant:\n{generation}"
 stop_phrases: []

nemo_skills/inference/prompt/openmathinstruct/text_masked_base.yaml

@@ -1,5 +1,5 @@
 few_shot_examples:
-  template: "Question:\n{question}\n\n{context}My masked solution:\n{generation}\n\n\n\n\n\n"
+  template: "Question:\n{question}{context}\n\nMy masked solution:\n{generation}\n\n\n\n\n\n"
 system: ""
 user: "\
   Here are some examples of questions, solutions, \
@@ -8,6 +8,6 @@ user: "\
   mentioned in the solution that are not mentioned in the question. \
   Mask the intermediate values calculated in the solution via capitalized variables, such as M, N, M_01.\n\n\
   {examples}\
-  Question:\n{question}\n\n{context}"
+  Question:\n{question}{context}\n\n"
 prompt_template: "{system}{user}My masked solution:\n{generation}"
 stop_phrases: ["\n\n\n\n\n\n"]

nemo_skills/inference/prompt/question_generation/augmentation.yaml

@@ -1,18 +1,18 @@
 few_shot_examples:
-  template: "Question:\n{question}\n\n{context}Write another question similar to this one:\n{augmented_question}\n\n\n\n\n\n"
+  template: "Question:\n{question}{context}\n\nWrite another question similar to this one:\n{augmented_question}\n\n\n\n\n\n"
 
 system: ""
 
 user: |-
   You are an AI assistant that excels at creating similar questions. Follow the given examples.
 
   {examples}Question:
-  {question}
-
-  {context}
+  {question}{context}
 
 prompt_template: |-
-  {system}{user}Write another question similar to this one:
+  {system}{user}
+
+  Write another question similar to this one:
   {generation}
 
 stop_phrases: ["\n\n\n\n\n\n"]

nemo_skills/inference/prompt/question_generation/rephrasing.yaml

@@ -1,18 +1,18 @@
 few_shot_examples:
-  template: "Question:\n{question}\n\n{context}Rephrase the above question:\n{rephrased_question}\n\n\n\n\n\n"
+  template: "Question:\n{question}{context}\n\nRephrase the above question:\n{rephrased_question}\n\n\n\n\n\n"
 
 system: ""
 
 user: |-
   You are an AI assistant that excels at rephrasing questions. Follow the given examples.
 
   {examples}Question:
-  {question}
-
-  {context}
+  {question}{context}
 
 prompt_template: |-
-  {system}{user}Rephrase the above question:
+  {system}{user}
+
+  Rephrase the above question:
   {generation}
 
 stop_phrases: ["\n\n\n\n\n\n"]

nemo_skills/inference/prompt/utils.py

@@ -61,7 +61,7 @@ def retrieve(self, query: str, top_k: int = 1):
 
 @nested_dataclass
 class FewShotExamplesConfig:
-    template: str = MISSING
+    template: str = ""
     num_few_shots: int = 0
 
     examples_type: Optional[str] = None
@@ -241,11 +241,11 @@ def get_prompt_config(prompt_type: str) -> PromptConfig:
 # and so most often will require changes to the code anyway
 context_templates = {
     "empty": "",
-    "reference_solution": "\n\nReference solution (do not copy it):\n{reference_solution}\n\n",
-    "masked_solution": "\n\nReference solution:\n{masked_reference_solution}\n\n",
-    "table": "\n\nUse the following table to answer the question:\n{table}\n\n",
+    "reference_solution": "\n\nReference solution (do not copy it):\n{reference_solution}",
+    "masked_solution": "\n\nReference solution:\n{masked_reference_solution}",
+    "table": "\n\nUse the following table to answer the question:\n{table}",
     "table_solution": (
         "\n\nUse the following table to answer the question:\n{table}\n"
-        "Reference solution (do not copy it):\n{reference_solution}\n\n"
+        "Reference solution (do not copy it):\n{reference_solution}"
     ),
 }

tests/test_prompts.py

@@ -111,9 +111,7 @@ def test_llama3_instruct_prompt():
 
     expected_prompt = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>
 
-You are Meta AI, a sophisticated and energetic AI Assistant. You excel at solving mathematical problems.
-
-You will look at the examples provided by the user and try to follow the solution format as much as possible.<|eot_id|><|start_header_id|>user<|end_header_id|>
+<|eot_id|><|start_header_id|>user<|end_header_id|>
 
 Here are some examples of questions and solutions followed by a new question that you need to solve.
 Make sure to put the answer (and only answer) inside \\boxed{}.
@@ -243,8 +241,11 @@ def test_nemotron_zeroshot_prompt():
     expected_prompt = """<extra_id_0>System
 
 <extra_id_1>User
-2 + 2 = ?
-<extra_id_1>Assistant
+You are an expert in math. Given the math question below, I want you to reason through the steps and then give a final answer.
+
+Your final answer should be inside \\boxed{}.
+
+2 + 2 = ?<extra_id_1>Assistant
 """
     assert prompt.build_string({'question': '2 + 2 = ?'}) == expected_prompt
 
@@ -287,8 +288,6 @@ def test_nemotron_fewshot_prompt():
 Question:
 2 + 2 = ?
 
-
-
 Don't forget that your final answer should be inside \\boxed{}!
 <extra_id_1>Assistant
 """
@@ -343,8 +342,6 @@ def test_nemotron_fewshot_prompt_reference():
 Reference solution (do not copy it):
 What should I do??
 
-
-
 Don't forget that your final answer should be inside \\boxed{}!
 <extra_id_1>Assistant
 """