add comprehensive tests

Author: voorhs

tests/embedder/test_basic.py

@@ -0,0 +1,61 @@
+import numpy as np
+import pytest
+
+from autointent._wrappers.embedder import Embedder
+from autointent.configs._transformers import EmbedderConfig
+
+
+@pytest.fixture
+def simple_embedder_config():
+    """Create a simple embedder config for testing."""
+    return EmbedderConfig(
+        model_name="sergeyzh/rubert-tiny-turbo",
+        batch_size=4,
+        device="cpu",
+        use_cache=False,
+    )
+
+
+def test_embedding_calculation(simple_embedder_config):
+    """Test basic embedding calculation functionality."""
+    embedder = Embedder(simple_embedder_config)
+    test_utterances = ["Hello world", "Test sentence", "Another example"]
+
+    embeddings = embedder.embed(test_utterances)
+
+    assert embeddings.shape[0] == len(test_utterances)
+    assert np.allclose(np.linalg.norm(embeddings, axis=1), 1.0, atol=1e-5)  # normalized
+
+
+def test_embedding_reproducibility(simple_embedder_config):
+    """Test that embeddings are reproducible for same input."""
+    embedder = Embedder(simple_embedder_config)
+    test_utterances = ["Hello world", "Test sentence"]
+
+    embeddings1 = embedder.embed(test_utterances)
+    embeddings2 = embedder.embed(test_utterances)
+
+    np.testing.assert_allclose(embeddings1, embeddings2, rtol=1e-5)
+
+
+def test_single_utterance(simple_embedder_config):
+    """Test embedding calculation for single utterance."""
+    embedder = Embedder(simple_embedder_config)
+
+    embeddings = embedder.embed(["Single test sentence"])
+    assert embeddings.shape[0] == 1
+    assert np.allclose(np.linalg.norm(embeddings[0]), 1.0, atol=1e-5)
+
+
+def test_similarity_symmetry():
+    """Test that similarity is symmetric for cosine similarity."""
+    config = EmbedderConfig(model_name="sergeyzh/rubert-tiny-turbo", similarity_fn_name="cosine", use_cache=False)
+    embedder = Embedder(config)
+
+    utterances = ["Hello world", "Test sentence"]
+    embeddings = embedder.embed(utterances)
+
+    sim1 = embedder.similarity(embeddings[:1], embeddings[1:])
+    sim2 = embedder.similarity(embeddings[1:], embeddings[:1])
+
+    np.testing.assert_allclose(sim1, sim2.T, rtol=1e-5)

tests/embedder/test_caching.py

@@ -0,0 +1,57 @@
+import time
+from unittest.mock import patch
+
+import numpy as np
+
+from autointent._wrappers.embedder import Embedder
+from autointent.configs._transformers import EmbedderConfig
+
+
+def test_caching_enabled():
+    """Test that caching works when enabled."""
+    config = EmbedderConfig(
+        model_name="sergeyzh/rubert-tiny-turbo",
+        use_cache=True,
+        device="cpu",
+    )
+    embedder = Embedder(config)
+    test_utterances = ["Cache test sentence"]
+
+    # Mock the actual embedding calculation to verify caching
+    with patch.object(embedder, "_load_model") as mock_load:
+        mock_model = mock_load.return_value
+        mock_model.encode.return_value = np.array([[0.1, 0.2, 0.3]])
+
+        # First call should trigger model loading
+        start_time = time.time()
+        embeddings1 = embedder.embed(test_utterances)
+        first_call_time = time.time() - start_time
+
+        # Second call should use cache (model.encode shouldn't be called again)
+        start_time = time.time()
+        embeddings2 = embedder.embed(test_utterances)
+        second_call_time = time.time() - start_time
+
+        # Verify results are the same
+        np.testing.assert_allclose(embeddings1, embeddings2, rtol=1e-5)
+
+        assert (
+            second_call_time < first_call_time / 5
+        ), f"Second call ({second_call_time:.4f}s) should be much faster than first call ({first_call_time:.4f}s)"
+
+
+def test_caching_disabled():
+    """Test behavior when caching is disabled."""
+    config = EmbedderConfig(
+        model_name="sergeyzh/rubert-tiny-turbo",
+        use_cache=False,
+        device="cpu",
+    )
+    embedder = Embedder(config)
+    test_utterances = ["No cache test"]
+
+    embeddings1 = embedder.embed(test_utterances)
+    embeddings2 = embedder.embed(test_utterances)
+
+    # Should still be the same since same model/input
+    np.testing.assert_allclose(embeddings1, embeddings2, rtol=1e-5)

