From 066c521ea2c2df81a744d34d63ffa8b6b6b30854 Mon Sep 17 00:00:00 2001 From: Axy Date: Thu, 27 Aug 2026 18:15:27 +0200 Subject: [PATCH] Answer bulk impl --- src/rag/__init__.py | 41 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/src/rag/__init__.py b/src/rag/__init__.py index 4157d8b..8dab130 100644 --- 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) -- 2.53.0