Update src/RAG/rag_pipeline.py

src/RAG/rag_pipeline.py

@@ -37,14 +37,18 @@ class TOSAssistant:
         print("Document ingested and Vector Store created.")
 
     def format_llama3_prompt(self, system_msg, user_msg):
+        # return (
+        #     f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n"
+        #     f"{system_msg}<|eot_id|>"
+        #     f"<|start_header_id|>user<|end_header_id|>\n"
+        #     f"{user_msg}<|eot_id|>"
+        #     f"<|start_header_id|>assistant<|end_header_id|>\n"
+        # )
         return (
-            f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n"
-            f"{system_msg}<|eot_id|>"
-            f"<|start_header_id|>user<|end_header_id|>\n"
-            f"{user_msg}<|eot_id|>"
-            f"<|start_header_id|>assistant<|end_header_id|>\n"
+            f"<|im_start|>system\n{system_msg}<|im_end|>\n"
+            f"<|im_start|>user\n{user_msg}<|im_end|>\n"
+            f"<|im_start|>assistant\n"
         )
-    
     def generate_global_summary(self):
         if not self.full_text:
             return "No document loaded."