tests/embedder/test_dump_load.py

@@ -0,0 +1,81 @@
+import tempfile
+from pathlib import Path
+
+import numpy as np
+from sentence_transformers import SentenceTransformer
+
+from autointent._wrappers.embedder import Embedder
+from autointent.configs._transformers import EmbedderConfig
+
+
+def test_load_from_disk():
+    model = SentenceTransformer("sergeyzh/rubert-tiny-turbo")
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        model.save(str(Path(tmp_dir) / "weights"))
+        embedder = Embedder(EmbedderConfig(model_name=str(Path(tmp_dir) / "weights")))
+        predictions = embedder.embed(["hi!"])
+        embedder.dump(Path(tmp_dir) / "embedder")
+        embedder_loaded = Embedder.load(Path(tmp_dir) / "embedder")
+        predictions_after = embedder_loaded.embed(["hi!"])
+
+    np.testing.assert_almost_equal(predictions_after, predictions, decimal=4)
+
+
+def test_dump_load_cycle():
+    """Test complete dump/load cycle preserves functionality."""
+    original_config = EmbedderConfig(
+        model_name="sergeyzh/rubert-tiny-turbo",
+        default_prompt="Test prompt:",
+        similarity_fn_name="cosine",
+        batch_size=4,
+        use_cache=False,
+    )
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_path = Path(temp_dir)
+
+        # Create and test original embedder
+        embedder_original = Embedder(original_config)
+        test_utterances = ["Test sentence for persistence"]
+        original_embeddings = embedder_original.embed(test_utterances)
+
+        # Dump embedder
+        embedder_original.dump(temp_path)
+
+        # Load embedder
+        embedder_loaded = Embedder.load(temp_path)
+
+        # Test that loaded embedder works the same
+        loaded_embeddings = embedder_loaded.embed(test_utterances)
+        np.testing.assert_allclose(original_embeddings, loaded_embeddings, rtol=1e-5)
+
+        # Test configuration preservation
+        assert embedder_loaded.config.model_name == original_config.model_name
+        assert embedder_loaded.config.default_prompt == original_config.default_prompt
+        assert embedder_loaded.config.similarity_fn_name == original_config.similarity_fn_name
+
+
+def test_load_with_config_override():
+    """Test loading with configuration override."""
+    original_config = EmbedderConfig(
+        model_name="sergeyzh/rubert-tiny-turbo",
+        batch_size=8,
+        use_cache=False,
+    )
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_path = Path(temp_dir)
+
+        # Create and dump original
+        embedder_original = Embedder(original_config)
+        embedder_original.dump(temp_path)
+
+        # Load with override
+        override_config = EmbedderConfig(batch_size=16)
+        embedder_loaded = Embedder.load(temp_path, override_config)
+
+        # Verify override took effect
+        assert embedder_loaded.config.batch_size == 16
+        # Verify original config preserved where not overridden
+        assert embedder_loaded.config.model_name == original_config.model_name

tests/embedder/test_fine_tuning.py tests/embedder/test_fine_tuned.pytests/embedder/test_fine_tuning.py renamed to tests/embedder/test_fine_tuned.py

@@ -22,15 +22,13 @@ def test_model_updates_after_training(dataset):
 
     train_config = EmbedderFineTuningConfig(epoch_num=3, batch_size=8)
     embedder = Embedder(embedder_config)
-    embedder._load_model()
+    embedder._model = embedder._load_model()
 
-    for param in embedder.embedding_model.parameters():
+    for param in embedder._model.parameters():
         assert param.requires_grad, "All trainable parameters should have requires_grad=True"
 
     original_weights = [
-        param.data.detach().cpu().numpy().copy()
-        for param in embedder.embedding_model.parameters()
-        if param.requires_grad
+        param.data.detach().cpu().numpy().copy() for param in embedder._model.parameters() if param.requires_grad
     ]
     embedder.train(
         utterances=data_handler.train_utterances(0)[:1000],
@@ -39,9 +37,7 @@ def test_model_updates_after_training(dataset):
     )
 
     trained_weights = [
-        param.data.detach().cpu().numpy().copy()
-        for param in embedder.embedding_model.parameters()
-        if param.requires_grad
+        param.data.detach().cpu().numpy().copy() for param in embedder._model.parameters() if param.requires_grad
     ]
 
     weights_changed = any(