fix: add token dropout to pair trainer (#393)

Author: stephantul

model2vec/train/pairs.py

@@ -279,6 +279,7 @@ def fit(
         validation_steps: int | None = None,
         random_seed: int = DEFAULT_RANDOM_SEED,
         temperature: float = 0.05,
+        token_dropout: float = 0.0,
     ) -> T:
         """Fit a model that embeds paired texts close together.
 
@@ -313,6 +314,8 @@ def fit(
         :param validation_steps: The number of steps to run validation for. If None, validation steps are estimated from the data.
         :param random_seed: The random seed to use. Defaults to 42.
         :param temperature: The temperature of the InfoNCE loss.
+        :param token_dropout: The fraction of tokens to randomly drop from each training text.
+            Has no effect during validation. Must be in the range [0, 1).
         :return: The fitted model.
         :raises ValueError: If `batch_size` is smaller than 2.
         """
@@ -339,6 +342,7 @@ def fit(
             max_epochs=max_epochs,
             device=device,
             validation_steps=validation_steps,
+            token_dropout=token_dropout,
         )
 
         return self

tests/test_trainable.py

@@ -833,6 +833,19 @@ def test_pair_similarity_fit_rejects_invalid_temperature_and_batch_size(
         model.fit(text_a, text_b, test_size=0.5, batch_size=1)
 
 
+def test_pair_similarity_fit_token_dropout(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None:
+    """fit() passes token_dropout to training and rejects values outside [0, 1)."""
+    model = StaticModelForPairSimilarity(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer)
+    text_a, text_b = ["word1", "word2", "word3", "word1 word2"], ["word2", "word3", "word1", "word3"]
+    with pytest.raises(ValueError, match="token_dropout"):
+        model.fit(text_a, text_b, test_size=0.5, token_dropout=1.0)
+    with pytest.raises(ValueError, match="token_dropout"):
+        model.fit(text_a, text_b, test_size=0.5, token_dropout=-0.1)
+
+    model.fit(text_a, text_b, test_size=0.5, token_dropout=0.3, max_epochs=1)
+    assert model.token_dropout == 0.3
+
+
 def test_pairdataset_drops_single_pair_batches() -> None:
     """A final batch with a single pair is dropped, unless it is the only pair."""
     dataset = PairDataset(ColumnRows(text_a=[[1], [2], [3]], text_b=[[1], [2], [3]]), _pretokenized)