Tidying up generation and stuff

diff --git a/src/rag/answering.py b/src/rag/answering.py
index acc5898f2ebeafb3d93045b485ea0dd283481a48..9f526f3b4a4bf3d181542d8b7d249ef88ffd7bea 100644 (file)
--- a/src/rag/answering.py
+++ b/src/rag/answering.py
@@ -1,6 +1,7 @@
 """Answering questions using LLMs."""
 
 import logging
+import re
 from collections.abc import Callable
 from typing import Any, cast, override
 
@@ -10,13 +11,13 @@ from rag.models import AnsweredQuestion, RetrievedQuestion, Source
 class InferenceBackend:
     """A simple inference backend using transformers."""
 
-    SYSTEM_PROMPT: str = """You are a codebase and document search assistant.
-    You answer querries factually using the provided sources.
-    Answer directly with no reasoning, no <think> tags.
-    Answer simply, in a few short sentences, in a single line.
-    Assume the user cannot see the sources, they are your knowledge, \
-    do not mention them.
-    End your answer with a newline."""
+    SYSTEM_PROMPT: str = """
+    You are a codebase and document search assistant.
+    You answer queries factually using the provided sources exclusitvely.
+    Answer directly and simply.
+    Give a full answer without redirecting them to another document.
+    If you cannot answer adequately from the sources, say so instead.
+    """
 
     def __init__(self, model: str) -> None:
         """Initialize the backend fetching the given model from huggingface."""
@@ -44,23 +45,24 @@ class InferenceBackend:
             )
         return None
 
-    def _prompt(self, question: RetrievedQuestion) -> str:
+    def _prompt(self, question: RetrievedQuestion) -> Any:
         sources = "\n\n".join(
-            s
+            str(e) + ":\n" + s
             for e in question.retrieved_sources
             if (s := self._fetch_range(e))
         )
-        sep = "\n\n\n"
-        return (
-            "# System instructions\n"
-            + self.SYSTEM_PROMPT
-            + sep
-            + "# Context\n"
-            + sources
-            + "# Question\n"
-            + question.question
-            + sep
-            + "# Answer\n"
+
+        messages = [
+            {"role": "system", "content": self.SYSTEM_PROMPT},
+            {"role": "system", "content": "CONTEXT:\n" + sources},
+            {"role": "user", "content": question.question},
+        ]
+        return self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_tensors="pt",
+            enable_thinking=False,
         )
 
     def _answer(
@@ -86,35 +88,31 @@ class InferenceBackend:
                     return
                 s: str = cast(
                     str,
-                    self.backend.tokenizer.decode(
-                        value, skip_special_tokens=True
-                    ),
+                    self.backend.tokenizer.decode(value),
                 )
-                self.cb(s.rstrip("\n"))
+                self.cb(s)
 
             @override
             def end(self) -> None:
                 self.cb("\n")
 
-        prompt = self._prompt(question)
-        model_inputs = self.tokenizer([prompt], return_tensors="pt").to(
-            self.model.device
-        )
+        model_inputs = self._prompt(question)
         generated_ids = self.model.generate(
             **model_inputs,
             max_new_tokens=max_tokens,
-            stop_strings="\n",
-            tokenizer=self.tokenizer,
             streamer=Streamer(self, cb),
         )
-        full_response: str = cast(
+        response: str = cast(
             str,
             self.tokenizer.batch_decode(
-                generated_ids,
-                skip_special_tokens=True,
+                generated_ids, skip_special_tokens=True
             )[0],
         )
-        return full_response[len(prompt) :].strip()
+        answer = response.split("assistant\n")[-1]
+        clean_answer = re.sub(
+            r"<think>.*</think>", "", answer, flags=re.DOTALL
+        )
+        return clean_answer
 
     def answer(
         self,

diff --git a/src/rag/models.py b/src/rag/models.py
index 7508598492f18ecb62c9b7996dd7c48e71e7892a..279893e85934bed9ad62b6c591a0e31386f34a10 100644 (file)
--- a/src/rag/models.py
+++ b/src/rag/models.py
@@ -69,8 +69,6 @@ class DatasetAnswer(BaseModel):
     question: str
     answer: str
     sources: list[Source]
-    difficulty: str
-    is_valid: bool
 
 
 class AnswerSet(BaseModel):