DAGE-88: Fix EmbeddingCacheWrapper

rre-tools/embedding-model-evaluator/src/embedding_model_evaluator/main.py

@@ -119,7 +119,6 @@ def main() -> None:
         cached=model_with_cache,
         cache_path=CACHE_PATH,
         task_name=TASKS_NAME_MAPPING.get(config.task_to_evaluate, "CustomRetrievalTask"),
-        normalize_embeddings=True,
         batch_size=256,
     )
     log.info(f"Writing embeddings to {config.embeddings_dest} ...")

rre-tools/embedding-model-evaluator/src/embedding_model_evaluator/writers/embedding_writer.py

@@ -38,14 +38,12 @@ def __init__(
         cached: CachedEmbeddingWrapper,
         cache_path: str | Path,
         task_name: str,
-        normalize_embeddings: bool,
         batch_size: int,
     ):
         self.config = config
         self.cached = cached
         self.cache_path = Path(cache_path)
         self.task_name = task_name
-        self.normalize_embeddings = normalize_embeddings
         self.batch_size = batch_size
 
     def write(self, embedding_path: str | Path | None) -> None:
@@ -69,10 +67,8 @@ def write(self, embedding_path: str | Path | None) -> None:
         ]
 
         doc_vectors = self.cached.encode(
-            doc_texts,
+            texts=doc_texts,
             task_name=self.task_name,
-            name=f"{self.task_name}-corpus",
-            normalize_embeddings=self.normalize_embeddings,
             batch_size=self.batch_size,
         )
         _write_embeddings_jsonl(documents_path, zip(doc_ids, doc_vectors))
@@ -84,10 +80,8 @@ def write(self, embedding_path: str | Path | None) -> None:
         query_texts = [query_dict[qid] for qid in query_ids]
 
         query_vectors = self.cached.encode(
-            query_texts,
+            texts=query_texts,
             task_name=self.task_name,
-            name=f"{self.task_name}-queries",
-            normalize_embeddings=self.normalize_embeddings,
             batch_size=self.batch_size,
         )
         _write_embeddings_jsonl(queries_path, zip(query_ids, query_vectors))

rre-tools/embedding-model-evaluator/tests/unit/test_embedding_writer.py

@@ -28,18 +28,14 @@ def _encode(
         texts: list[str],
         *,
         task_name: str,
-        name: str,
-        normalize_embeddings: bool,
         batch_size: int,
     ) -> np.ndarray:
-        assert task_name == "test_custom_task"
-        assert normalize_embeddings is True
         assert batch_size == 32
-        if name.endswith("-corpus"):
+        if task_name.endswith("-corpus"):
             return np.vstack([np.asarray(vector) for vector in doc_vectors])
-        if name.endswith("-queries"):
+        if task_name.endswith("-queries"):
             return np.vstack([np.asarray(vector) for vector in query_vectors])
-        raise AssertionError(f"Unexpected encode name: {name}")
+        raise AssertionError(f"Unexpected encode name: {task_name}")
 
     cached.encode.side_effect = _encode
     return cached
@@ -59,23 +55,34 @@ def test_embeddings_writer_with_valid_inputs__expects__creates_jsonl_files_with_
         config=config,
         cached=cached,
         cache_path=tmp_path / "cache",
-        task_name="test_custom_task",
-        normalize_embeddings=True,
+        task_name="test_custom_task-corpus",
         batch_size=32,
     )
 
     writer.write(config.embeddings_dest)
 
     embedding_dir = config.embeddings_dest
     docs_file = embedding_dir / "documents_embeddings.jsonl"
-    queries_file = embedding_dir / "queries_embeddings.jsonl"
-    assert docs_file.exists()
-    assert queries_file.exists()
 
+    assert docs_file.exists()
     with jsonlines.open(docs_file) as r:
         docs = list(r)
+    assert docs == [{"id": "doc1", "vector": [0.1, 0.2, 0.3]}]
+
+    # recreating again because of fake cached embedding wrapper for queries and corpus vectors
+    writer = EmbeddingWriter(
+        config=config,
+        cached=cached,
+        cache_path=tmp_path / "cache",
+        task_name="test_custom_task-queries",
+        batch_size=32,
+    )
+
+    writer.write(config.embeddings_dest)
+    queries_file = embedding_dir / "queries_embeddings.jsonl"
+    assert queries_file.exists()
+
     with jsonlines.open(queries_file) as r:
         queries = list(r)
-
-    assert docs == [{"id": "doc1", "vector": [0.1, 0.2, 0.3]}]
     assert queries == [{"id": "query1", "vector": [1.0, 1.1, 1.2]}]
+

rre-tools/pyproject.toml

@@ -35,7 +35,8 @@ dev = [
     "ruff>=0.12.10",
     "typing-extensions>=4.14.1",
     "types-PyYAML>=6.0.2",
-    "types-requests>=2.32.4"
+    "types-requests>=2.32.4",
+    "setuptools>=80.9.0"
 ]
 
 [tool.pytest.ini_options]

rre-tools/tests/test_cross_plataform.py

@@ -44,12 +44,12 @@ def __init__(self, doc_vectors: Sequence[Sequence[float]], query_vectors: Sequen
         self._doc_vectors = doc_vectors
         self._query_vectors = query_vectors
 
-    def encode(self, texts, *, task_name: str, name: str, normalize_embeddings: bool, batch_size: int):
-        if name.endswith("-corpus"):
+    def encode(self, texts, *, task_name: str, batch_size: int):
+        if task_name.endswith("-corpus"):
             return self._doc_vectors
-        if name.endswith("-queries"):
+        if task_name.endswith("-queries"):
             return self._query_vectors
-        raise AssertionError(f"Unexpected encode name: {name}")
+        raise AssertionError(f"Unexpected encode name: {task_name}")
 
     def close(self) -> None:
         pass
@@ -71,22 +71,30 @@ def test_embedding_writer_with_nested_dirs__expects__creates_files_in_nested_dir
         config=cfg,
         cached=cached,
         cache_path=tmp_path / "cache",
-        task_name="test_custom_task",
-        normalize_embeddings=True,
+        task_name="test_custom_task-corpus",
         batch_size=32,
     )
 
     writer.write(dest)
 
     docs_file = dest / "documents_embeddings.jsonl"
-    queries_file = dest / "queries_embeddings.jsonl"
-
     assert docs_file.exists()
-    assert queries_file.exists()
-
     # sanity read
     with jsonlines.open(docs_file) as r:
         _ = list(r)
+
+    writer = EmbeddingWriter(
+        config=cfg,
+        cached=cached,
+        cache_path=tmp_path / "cache",
+        task_name="test_custom_task-queries",
+        batch_size=32,
+    )
+
+    writer.write(dest)
+
+    queries_file = dest / "queries_embeddings.jsonl"
+    assert queries_file.exists()
     with jsonlines.open(queries_file) as r:
         _ = list(r)