fix multi-shot prompt

cicl-stanford · Jul 16, 2023 · fe647d6 · fe647d6
1 parent 5ed88e3
commit fe647d6
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/code/src/bigtom.py b/code/src/bigtom.py
@@ -95,11 +95,12 @@ def gen_chat(args):
                 examples.append(example)
         random.shuffle(examples)
 
-        # 2-shots by default
-        messages = [system_message, human_message_0]
+        # 3-shots by default
+        messages = [system_message]
         for i in range(args.num_shots):
+            messages.append(human_message_0)
             messages.append(AIMessage(content=response_template.format(**examples[i])))
-            messages.append(human_message_1)
+        messages.append(human_message_1)
 
         if args.verbose:
             print(f"------ messages ------")