Only missing flake8 and readme

diff --git a/src/rag/__init__.py b/src/rag/__init__.py
index f0b6e631b97c9c7ef63e13348176793c02e643c6..681a1c3aed7d014ae1841d5211f67f2b1725fcb8 100644 (file)
--- a/src/rag/__init__.py
+++ b/src/rag/__init__.py
@@ -1,16 +1,16 @@
 import logging
 import os
 from collections.abc import Callable
-from sys import stderr
 
 import fire
 import pydantic
-from pydantic_core import ValidationError
 from tqdm import tqdm
 
 from rag.answering import InferenceBackend
 from rag.chunking import FileType, chunk_file
+from rag.evaluate import sources_overlap
 from rag.models import (
+    AnswerSet,
     RagDataset,
     RetrievedQuestion,
     SearchAnswers,
@@ -41,7 +41,9 @@ class RAG:
             with open(path) as f:
                 return cb(f.read())
         except Exception as e:
-            logging.getLogger(__name__).error(f"Failed to read {name}: {e}")
+            logging.getLogger(__name__).error(
+                f"Failed to read {name} ({path}): {e}"
+            )
             exit(1)
 
     @staticmethod
@@ -55,7 +57,9 @@ class RAG:
                 with open(path, "w") as f:
                     f.write(data)
         except Exception as e:
-            logging.getLogger(__name__).error(f"Failed to write {name}: {e}")
+            logging.getLogger(__name__).error(
+                f"Failed to write {name} ({path}): {e}"
+            )
             exit(1)
 
     def index(
@@ -65,9 +69,14 @@ class RAG:
         index: str = "data/processed",
         max_chunk_size: int = 2000,
     ) -> None:
+        if max_chunk_size < 1:
+            logging.getLogger(__name__).error(
+                f"Invalid max chunk size {max_chunk_size}"
+            )
+            exit(1)
         chunks = {}
         for dirpath, _, files in tqdm(
-            list(os.walk(data)), desc="Collecting files"
+            list(os.walk(data)), desc="Chunking files"
         ):
             for file in files:
                 path = dirpath + "/" + file
@@ -114,6 +123,9 @@ class RAG:
     def search(
         self, query: str, /, k: int = 5, index: str = "data/processed"
     ) -> None:
+        if k < 1:
+            logging.getLogger(__name__).error(f"Invalid k {k}")
+            exit(1)
         for source in self._search(query, k, index):
             print(source)
 
@@ -125,6 +137,9 @@ class RAG:
         k: int = 5,
         index: str = "data/processed",
     ) -> None:
+        if k < 1:
+            logging.getLogger(__name__).error(f"Invalid k {k}")
+            exit(1)
         bm25 = self._bm25(index)
         dataset = self._readf(dataset_path, RagDataset.model_validate_json)
         result = SearchResults(
@@ -159,6 +174,14 @@ class RAG:
         index: str = "data/processed",
         max_tokens: int = 100,
     ) -> None:
+        if k < 1:
+            logging.getLogger(__name__).error(f"Invalid k {k}")
+            exit(1)
+        if max_tokens < 1:
+            logging.getLogger(__name__).error(
+                f"Invalid max tokens {max_tokens}"
+            )
+            exit(1)
         sources = self._search(query, k, index)
         inference = self._inference()
         print("Sources:")
@@ -176,10 +199,13 @@ class RAG:
         student_search_result_path: str,
         save_directory: str,
         /,
-        index: str = "data/processed",
         max_tokens: int = 100,
     ) -> None:
-        bm25 = self._bm25(index)
+        if max_tokens < 1:
+            logging.getLogger(__name__).error(
+                f"Invalid max tokens {max_tokens}"
+            )
+            exit(1)
         dataset = self._readf(
             student_search_result_path,
             SearchResults.model_validate_json,
@@ -201,3 +227,59 @@ class RAG:
             + os.path.basename(student_search_result_path),
             result.model_dump_json(),
         )
+
+    def _evaluate(
+        self, answers: SearchResults, references: AnswerSet, k: int, iou: float
+    ) -> None:
+        count = len(references.rag_questions)
+        success = sum(
+            1
+            for answer, reference in zip(
+                answers.search_results,
+                references.rag_questions,
+                strict=False,
+            )
+            if any(
+                sources_overlap(a, b, iou)
+                for i, a in enumerate(answer.retrieved_sources)
+                for b in reference.sources
+                if i < k
+            )
+        )
+        print(
+            f"recall@{k}: {success / count * 100.0:.2f}% ({success} / {count})"
+        )
+
+    def evaluate(
+        self,
+        student_search_result_path: str,
+        dataset_path: str,
+        /,
+        k: int = 5,
+        iou: float = 0.05,
+    ) -> None:
+        if k < 1:
+            logging.getLogger(__name__).error(f"Invalid k {k}")
+            exit(1)
+        answers = self._readf(
+            student_search_result_path,
+            SearchResults.model_validate_json,
+            name="student search results",
+        )
+        references = self._readf(
+            dataset_path,
+            AnswerSet.model_validate_json,
+            name="reference dataset",
+        )
+        if len(answers.search_results) != len(references.rag_questions):
+            print("Mismatch between search result and reference set lenghts!")
+            exit(1)
+        if len(answers.search_results) == 0:
+            print("Empty dataset, grade would not make sense.")
+            exit(0)
+        prev_i = 1
+        i = 1
+        while i < k:
+            self._evaluate(answers, references, i, iou)
+            prev_i, i = i, i + prev_i
+        self._evaluate(answers, references, k, iou)

