modify chat

llama3_chat.py

@@ -20,7 +20,7 @@ def main(
         max_seq_len=max_seq_len,
         max_batch_size=max_batch_size,
     )
-    dialogs: List[Dialog] = [[]] # Start with an empty dialog
+    dialogs: List[Dialog] = [] # Start with an empty dialog
     try:
         # Continue util the user decides to stop
         while True:
@@ -28,19 +28,21 @@ def main(
             # Allow the user to quit the dialogue
             if user_input.lower() in ['stop', 'exit']:
                 break
-            dialogs[0].append({"role": "user", "content": user_input})
+            dialogs.append({"role": "user", "content": user_input})
             # Generate response based on the current dialog context
             results  = generator.chat_completion(
                 [dialogs],
                 max_gen_len=max_gen_len,
                 temperature=temperature,
-                top_p=top_p,)[0]
-            response = results['generation']['content']
+                top_p=top_p,
+            )
+            response = results[0]['generation']['content']
             print(f"Assistant: {response}\n")
             # Append the generated response to the dialog
-            dialogs[0].append({"role": "assistant", "content": response})
+            dialogs.append({"role": "assistant", "content": response})
     except KeyboardInterrupt:
         print("Exiting dialogue.")
 
+
 if __name__ == "__main__":
     fire.Fire(main)