Fixed the *totally awesome* pydantic model usage that's required

diff --git a/src/rag/__init__.py b/src/rag/__init__.py
index 76919b7c39f6c56dea5bce8869268edf9c67489f..17aa61b2c77b7f59366aef1f9733ec3e35d19f7a 100644 (file)
--- a/src/rag/__init__.py
+++ b/src/rag/__init__.py
@@ -8,8 +8,10 @@ from sys import stderr
 
 import fire
 import pydantic
+from pydantic.v1 import BaseModel
 from tqdm import tqdm
 
+import rag._stoopid
 from rag.answering import InferenceBackend
 from rag.chunking import FileType, chunk_file
 from rag.evaluate import sources_overlap
@@ -49,10 +51,18 @@ class RAG:
         self._inference_store: InferenceBackend | None = None
 
     @staticmethod
-    def _readf[T](path: str, cb: Callable[[str], T], name: str = "file") -> T:
+    def _readf[T](
+        path: str,
+        cb: Callable[[str], T],
+        name: str = "file",
+        stp: Callable[[str], str] | None = None,
+    ) -> T:
         try:
             with open(path) as f:
-                return cb(f.read())
+                s = f.read()
+                if stp:
+                    s = stp(s)
+                return cb(s)
         except Exception as e:
             logging.getLogger(__name__).error(
                 f"Failed to read {name} ({path}): {e}"
@@ -60,15 +70,20 @@ class RAG:
             exit(1)
 
     @staticmethod
-    def _writef(path: str, data: bytes | str, name: str = "file") -> None:
+    def _writef(
+        path: str,
+        data: bytes | str,
+        name: str = "file",
+        stp: Callable[[str], str] | None = None,
+    ) -> None:
         try:
             os.makedirs(os.path.dirname(path), exist_ok=True)
             if isinstance(data, bytes):
-                with open(path, "wb") as f:
-                    f.write(data)
-            else:
-                with open(path, "w") as f:
-                    f.write(data)
+                data = data.decode()
+            if stp:
+                data = stp(data)
+            with open(path, "w") as f:
+                f.write(data)
         except Exception as e:
             logging.getLogger(__name__).error(
                 f"Failed to write {name} ({path}): {e}"
@@ -157,7 +172,13 @@ class RAG:
             logging.getLogger(__name__).error(f"Invalid k {k}")
             exit(1)
         bm25 = self._bm25(index)
-        dataset = self._readf(dataset_path, RagQuestions.model_validate_json)
+        dataset = self._readf(
+            dataset_path,
+            RagQuestions.model_validate_json,
+            stp=lambda s: rag._stoopid.RagDataset.model_validate_json(
+                s
+            ).model_dump_json(),
+        )
         result = SearchResults(
             search_results=[
                 RetrievedQuestion(
@@ -175,6 +196,11 @@ class RAG:
             save_directory + "/" + os.path.basename(dataset_path),
             result.model_dump_json(),
             name="dataset",
+            stp=lambda s: (
+                rag._stoopid.StudentSearchResults.model_validate_json(
+                    s
+                ).model_dump_json()
+            ),
         )
 
     def _inference(self) -> InferenceBackend:
@@ -228,6 +254,11 @@ class RAG:
             student_search_result_path,
             SearchResults.model_validate_json,
             name="search results",
+            stp=lambda s: (
+                rag._stoopid.StudentSearchResults.model_validate_json(
+                    s
+                ).model_dump_json()
+            ),
         )
         inference = self._inference()
         result = SearchAnswers(
@@ -244,6 +275,11 @@ class RAG:
             + "/"
             + os.path.basename(student_search_result_path),
             result.model_dump_json(),
+            stp=lambda s: (
+                rag._stoopid.StudentSearchResultsAndAnswer.model_validate_json(
+                    s
+                ).model_dump_json()
+            ),
         )
 
     def _evaluate(
@@ -284,11 +320,19 @@ class RAG:
             student_search_result_path,
             SearchResults.model_validate_json,
             name="student search results",
+            stp=lambda s: (
+                rag._stoopid.StudentSearchResults.model_validate_json(
+                    s
+                ).model_dump_json()
+            ),
         )
         references = self._readf(
             dataset_path,
             AnswerSet.model_validate_json,
             name="reference dataset",
+            stp=lambda s: rag._stoopid.RagDataset.model_validate_json(
+                s
+            ).model_dump_json(),
         )
         if len(answers.search_results) != len(references.rag_questions):
             print("Mismatch between search result and reference set lenghts!")