diff --git a/src/rag/answering.py b/src/rag/answering.py
index f201793b47579af1e99b0684bd969ccbf40cd6ab..008d13dd47ec3e46c701d064c54f663741e2dc67 100644 (file)
--- a/src/rag/answering.py
+++ b/src/rag/answering.py
@@ -2,9 +2,6 @@ import logging
 from collections.abc import Callable
 from typing import Any, cast
 
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from transformers.generation import BaseStreamer  # type:ignore
-
 from rag.models import AnsweredQuestion, RetrievedQuestion, Source
 
 
@@ -17,28 +14,9 @@ class InferenceBackend:
     do not mention them.
     End your answer with a newline."""
 
-    class _Streamer(BaseStreamer):
-        def __init__(
-            self, backend: "InferenceBackend", cb: Callable[[str], None]
-        ) -> None:
-            self.backend = backend
-            self.cb = cb
-            self.skip = True
-
-        def put(self, value: Any) -> None:
-            if self.skip:
-                self.skip = False
-                return
-            s: str = cast(
-                str,
-                self.backend.tokenizer.decode(value, skip_special_tokens=True),
-            )
-            self.cb(s.rstrip("\n"))
-
-        def end(self) -> None:
-            self.cb("\n")
-
     def __init__(self, model: str) -> None:
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
         self.model: Any = AutoModelForCausalLM.from_pretrained(
             model, device_map="auto"
         )
@@ -86,6 +64,31 @@ class InferenceBackend:
         max_tokens: int,
         cb: Callable[[str], None],
     ) -> str:
+        from transformers.generation import BaseStreamer  # type: ignore
+
+        class Streamer(BaseStreamer):
+            def __init__(
+                self, backend: "InferenceBackend", cb: Callable[[str], None]
+            ) -> None:
+                self.backend = backend
+                self.cb = cb
+                self.skip = True
+
+            def put(self, value: Any) -> None:
+                if self.skip:
+                    self.skip = False
+                    return
+                s: str = cast(
+                    str,
+                    self.backend.tokenizer.decode(
+                        value, skip_special_tokens=True
+                    ),
+                )
+                self.cb(s.rstrip("\n"))
+
+            def end(self) -> None:
+                self.cb("\n")
+
         prompt = self._prompt(question)
         model_inputs = self.tokenizer([prompt], return_tensors="pt").to(
             self.model.device
@@ -95,7 +98,7 @@ class InferenceBackend:
             max_new_tokens=max_tokens,
             stop_strings="\n",
             tokenizer=self.tokenizer,
-            streamer=self._Streamer(self, cb),
+            streamer=Streamer(self, cb),
         )
         full_response = self.tokenizer.batch_decode(
             generated_ids,

diff --git a/src/rag/chunking.py b/src/rag/chunking.py
index aabaf7366f87d0a155cd4c546b5ec38c806951be..e14c6c35828709cf7218bde9b8e0bb6d58151b00 100644 (file)
--- a/src/rag/chunking.py
+++ b/src/rag/chunking.py
@@ -43,7 +43,10 @@ def chunk_by_depth(
                 chunks.pop()
             chunks.append(line)
     if s.count("\n") <= 1:
-        chunks = s.split()
+        chunks = s.split(" ")
+        chunks = [s for s in chunks if s]
+        for i in range(len(chunks) - 1):
+            chunks[i] = chunks[i] + " "
         if len(chunks) <= 1:
             chunks = [
                 s[i : min(len(chunks), i + max_chunk_size)]

diff --git a/src/rag/evaluate.py b/src/rag/evaluate.py
new file mode 100644 (file)
index 0000000..090ff5a
--- /dev/null
+++ b/src/rag/evaluate.py
@@ -0,0 +1,14 @@
+from rag.models import Source
+
+
+def sources_overlap(a: Source, b: Source, min_iou: float) -> bool:
+    union_start = min(a.first_character_index, b.first_character_index)
+    union_end = max(a.last_character_index, b.last_character_index) + 1
+    union = union_end - union_start
+    if union < 1:
+        union = 1
+    intersection_start = max(a.first_character_index, b.first_character_index)
+    intersection_end = min(a.last_character_index, b.last_character_index)
+    intersection = intersection_end - intersection_start
+    iou = intersection / union
+    return iou >= min_iou

diff --git a/src/rag/models.py b/src/rag/models.py
index cb67b1794273f823766327a2bd10f9d55aab6708..4ba9f7854034c734448d06be9f162c7a86c9d81c 100644 (file)
--- a/src/rag/models.py
+++ b/src/rag/models.py
@@ -41,3 +41,16 @@ class SearchResults(BaseModel):
 class SearchAnswers(BaseModel):
     search_results: list[AnsweredQuestion]
     k: int
+
+
+class DatasetAnswer(BaseModel):
+    question_id: str
+    question: str
+    answer: str
+    sources: list[Source]
+    difficulty: str
+    is_valid: bool
+
+
+class AnswerSet(BaseModel):
+    rag_questions: list[DatasetAnswer]