fixed: dump & load modules and added tests

Author: Darinochka

autointent/modules/embedding/_retrieval.py

@@ -23,6 +23,15 @@ class VectorDBMetadata(BaseMetadataDict):
     max_length: int | None
 
 
+class ClassifierMetadata(BaseMetadataDict):
+    """Metadata class for LogisticRegressionCV and LabelEncoder."""
+
+    coef_: list[list[float]]
+    intercept_: list[float]
+    params: dict[str, any]
+    classes: list[str]
+
+
 class LogRegEmbedding(EmbeddingModule):
     r"""
     Module for managing classification operations using logistic regression.
@@ -63,6 +72,7 @@ class LogRegEmbedding(EmbeddingModule):
 
     """
 
+    vector_index: VectorIndex
     classifier: LogisticRegressionCV
     label_encoder: LabelEncoder
     name = "logreg"
@@ -201,33 +211,29 @@ def clear_cache(self) -> None:
 
     def dump(self, path: str) -> None:
         """
-        Save the module's metadata and model parameters to a specified directory.
+        Save the module's metadata, classifier parameters, and label encoder to a specified directory.
 
         :param path: Path to the directory where assets will be dumped.
         """
-        metadata = VectorDBMetadata(
+        self.metadata = VectorDBMetadata(
             batch_size=self.batch_size,
             max_length=self.max_length,
-            db_dir=self.db_dir,
+            db_dir=str(self.db_dir),
         )
 
         dump_dir = Path(path)
-        with (dump_dir / "metadata.json").open("w") as file:
-            json.dump(metadata.__dict__, file, indent=4)
-
-        model_path = dump_dir / "logreg_model.json"
-        with model_path.open("w") as file:
-            json.dump(
-                {
-                    "coef": self.classifier.coef_.tolist(),
-                    "intercept": self.classifier.intercept_.tolist(),
-                    "classes": self.label_encoder.classes_.tolist(),
-                },
-                file,
-                indent=4,
-            )
-
-        super().dump(path)
+        with (dump_dir / self.metadata_dict_name).open("w") as file:
+            json.dump(self.metadata, file, indent=4)
+        self.vector_index.dump(dump_dir)
+
+        self.classifier_metadata = ClassifierMetadata(
+            coef_=self.classifier.coef_.tolist(),
+            intercept_=self.classifier.intercept_.tolist(),
+            classes=self.label_encoder.classes_.tolist(),
+            params=self.classifier.get_params(),
+        )
+        with (dump_dir / "classifier.json").open("w") as file:
+            json.dump(self.classifier_metadata, file, indent=4)
 
     def load(self, path: str) -> None:
         """
@@ -236,24 +242,28 @@ def load(self, path: str) -> None:
         :param path: Path to the directory containing the dumped assets.
         """
         dump_dir = Path(path)
+        with (dump_dir / self.metadata_dict_name).open() as file:
+            self.metadata: VectorDBMetadata = json.load(file)
 
-        with (dump_dir / "metadata.json").open() as file:
-            metadata_dict = json.load(file)
-            self.batch_size = metadata_dict.get("batch_size", self.batch_size)
-            self.max_length = metadata_dict.get("max_length", self.max_length)
-            self._db_dir = metadata_dict.get("db_dir", self._db_dir)
-
-        model_path = dump_dir / "logreg_model.json"
-        with model_path.open() as file:
-            model_data = json.load(file)
-            self.classifier = LogisticRegressionCV()
-            self.k = model_data["k"]
-            self.classifier.coef_ = [model_data["coef"]]
-            self.classifier.intercept_ = model_data["intercept"]
-            self.label_encoder = LabelEncoder()
-            self.label_encoder.classes_ = model_data["classes"]
-
-        super().load(path)
+        vector_index_client = VectorIndexClient(
+            embedder_device=self.embedder_device,
+            db_dir=self.metadata["db_dir"],
+            embedder_batch_size=self.metadata["batch_size"],
+            embedder_max_length=self.metadata["max_length"],
+            embedder_use_cache=self.embedder_use_cache,
+        )
+        self.vector_index = vector_index_client.get_index(self.embedder_name)
+
+        with (dump_dir / "classifier.json").open() as file:
+            self.classifier_metadata: ClassifierMetadata = json.load(file)
+
+        self.classifier = LogisticRegressionCV()
+        self.classifier.set_params(**self.classifier_metadata["params"])
+        self.classifier.coef_ = self.classifier_metadata["coef_"]
+        self.classifier.intercept_ = self.classifier_metadata["intercept_"]
+
+        self.label_encoder = LabelEncoder()
+        self.label_encoder.classes_ = self.classifier_metadata["classes"]
 
     def predict(self, utterances: list[str]) -> list[int | list[int]]:
         """

tests/modules/retrieval/test_logreg.py

@@ -0,0 +1,73 @@
+import shutil
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import numpy as np
+
+from autointent.modules.embedding import LogRegEmbedding
+from tests.conftest import setup_environment
+
+
+def test_get_assets_returns_correct_artifact_for_logreg():
+    db_dir, dump_dir, logs_dir = setup_environment()
+    module = LogRegEmbedding(k=5, embedder_name="sergeyzh/rubert-tiny-turbo", db_dir=db_dir)
+    artifact = module.get_assets()
+    assert artifact.embedder_name == "sergeyzh/rubert-tiny-turbo"
+
+
+def test_fit_trains_model():
+    db_dir, dump_dir, logs_dir = setup_environment()
+    module = LogRegEmbedding(k=5, embedder_name="sergeyzh/rubert-tiny-turbo", db_dir=db_dir)
+
+    utterances = ["hello", "goodbye", "hi", "bye", "bye", "hello", "welcome", "hi123", "hiii", "bye-bye", "bye!"]
+    labels = [0, 1, 0, 1, 1, 0, 0, 0, 0, 1, 1]
+    module.fit(utterances, labels)
+
+    assert module.classifier.coef_ is not None
+    assert len(module.classifier.coef_) > 0
+    assert module.label_encoder.classes_.tolist() == [0, 1]
+
+
+def test_score_evaluates_model():
+    db_dir, dump_dir, logs_dir = setup_environment()
+    module = LogRegEmbedding(k=5, embedder_name="sergeyzh/rubert-tiny-turbo", db_dir=db_dir)
+
+    utterances = ["hello", "goodbye", "hi", "bye", "bye", "hello", "welcome", "hi123", "hiii", "bye-bye", "bye!"]
+    labels = [0, 1, 0, 1, 1, 0, 0, 0, 0, 1, 1]
+    module.fit(utterances, labels)
+
+    mock_context = MagicMock()
+    mock_context.data_handler.test_utterances.return_value = ["hello", "goodbye"]
+    mock_context.data_handler.test_labels.return_value = [0, 1]
+
+    def mock_metric_fn(true_labels, predicted_labels):
+        return sum(1 for t, p in zip(true_labels, predicted_labels[0], strict=False) if t == p) / len(true_labels)
+
+    score = module.score(mock_context, split="test", metric_fn=mock_metric_fn)
+
+    assert 0 <= score <= 1
+    assert score > 0
+
+
+def test_dump_and_load_preserves_model_state():
+    db_dir, dump_dir, logs_dir = setup_environment()
+    module = LogRegEmbedding(k=5, embedder_name="sergeyzh/rubert-tiny-turbo", db_dir=db_dir)
+
+    utterances = ["hello", "goodbye", "hi", "bye", "bye", "hello", "welcome", "hi123", "hiii", "bye-bye", "bye!"]
+    labels = [0, 1, 0, 1, 1, 0, 0, 0, 0, 1, 1]
+    module.fit(utterances, labels)
+
+    dump_path = Path(dump_dir)
+    dump_path.mkdir(parents=True, exist_ok=True)
+    module.dump(str(dump_path))
+
+    loaded_module = LogRegEmbedding(k=5, embedder_name="sergeyzh/rubert-tiny-turbo", db_dir=db_dir)
+    loaded_module.load(str(dump_path))
+    epsilon = 1e-6
+
+    assert np.allclose(loaded_module.classifier.coef_, module.classifier.coef_, atol=epsilon)
+    assert np.allclose(loaded_module.classifier.intercept_, module.classifier.intercept_, atol=epsilon)
+    assert np.array_equal(np.array(loaded_module.label_encoder.classes_), np.array(module.label_encoder.classes_))
+    assert loaded_module.embedder_name == module.embedder_name
+
+    shutil.rmtree(dump_path)

tests/modules/retrieval/test_retrieval.py

@@ -0,0 +1,32 @@
+import shutil
+from pathlib import Path
+
+from autointent.modules.embedding import RetrievalEmbedding
+from tests.conftest import setup_environment
+
+
+def test_get_assets_returns_correct_artifact():
+    db_dir, dump_dir, logs_dir = setup_environment()
+    module = RetrievalEmbedding(k=5, embedder_name="sergeyzh/rubert-tiny-turbo", db_dir=db_dir)
+    artifact = module.get_assets()
+    assert artifact.embedder_name == "sergeyzh/rubert-tiny-turbo"
+
+
+def test_dump_and_load_preserves_model_state():
+    db_dir, dump_dir, logs_dir = setup_environment()
+    module = RetrievalEmbedding(k=5, embedder_name="sergeyzh/rubert-tiny-turbo", db_dir=db_dir)
+
+    utterances = ["hello", "goodbye", "hi", "bye", "bye", "hello", "welcome", "hi123", "hiii", "bye-bye", "bye!"]
+    labels = [0, 1, 0, 1, 1, 0, 0, 0, 0, 1, 1]
+    module.fit(utterances, labels)
+
+    dump_path = Path(dump_dir)
+    dump_path.mkdir(parents=True, exist_ok=True)
+    module.dump(str(dump_path))
+
+    loaded_module = RetrievalEmbedding(k=5, embedder_name="sergeyzh/rubert-tiny-turbo", db_dir=db_dir)
+    loaded_module.load(str(dump_path))
+
+    assert loaded_module.embedder_name == module.embedder_name
+
+    shutil.rmtree(dump_path)