tests: fix tests

MinishLab · stephantul · Feb 7, 2025 · Dec 22, 2024 · Dec 22, 2024 · Dec 22, 2024
commit 8ab8456f4c37edbefcb4d89910755dabe65c85b7
diff --git a/model2vec/train/classifier.py b/model2vec/train/classifier.py
@@ -2,6 +2,7 @@
 
 import logging
 from collections import Counter
+from tempfile import TemporaryDirectory
 
 import lightning as pl
 import numpy as np
@@ -148,23 +149,26 @@
            val_check_interval = None
            check_val_every_epoch = 1
         else:
             val_check_interval = max(250, 2 * len(val_dataset) // batch_size)
             check_val_every_epoch = None
-        trainer = pl.Trainer(
-            max_epochs=500,
-            callbacks=callbacks,
-            val_check_interval=val_check_interval,
-            check_val_every_n_epoch=check_val_every_epoch,
-            accelerator=device,
-        )
 
-        trainer.fit(
-            c,
-            train_dataloaders=train_dataset.to_dataloader(shuffle=True, batch_size=batch_size),
-            val_dataloaders=val_dataset.to_dataloader(shuffle=False, batch_size=batch_size),
-        )
-        best_model_path = trainer.checkpoint_callback.best_model_path  # type: ignore
-        best_model_weights = torch.load(best_model_path, weights_only=True)
+        with TemporaryDirectory() as tempdir:
+            trainer = pl.Trainer(
+                max_epochs=500,
+                callbacks=callbacks,
+                val_check_interval=val_check_interval,
+                check_val_every_n_epoch=check_val_every_epoch,
+                accelerator=device,
+                default_root_dir=tempdir,
+            )
+
+            trainer.fit(
+                c,
+                train_dataloaders=train_dataset.to_dataloader(shuffle=True, batch_size=batch_size),
+                val_dataloaders=val_dataset.to_dataloader(shuffle=False, batch_size=batch_size),
+            )
+            best_model_path = trainer.checkpoint_callback.best_model_path  # type: ignore
+            best_model_weights = torch.load(best_model_path, weights_only=True)
 
         state_dict = {}
         for weight_name, weight in best_model_weights["state_dict"].items():
@@ -181,7 +185,7 @@
        self.classes_ = classes

        if len(self.classes) != self.out_dim:
            self.out_dim = len(self.classes)

        self.head = self.construct_head()
        self.embeddings = nn.Embedding.from_pretrained(self.vectors.clone(), freeze=False, padding_idx=self.pad_id)
@@ -228,7 +232,14 @@
         for index, layer in enumerate([module for module in self.head if isinstance(module, nn.Linear)]):
             mlp_head.coefs_[index] = layer.weight.detach().cpu().numpy().T
             mlp_head.intercepts_[index] = layer.bias.detach().cpu().numpy()
+        # Below is necessary to ensure that the converted model works correctly.
+        # In scikit-learn, a binary classifier only has a single vector of output coefficients
+        # and a single intercept. We use two output vectors.
+        # To convert correctly, we need to set the outputs correctly, and fix the activation function.
+        # Make sure n_outputs is set to > 1.
         mlp_head.n_outputs_ = self.out_dim
+        # Set to softmax
+        mlp_head.out_activation_ = "softmax"
 
         return StaticModelPipeline(static_model, converted)
 
@@ -242,7 +253,7 @@

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Simple forward pass."""
        return self.model(x)

    def training_step(self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> torch.Tensor:
        """Simple training step using cross entropy loss."""

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -20,7 +20,7 @@
 @pytest.fixture(scope="session")
 def mock_tokenizer() -> Tokenizer:
     """Create a mock tokenizer."""
-    vocab = ["word1", "word2", "word3", "[UNK]", "[PAD]"]
+    vocab = ["[PAD]", "word1", "word2", "word3", "[UNK]"]
     unk_token = "[UNK]"
 
     model = WordLevel(vocab={word: idx for idx, word in enumerate(vocab)}, unk_token=unk_token)
@@ -81,21 +81,18 @@ def mock_config() -> dict[str, str]:
 
 
 @pytest.fixture(scope="session")
-def mock_inference_pipeline(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> StaticModelPipeline:
+def mock_inference_pipeline(mock_trained_pipeline: StaticModelForClassification) -> StaticModelPipeline:
     """Mock pipeline."""
-    encoder = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer, config={})
-    encoded = encoder.encode(["dog", "cat"])
-    labels = ["a", "b"]
-    head = make_pipeline(MLPClassifier(random_state=12)).fit(encoded, labels)
-
-    return StaticModelPipeline(encoder, head=head)
+    return mock_trained_pipeline.to_pipeline()
 
 
 @pytest.fixture(scope="session")
-def mock_trained_pipeline(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> StaticModelForClassification:
+def mock_trained_pipeline() -> StaticModelForClassification:
     """Mock staticmodelforclassification."""
-    vectors_torched = torch.from_numpy(mock_vectors).float()
-    s = StaticModelForClassification(vectors=vectors_torched, tokenizer=mock_tokenizer).to("cpu")
+    tokenizer = AutoTokenizer.from_pretrained("tests/data/test_tokenizer").backend_tokenizer
+    torch.random.manual_seed(42)
+    vectors_torched = torch.randn(len(tokenizer.get_vocab()), 12)
+    s = StaticModelForClassification(vectors=vectors_torched, tokenizer=tokenizer, hidden_dim=12).to("cpu")
     s.fit(["dog", "cat"], ["a", "b"], device="cpu")
 
     return s
diff --git a/tests/test_inference.py b/tests/test_inference.py
@@ -10,25 +10,25 @@
 
 def test_init_predict(mock_inference_pipeline: StaticModelPipeline) -> None:
     """Test successful initialization of StaticModelPipeline."""
-    assert mock_inference_pipeline.predict("dog").tolist() == ["a"]
-    assert mock_inference_pipeline.predict(["dog"]).tolist() == ["a"]
+    assert mock_inference_pipeline.predict("dog").tolist() == ["b"]
+    assert mock_inference_pipeline.predict(["dog"]).tolist() == ["b"]
 
 
 def test_init_predict_proba(mock_inference_pipeline: StaticModelPipeline) -> None:
     """Test successful initialization of StaticModelPipeline."""
-    assert mock_inference_pipeline.predict_proba("dog").argmax() == 0
-    assert mock_inference_pipeline.predict_proba(["dog"]).argmax(1).tolist() == [0]
+    assert mock_inference_pipeline.predict_proba("dog").argmax() == 1
+    assert mock_inference_pipeline.predict_proba(["dog"]).argmax(1).tolist() == [1]
 
 
 def test_roundtrip_save(mock_inference_pipeline: StaticModelPipeline) -> None:
     """Test saving and loading the pipeline."""
     with TemporaryDirectory() as temp_dir:
         mock_inference_pipeline.save_pretrained(temp_dir)
         loaded = StaticModelPipeline.from_pretrained(temp_dir)
-        assert loaded.predict("dog") == ["a"]
-        assert loaded.predict(["dog"]) == ["a"]
-        assert loaded.predict_proba("dog").argmax() == 0
-        assert loaded.predict_proba(["dog"]).argmax(1).tolist() == [0]
+        assert loaded.predict("dog") == ["b"]
+        assert loaded.predict(["dog"]) == ["b"]
+        assert loaded.predict_proba("dog").argmax() == 1
+        assert loaded.predict_proba(["dog"]).argmax(1).tolist() == [1]
 
 
 @patch("model2vec.inference.model._DEFAULT_TRUST_PATTERN", re.compile("torch"))

diff --git a/tests/test_trainable.py b/tests/test_trainable.py
@@ -69,7 +69,7 @@ def test_init_classifier_from_model(mock_vectors: np.ndarray, mock_tokenizer: To
 def test_encode(mock_trained_pipeline: StaticModelForClassification) -> None:
     """Test the encode function."""
     result = mock_trained_pipeline._encode(torch.tensor([[0, 1], [1, 0]]).long())
-    assert result.shape == (2, 2)
+    assert result.shape == (2, 12)
     assert torch.allclose(result[0], result[1])
 
 
@@ -111,23 +111,29 @@ def test_textdataset_init_incorrect() -> None:
 
 def test_predict(mock_trained_pipeline: StaticModelForClassification) -> None:
     """Test the predict function."""
-    result = mock_trained_pipeline.predict(["dog dog", "cat"]).tolist()
-    assert result == ["a", "a"]
+    result = mock_trained_pipeline.predict(["dog cat", "dog"]).tolist()
+    assert result == ["b", "b"]
 
 
 def test_predict_proba(mock_trained_pipeline: StaticModelForClassification) -> None:
     """Test the predict function."""
-    result = mock_trained_pipeline.predict_proba(["dog dog", "cat"])
+    result = mock_trained_pipeline.predict_proba(["dog cat", "dog"])
     assert result.shape == (2, 2)
 
 
 def test_convert_to_pipeline(mock_trained_pipeline: StaticModelForClassification) -> None:
     """Convert a model to a pipeline."""
     mock_trained_pipeline.eval()
     pipeline = mock_trained_pipeline.to_pipeline()
-    a = pipeline.predict(["dog dog", "cat"]).tolist()
-    b = mock_trained_pipeline.predict(["dog dog", "cat"]).tolist()
+    encoded_pipeline = pipeline.model.encode(["dog cat", "dog"])
+    encoded_model = mock_trained_pipeline(mock_trained_pipeline.tokenize(["dog cat", "dog"]))[1].detach().numpy()
+    assert np.allclose(encoded_pipeline, encoded_model)
+    a = pipeline.predict(["dog cat", "dog"]).tolist()
+    b = mock_trained_pipeline.predict(["dog cat", "dog"]).tolist()
     assert a == b
+    p1 = pipeline.predict_proba(["dog cat", "dog"])
+    p2 = mock_trained_pipeline.predict_proba(["dog cat", "dog"])
+    assert np.allclose(p1, p2)
 
 
 def test_train_test_split() -> None: