Answer bulk impl

diff --git a/src/rag/__init__.py b/src/rag/__init__.py
index 4157d8bdb66a60ea342ed09e348092813543ef27..8dab130adcc7407a642eb0948edfdcdc5aad1690 100644 (file)
--- a/src/rag/__init__.py
+++ b/src/rag/__init__.py
@@ -12,6 +12,7 @@ from rag.chunking import FileType, chunk_file
 from rag.models import (
     RagDataset,
     RetrievedQuestion,
+    SearchAnswers,
     SearchResults,
     Source,
 )
@@ -165,3 +166,43 @@ class RAG:
             max_tokens,
             cb=lambda s: print(s, end="", flush=True),
         )
+
+    def answer_dataset(
+        self,
+        student_search_result_path: str,
+        save_directory: str,
+        /,
+        index: str = "data/processed",
+        max_tokens: int = 100,
+    ) -> None:
+        bm25 = self._bm25(index)
+        try:
+            with open(student_search_result_path) as f:
+                dataset = SearchResults.model_validate_json(f.read())
+        except (OSError, ValidationError) as e:
+            logging.getLogger(__name__).error(
+                f"Failed to open search results {e}"
+            )
+            exit(1)
+        inference = self._inference()
+        result = SearchAnswers(
+            search_results=[
+                inference.answer(question, max_tokens)
+                for question in tqdm(
+                    dataset.search_results, desc="Generating answers"
+                )
+            ],
+            k=dataset.k,
+        )
+        try:
+            os.makedirs(save_directory, exist_ok=True)
+            with open(
+                save_directory
+                + "/"
+                + os.path.basename(student_search_result_path),
+                "w",
+            ) as f:
+                f.write(result.model_dump_json())
+        except (OSError, ValidationError) as e:
+            logging.getLogger(__name__).error(f"Failed to open dataset {e}")
+            exit(1)