diff --git a/model2vec/train/pairs.py b/model2vec/train/pairs.py index e1f46fd..0f910bd 100644 --- a/model2vec/train/pairs.py +++ b/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 diff --git a/tests/test_trainable.py b/tests/test_trainable.py index 00eb7a4..d8afdc2 100644 --- a/tests/test_trainable.py +++ b/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)