Fix broken stuff

1 files changed, 6 insertions, 4 deletions

diff --git a/rag/generator/prompt.py b/rag/generator/prompt.py
index cedf610..75966e8 100644
--- a/rag/generator/prompt.py
+++ b/rag/generator/prompt.py
@@ -29,11 +29,13 @@ class Prompt:
             return f"{self.query}\n\n{ANSWER_INSTRUCTION}"
         else:
             return (
-                "Context information is below.\n"
+                "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n"
+                "Using the information contained in the context, give a comprehensive answer to the question.\n"
+                "If the answer cannot be deduced from the context, do not give an answer.\n\n"
+                "Context:\n"
                 "---\n"
                 f"{self.__context(self.documents)}\n\n"
                 "---\n"
-                f"{ANSWER_INSTRUCTION}"
-                f"Query: {self.query.strip()}\n\n"
-                "Answer:"
+                f"Question: {self.query}<|eot_id|>\n"
+                "<|start_header_id|>assistant<|end_header_id|>\n\n"
             )