From e935ad76e09f38508cc6e968b933c392739f5838 Mon Sep 17 00:00:00 2001 From: stephantul Date: Thu, 17 Sep 2026 07:57:33 +0200 Subject: [PATCH 1/2] feat: add pairwise trainer --- model2vec/train/README.md | 21 ++++ model2vec/train/__init__.py | 8 +- model2vec/train/base.py | 29 +++-- model2vec/train/dataset.py | 54 +++++++++ model2vec/train/pairs.py | 226 ++++++++++++++++++++++++++++++++++++ model2vec/train/utils.py | 9 +- tests/conftest.py | 22 +++- tests/test_trainable.py | 146 ++++++++++++++++++++++- 8 files changed, 499 insertions(+), 16 deletions(-) create mode 100644 model2vec/train/pairs.py diff --git a/model2vec/train/README.md b/model2vec/train/README.md index e5aab12..98342f4 100644 --- a/model2vec/train/README.md +++ b/model2vec/train/README.md @@ -98,6 +98,27 @@ print(classification_report) The scores are competitive with the popular [roberta-base-go_emotions](https://huggingface.co/SamLowe/roberta-base-go_emotions) model, while our model is orders of magnitude faster. +## Pair similarity + +`StaticModelForPairSimilarity` trains a model to embed pairs of related texts (e.g. queries and their matching documents) close together, by encoding both sides with the same model and minimizing the cosine distance between them: + +```python +from model2vec.train import StaticModelForPairSimilarity + +model = StaticModelForPairSimilarity.from_pretrained(model_name="minishlab/potion-base-32M") +model.fit(text_a=["how tall is the eiffel tower?"], text_b=["the eiffel tower is 330 meters tall."]) +``` + +Pairs can also be labeled: pairs labeled `1` are pushed together (cosine similarity towards 1), while pairs labeled `0` are pushed towards a cosine similarity of 0. If `labels` is omitted, every pair is treated as positive: + +```python +model.fit( + text_a=["how tall is the eiffel tower?", "how tall is the eiffel tower?"], + text_b=["the eiffel tower is 330 meters tall.", "paris is the capital of france."], + labels=[1, 0], +) +``` + # Persistence You can turn a classifier into a lightweight inference pipeline, as follows: diff --git a/model2vec/train/__init__.py b/model2vec/train/__init__.py index c97a1f9..7a93ba1 100644 --- a/model2vec/train/__init__.py +++ b/model2vec/train/__init__.py @@ -6,7 +6,13 @@ importable(extra_dependency, _REQUIRED_EXTRA) from model2vec.train.classifier import StaticModelForClassification +from model2vec.train.pairs import StaticModelForPairSimilarity from model2vec.train.regression import StaticModelForRegression from model2vec.train.similarity import StaticModelForSimilarity -__all__ = ["StaticModelForClassification", "StaticModelForSimilarity", "StaticModelForRegression"] +__all__ = [ + "StaticModelForClassification", + "StaticModelForSimilarity", + "StaticModelForRegression", + "StaticModelForPairSimilarity", +] diff --git a/model2vec/train/base.py b/model2vec/train/base.py index 9c5eaa1..2c1e85e 100644 --- a/model2vec/train/base.py +++ b/model2vec/train/base.py @@ -13,7 +13,7 @@ from model2vec.inference import StaticModelPipeline from model2vec.model import DEFAULT_MAX_LENGTH, PathLike, StaticModel, _get_unk_token_id -from model2vec.train.dataset import TextDataset +from model2vec.train.dataset import PairDataset, TextDataset from model2vec.train.trainer import MetricsFn, default_metrics, resolve_device, run_training_loop from model2vec.train.utils import ( get_probable_pad_token_id, @@ -332,8 +332,8 @@ def _train( self, loss_function: nn.Module, learning_rate: float, - train_dataset: TextDataset, - val_dataset: TextDataset, + train_dataset: TextDataset | PairDataset, + val_dataset: TextDataset | PairDataset, batch_size: int, early_stopping_patience: int | None, min_epochs: int | None, @@ -391,17 +391,16 @@ def _determine_val_check_interval( return val_check_interval, check_val_every_epoch - def _prepare_dataset(self, X: list[str], y: torch.Tensor, max_length: int | None) -> TextDataset: - """Prepare a dataset. + def _tokenize_texts(self, X: list[str], max_length: int | None) -> list[list[int]]: + """Tokenize a list of texts into lists of token ids. - :param X: The texts. - :param y: The labels. + :param X: The texts to tokenize. :param max_length: The maximum length of the input in tokens. If this is None, no truncation is done. - :return: A TextDataset. + :return: The tokenized texts. """ batch_size = 1024 tokenized: list[list[int]] = [] - for batch_idx in trange(0, len(X), 1024, desc="Tokenizing data"): + for batch_idx in trange(0, len(X), batch_size, desc="Tokenizing data"): batch = X[batch_idx : batch_idx + batch_size] if max_length is not None: truncate_length = max_length * 10 @@ -409,7 +408,17 @@ def _prepare_dataset(self, X: list[str], y: torch.Tensor, max_length: int | None encoded = self.tokenizer.encode_batch_fast(batch, add_special_tokens=False) tokenized.extend([self._remove_unk(encoding.ids)[:max_length] for encoding in encoded]) - return TextDataset(tokenized, y, pad_id=self.pad_id) + return tokenized + + def _prepare_dataset(self, X: list[str], y: torch.Tensor, max_length: int | None) -> TextDataset: + """Prepare a dataset. + + :param X: The texts. + :param y: The labels. + :param max_length: The maximum length of the input in tokens. If this is None, no truncation is done. + :return: A TextDataset. + """ + return TextDataset(self._tokenize_texts(X, max_length), y, pad_id=self.pad_id) def _labels_to_tensor(self, labels: Any) -> torch.Tensor: """Turn the labels into a tensor.""" diff --git a/model2vec/train/dataset.py b/model2vec/train/dataset.py index bad58fa..0d7627e 100644 --- a/model2vec/train/dataset.py +++ b/model2vec/train/dataset.py @@ -38,3 +38,57 @@ def collate_fn(self, batch: list[tuple[list[list[int]], int]]) -> tuple[torch.Te def to_dataloader(self, shuffle: bool, batch_size: int = 32) -> DataLoader: """Convert the dataset to a DataLoader.""" return DataLoader(self, collate_fn=self.collate_fn, shuffle=shuffle, batch_size=batch_size) + + +class PairDataset(Dataset): + def __init__( + self, + tokenized_texts_a: list[list[int]], + tokenized_texts_b: list[list[int]], + labels: list[int] | torch.Tensor | None = None, + pad_id: int = 0, + ) -> None: + """A dataset of aligned text pairs. + + :param tokenized_texts_a: The tokenized first half of each pair. Each text is a list of token ids. + :param tokenized_texts_b: The tokenized second half of each pair. Each text is a list of token ids. + :param labels: The label for each pair: 1 if the pair should be pushed together, 0 if it should be + pushed towards a cosine similarity of 0. If None, every pair is labeled 1. + :param pad_id: The id used to pad batches. Must match the `pad_id` of the model being trained. + :raises ValueError: If the two halves don't have the same number of texts, or if `labels` doesn't + have one entry per pair. + """ + if len(tokenized_texts_a) != len(tokenized_texts_b): + raise ValueError("The two halves of a pair dataset must have the same number of texts.") + if labels is not None and len(labels) != len(tokenized_texts_a): + raise ValueError("labels must have one entry per pair.") + self.tokenized_texts_a = tokenized_texts_a + self.tokenized_texts_b = tokenized_texts_b + self.labels = torch.ones(len(tokenized_texts_a)) if labels is None else torch.as_tensor(labels).float() + self.pad_id = pad_id + + def __len__(self) -> int: + """Return the length of the dataset.""" + return len(self.tokenized_texts_a) + + def __getitem__(self, index: int) -> tuple[list[int], list[int], torch.Tensor]: + """Gets an item.""" + return self.tokenized_texts_a[index], self.tokenized_texts_b[index], self.labels[index] + + def collate_fn(self, batch: list[tuple[list[int], list[int], torch.Tensor]]) -> tuple[torch.Tensor, torch.Tensor]: + """Collate function. + + Both halves are padded together so they end up with the same sequence length, then + stacked into a single (2, batch_size, seq_len) tensor. + """ + texts_a, texts_b, labels = zip(*batch) + + tensors: list[torch.Tensor] = [torch.LongTensor(x) for x in (*texts_a, *texts_b)] + padded = pad_sequence(tensors, batch_first=True, padding_value=self.pad_id) + padded_a, padded_b = padded[: len(texts_a)], padded[len(texts_a) :] + + return torch.stack([padded_a, padded_b]), torch.stack(labels) + + def to_dataloader(self, shuffle: bool, batch_size: int = 32) -> DataLoader: + """Convert the dataset to a DataLoader.""" + return DataLoader(self, collate_fn=self.collate_fn, shuffle=shuffle, batch_size=batch_size) diff --git a/model2vec/train/pairs.py b/model2vec/train/pairs.py new file mode 100644 index 0000000..07f9254 --- /dev/null +++ b/model2vec/train/pairs.py @@ -0,0 +1,226 @@ +from __future__ import annotations + +import logging +from typing import TypeVar + +import torch +from tokenizers import Tokenizer +from torch import nn + +from model2vec.model import DEFAULT_MAX_LENGTH +from model2vec.train.base import BaseFinetuneable +from model2vec.train.dataset import PairDataset +from model2vec.train.utils import DEFAULT_RANDOM_SEED, seed_everything, train_test_split + +logger = logging.getLogger(__name__) + + +class PairCosineLoss(nn.Module): + def __call__(self, head_out: tuple[torch.Tensor, torch.Tensor], y: torch.Tensor) -> torch.Tensor: + """Returns the cosine loss between the two encoded halves of a pair batch, per pair label. + + Pairs labeled 1 are pushed towards a cosine similarity of 1, pairs labeled 0 are pushed + towards a cosine similarity of 0. + """ + out_a, out_b = head_out + out_a = torch.nn.functional.normalize(out_a, dim=1) + out_b = torch.nn.functional.normalize(out_b, dim=1) + cosine_sim = torch.sum(out_a * out_b, dim=1) + loss = torch.where(y == 1, 1 - cosine_sim, cosine_sim.abs()) + return loss.mean() + + +class StaticModelForPairSimilarity(BaseFinetuneable): + val_metric = "val_loss" + early_stopping_direction = "min" + + def __init__( + self, + *, + vectors: torch.Tensor, + tokenizer: Tokenizer, + n_layers: int = 1, + hidden_dim: int = 512, + out_dim: int | None = None, + pad_id: int = 0, + token_mapping: list[int] | None = None, + weights: torch.Tensor | None = None, + freeze: bool = False, + normalize: bool = True, + freeze_weights: bool = False, + max_length: int | None = DEFAULT_MAX_LENGTH, + ) -> None: + """Initialize a model that is trained to embed pairs of texts close together. + + :param vectors: The embeddings of the staticmodel. + :param tokenizer: The tokenizer. + :param n_layers: The number of layers in the head. + :param hidden_dim: The hidden dimension of the head. + :param out_dim: The output embedding dimension. If None, defaults to the input embedding dimension. + :param pad_id: The padding id. This is set to 0 in almost all model2vec models. + :param token_mapping: The token mapping. If None, the token mapping is set to the range of the number of vectors. + :param weights: The weights of the model. If None, the weights are initialized to zeros. + :param freeze: Whether to freeze the embeddings. This should be set to False in most cases. + :param normalize: Whether to normalize the embeddings. + :param freeze_weights: Whether to freeze the learned token weights. + :param max_length: The default maximum sequence length (in tokens) used to tokenize inputs. + """ + super().__init__( + vectors=vectors, + out_dim=out_dim if out_dim is not None else vectors.shape[1], + pad_id=pad_id, + tokenizer=tokenizer, + token_mapping=token_mapping, + weights=weights, + freeze=freeze, + hidden_dim=hidden_dim, + n_layers=n_layers, + normalize=normalize, + freeze_weights=freeze_weights, + max_length=max_length, + ) + + def forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: # type: ignore[override] + """Encode both halves of a pair batch through the shared embeddings and head. + + :param input_ids: A `(2, batch_size, seq_len)` tensor, stacking the two padded text sets. + :return: The head outputs for the first and second set of texts. + """ + out_a = self.head(self._encode(input_ids[0])) + out_b = self.head(self._encode(input_ids[1])) + return out_a, out_b + + def _check_pair_val_split( + self, + text_a: list[str], + text_b: list[str], + labels: list[int], + text_a_val: list[str] | None, + text_b_val: list[str] | None, + labels_val: list[int] | None, + test_size: float, + ) -> tuple[list[str], list[str], list[str], list[str], list[int], list[int]]: + if len(text_a) != len(text_b): + raise ValueError("text_a and text_b must have the same length.") + if len(labels) != len(text_a): + raise ValueError("labels must have the same length as text_a and text_b.") + if (text_a_val is not None) != (text_b_val is not None): + raise ValueError("Both text_a_val and text_b_val must be provided together, or neither.") + + if text_a_val is not None and text_b_val is not None: + if len(text_a_val) != len(text_b_val): + raise ValueError("text_a_val and text_b_val must have the same length.") + labels_val = [1] * len(text_a_val) if labels_val is None else labels_val + if len(labels_val) != len(text_a_val): + raise ValueError("labels_val must have the same length as text_a_val and text_b_val.") + return text_a, text_a_val, text_b, text_b_val, labels, labels_val + + pairs = list(zip(text_a, text_b)) + train_pairs, val_pairs, train_labels, val_labels = train_test_split(pairs, labels, test_size=test_size) + train_a, train_b = map(list, zip(*train_pairs)) if train_pairs else ([], []) + val_a, val_b = map(list, zip(*val_pairs)) if val_pairs else ([], []) + return train_a, val_a, train_b, val_b, train_labels, val_labels + + def _prepare_pair_dataset( + self, text_a: list[str], text_b: list[str], labels: list[int], max_length: int | None + ) -> PairDataset: + """Tokenize both halves of a pair dataset. + + :param text_a: The first half of each pair. + :param text_b: The second half of each pair. + :param labels: The label for each pair. + :param max_length: The maximum length of the input in tokens. If this is None, no truncation is done. + :return: A PairDataset. + """ + return PairDataset( + self._tokenize_texts(text_a, max_length), + self._tokenize_texts(text_b, max_length), + labels=labels, + pad_id=self.pad_id, + ) + + def fit( + self: T, + text_a: list[str], + text_b: list[str], + labels: list[int] | None = None, + learning_rate: float = 1e-3, + batch_size: int | None = None, + min_epochs: int | None = None, + max_epochs: int | None = -1, + early_stopping_patience: int | None = 5, + test_size: float = 0.1, + device: str = "auto", + text_a_val: list[str] | None = None, + text_b_val: list[str] | None = None, + labels_val: list[int] | None = None, + validation_steps: int | None = None, + random_seed: int = DEFAULT_RANDOM_SEED, + ) -> T: + """Fit a model that maximizes the cosine similarity between paired texts. + + This function trains the model with a plain torch training loop. Both `text_a` and `text_b` + are encoded with the same model. Pairs labeled 1 are pushed together, minimizing the cosine + distance between them. Pairs labeled 0 are pushed towards a cosine similarity of 0. We use + early stopping. After training, the weights of the best model are loaded back into the model. + + This function seeds everything with a seed of 42, so the results are reproducible. + It also splits the data into a train and validation set, again with a random seed. + + If `text_a_val` and `text_b_val` are not provided, the function will automatically + split the training data into a train and validation set using `test_size`. + + :param text_a: The first half of each training pair. + :param text_b: The second half of each training pair. + :param labels: The label for each training pair: 1 if the pair should be pushed together, 0 if + it should be pushed towards a cosine similarity of 0. If None, every pair is labeled 1. + :param learning_rate: The learning rate. + :param batch_size: The batch size. If None, a good batch size is chosen automatically. + :param min_epochs: The minimum number of epochs to train for. + :param max_epochs: The maximum number of epochs to train for. + If this is -1, the model trains until early stopping is triggered. + :param early_stopping_patience: The patience for early stopping. + If this is None, early stopping is disabled. + :param test_size: The test size for the train-test split. + :param device: The device to train on. If this is "auto", the device is chosen automatically. + :param text_a_val: The first half of each validation pair. + :param text_b_val: The second half of each validation pair. + :param labels_val: The label for each validation pair. If None, every validation pair is labeled 1. + :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. + :return: The fitted model. + """ + seed_everything(random_seed) + logger.info("Re-initializing model.") + + labels = [1] * len(text_a) if labels is None else labels + + train_a, val_a, train_b, val_b, train_labels, val_labels = self._check_pair_val_split( + text_a, text_b, labels, text_a_val, text_b_val, labels_val, test_size + ) + self._initialize() + + logger.info("Preparing train dataset.") + train_dataset = self._prepare_pair_dataset(train_a, train_b, train_labels, self.max_length) + logger.info("Preparing validation dataset.") + val_dataset = self._prepare_pair_dataset(val_a, val_b, val_labels, self.max_length) + + batch_size = self._determine_batch_size(batch_size, len(train_dataset)) + + self._train( + loss_function=PairCosineLoss(), + learning_rate=learning_rate, + train_dataset=train_dataset, + val_dataset=val_dataset, + batch_size=batch_size, + early_stopping_patience=early_stopping_patience, + min_epochs=min_epochs, + max_epochs=max_epochs, + device=device, + validation_steps=validation_steps, + ) + + return self + + +T = TypeVar("T", bound=StaticModelForPairSimilarity) diff --git a/model2vec/train/utils.py b/model2vec/train/utils.py index e887728..a0e31fa 100644 --- a/model2vec/train/utils.py +++ b/model2vec/train/utils.py @@ -3,7 +3,7 @@ import logging import random from collections import Counter, defaultdict -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, TypeVar import numpy as np import torch @@ -69,11 +69,14 @@ def _index(sequence: Any, indices: list[int]) -> Any: return sequence[indices] +X_co = TypeVar("X_co") + + def train_test_split( - X: list[str], + X: list[X_co], y: list, test_size: float, -) -> tuple[list[str], list[str], list, list]: +) -> tuple[list[X_co], list[X_co], list, list]: """Split the data. For single-label classification, stratification is attempted (if possible). diff --git a/tests/conftest.py b/tests/conftest.py index f2d6e0a..e1fd172 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -16,7 +16,12 @@ from model2vec.inference import StaticModelPipeline from model2vec.model import StaticModel -from model2vec.train import StaticModelForClassification, StaticModelForRegression, StaticModelForSimilarity +from model2vec.train import ( + StaticModelForClassification, + StaticModelForPairSimilarity, + StaticModelForRegression, + StaticModelForSimilarity, +) _TOKENIZER_TYPES = ["wordpiece", "bpe", "unigram"] @@ -233,6 +238,21 @@ def mock_trained_regression_pipeline() -> StaticModelForRegression: return model +@pytest.fixture(scope="session") +def mock_trained_pair_similarity_pipeline() -> StaticModelForPairSimilarity: + """Mock StaticModelForPairSimilarity.""" + tokenizer = AutoTokenizer.from_pretrained("tests/data/test_tokenizer").backend_tokenizer + torch.random.manual_seed(42) + vectors_torched = torch.randn(len(tokenizer.get_vocab()), 12) + model = StaticModelForPairSimilarity(vectors=vectors_torched, tokenizer=tokenizer, hidden_dim=12).to("cpu") + + text_a = ["dog", "cat"] + text_b = ["puppy", "kitten"] + model.fit(text_a, text_b) + + return model + + @pytest.fixture(scope="session") def mock_inference_pipeline_projector( mock_trained_similarity_pipeline: StaticModelForSimilarity, diff --git a/tests/test_trainable.py b/tests/test_trainable.py index fb61d73..83ba1ff 100644 --- a/tests/test_trainable.py +++ b/tests/test_trainable.py @@ -15,7 +15,8 @@ from model2vec.model import StaticModel from model2vec.train import StaticModelForClassification from model2vec.train.base import BaseFinetuneable -from model2vec.train.dataset import TextDataset +from model2vec.train.dataset import PairDataset, TextDataset +from model2vec.train.pairs import PairCosineLoss, StaticModelForPairSimilarity from model2vec.train.regression import StaticModelForRegression from model2vec.train.similarity import StaticModelForSimilarity from model2vec.train.trainer import _resolve_max_epochs, resolve_device, run_training_loop @@ -259,6 +260,149 @@ def test_convert_to_pipeline_regression(mock_trained_regression_pipeline: Static assert np.allclose(p1, p2, rtol=1e-5, atol=1e-4) +def test_pairdataset_init() -> None: + """Test the pair dataset init.""" + dataset = PairDataset([[0], [1]], [[2], [3]]) + assert len(dataset) == 2 + + +def test_pairdataset_init_incorrect() -> None: + """Test the pair dataset init with mismatched lengths.""" + with pytest.raises(ValueError): + PairDataset([[0]], [[2], [3]]) + + +def test_pairdataset_collate() -> None: + """Batches should stack the two padded halves into a single (2, batch, seq_len) tensor.""" + dataset = PairDataset([[1], [1, 2]], [[1, 2, 3], [1]], pad_id=0) + batch, y = next(iter(dataset.to_dataloader(shuffle=False, batch_size=2))) + assert batch.shape == (2, 2, 3) + assert y.shape == (2,) + assert torch.equal(batch[0], torch.tensor([[1, 0, 0], [1, 2, 0]])) + assert torch.equal(batch[1], torch.tensor([[1, 2, 3], [1, 0, 0]])) + + +def test_pairdataset_default_labels_are_positive() -> None: + """Without explicit labels, every pair defaults to label 1.""" + dataset = PairDataset([[1], [2]], [[3], [4]]) + assert torch.equal(dataset.labels, torch.tensor([1.0, 1.0])) + + +def test_pairdataset_custom_labels() -> None: + """Custom labels are stored and returned by the collate function.""" + dataset = PairDataset([[1], [2]], [[3], [4]], labels=[1, 0]) + _, y = next(iter(dataset.to_dataloader(shuffle=False, batch_size=2))) + assert torch.equal(y, torch.tensor([1.0, 0.0])) + + +def test_pairdataset_labels_mismatched_length() -> None: + """Labels must have one entry per pair.""" + with pytest.raises(ValueError): + PairDataset([[1], [2]], [[3], [4]], labels=[1]) + + +def test_pair_cosine_loss_pushes_towards_label() -> None: + """Label 1 pairs are pushed towards a cosine similarity of 1, label 0 pairs towards 0.""" + loss_fn = PairCosineLoss() + out_a = torch.tensor([[1.0, 0.0], [1.0, 0.0]]) + + identical = torch.tensor([[1.0, 0.0], [1.0, 0.0]]) + orthogonal = torch.tensor([[0.0, 1.0], [0.0, 1.0]]) + + assert loss_fn((out_a, identical), torch.tensor([1.0, 1.0])).item() == pytest.approx(0.0, abs=1e-6) + assert loss_fn((out_a, orthogonal), torch.tensor([1.0, 1.0])).item() == pytest.approx(1.0) + assert loss_fn((out_a, orthogonal), torch.tensor([0.0, 0.0])).item() == pytest.approx(0.0, abs=1e-6) + assert loss_fn((out_a, identical), torch.tensor([0.0, 0.0])).item() == pytest.approx(1.0) + + +def test_pair_similarity_out_dim_defaults_to_embed_dim(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """The output dimension defaults to the input embedding dimension when not specified.""" + s = StaticModelForPairSimilarity(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + assert s.out_dim == mock_vectors.shape[1] + + s = StaticModelForPairSimilarity( + vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer, out_dim=7 + ) + assert s.out_dim == 7 + + +def test_pair_similarity_forward(mock_trained_pair_similarity_pipeline: StaticModelForPairSimilarity) -> None: + """The forward pass should return one head output per half of the pair batch.""" + model = mock_trained_pair_similarity_pipeline + dataset = model._prepare_pair_dataset(["dog cat", "dog"], ["puppy", "kitten cat"], [1, 1], max_length=None) + batch, _ = next(iter(dataset.to_dataloader(shuffle=False, batch_size=2))) + + with torch.no_grad(): + out_a, out_b = model(batch) + assert out_a.shape == (2, model.out_dim) + assert out_b.shape == (2, model.out_dim) + + +def test_pair_similarity_mismatched_lengths( + mock_trained_pair_similarity_pipeline: StaticModelForPairSimilarity, +) -> None: + """text_a and text_b must have the same length.""" + with pytest.raises(ValueError): + mock_trained_pair_similarity_pipeline.fit(["dog", "cat"], ["puppy"]) + + +def test_pair_similarity_val_split_errors(mock_trained_pair_similarity_pipeline: StaticModelForPairSimilarity) -> None: + """Both validation halves must be provided together, or neither.""" + with pytest.raises(ValueError): + mock_trained_pair_similarity_pipeline.fit( + ["dog", "cat"], ["puppy", "kitten"], text_a_val=["dog"], text_b_val=None + ) + with pytest.raises(ValueError): + mock_trained_pair_similarity_pipeline.fit( + ["dog", "cat"], ["puppy", "kitten"], text_a_val=["dog", "cat"], text_b_val=["puppy"] + ) + + +def test_pair_similarity_labels_mismatched_length( + mock_trained_pair_similarity_pipeline: StaticModelForPairSimilarity, +) -> None: + """Labels must have one entry per training pair, and labels_val one entry per validation pair.""" + with pytest.raises(ValueError): + mock_trained_pair_similarity_pipeline.fit(["dog", "cat"], ["puppy", "kitten"], labels=[1]) + with pytest.raises(ValueError): + mock_trained_pair_similarity_pipeline.fit( + ["dog", "cat"], + ["puppy", "kitten"], + text_a_val=["dog"], + text_b_val=["puppy"], + labels_val=[1, 0], + ) + + +def test_pair_similarity_fit_with_labels(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """A model can be fit with a mix of positive and negative pair labels.""" + model = StaticModelForPairSimilarity(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + text_a = ["word1", "word2", "word3", "word1 word2"] + text_b = ["word2", "word3", "word1", "word3 word1"] + labels = [1, 1, 0, 0] + model.fit(text_a, text_b, labels=labels, early_stopping_patience=1, max_epochs=1) + + +def test_convert_to_pipeline_pair_similarity( + mock_trained_pair_similarity_pipeline: StaticModelForPairSimilarity, +) -> None: + """Convert a model to a pipeline.""" + mock_trained_pair_similarity_pipeline.eval() + pipeline = mock_trained_pair_similarity_pipeline.to_pipeline() + encoded_pipeline = pipeline.model.encode(["dog cat", "dog"]) + encoded_model = ( + mock_trained_pair_similarity_pipeline._encode( + mock_trained_pair_similarity_pipeline.tokenize(["dog cat", "dog"]) + ) + .detach() + .numpy() + ) + assert np.allclose(encoded_pipeline, encoded_model) + p1 = pipeline.predict(["dog cat", "dog"]) + p2 = mock_trained_pair_similarity_pipeline.encode(["dog cat", "dog"]) + assert np.allclose(p1, p2, rtol=1e-5, atol=1e-4) + + def test_train_test_split() -> None: """Test the train test split function.""" a, b, c, d = train_test_split(["0", "1", "2", "3"], ["1", "1", "0", "0"], 0.5) From 090854ed80f6446563ed3c6645f13ea2fec39515 Mon Sep 17 00:00:00 2001 From: stephantul Date: Thu, 17 Sep 2026 14:05:04 +0200 Subject: [PATCH 2/2] extra test --- tests/test_trainable.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/test_trainable.py b/tests/test_trainable.py index 83ba1ff..cb12f7e 100644 --- a/tests/test_trainable.py +++ b/tests/test_trainable.py @@ -374,6 +374,24 @@ def test_pair_similarity_labels_mismatched_length( ) +def test_pair_similarity_fit_with_explicit_val(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """A model can be fit with explicit validation pairs instead of an automatic split.""" + model = StaticModelForPairSimilarity(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) + text_a = ["word1", "word2", "word3", "word1 word2"] + text_b = ["word2", "word3", "word1", "word3 word1"] + labels = [1, 1, 0, 0] + model.fit( + text_a, + text_b, + labels=labels, + text_a_val=["word1"], + text_b_val=["word2"], + labels_val=[1], + early_stopping_patience=1, + max_epochs=1, + ) + + def test_pair_similarity_fit_with_labels(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: """A model can be fit with a mix of positive and negative pair labels.""" model = StaticModelForPairSimilarity(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer)