From 141a1038bbb3632f18ee50036a5346182c80e7f1 Mon Sep 17 00:00:00 2001 From: stephantul Date: Tue, 6 Oct 2026 09:33:44 +0200 Subject: [PATCH 1/3] speed up training --- model2vec/train/base.py | 138 +++++++++++++++++++--------------- model2vec/train/classifier.py | 12 +-- model2vec/train/dataset.py | 83 +++++++++++++------- model2vec/train/pairs.py | 33 ++++---- model2vec/train/similarity.py | 12 +-- model2vec/train/trainer.py | 2 +- tests/test_trainable.py | 129 ++++++++++++++++++++----------- 7 files changed, 252 insertions(+), 157 deletions(-) diff --git a/model2vec/train/base.py b/model2vec/train/base.py index 2a487e9..eaa181e 100644 --- a/model2vec/train/base.py +++ b/model2vec/train/base.py @@ -11,12 +11,18 @@ from datasets import Dataset as HFDataset from tokenizers import Encoding, Tokenizer from torch import nn -from torch.nn.utils.rnn import pad_sequence from tqdm import trange from model2vec.inference import StaticModelPipeline from model2vec.model import DEFAULT_MAX_LENGTH, PathLike, StaticModel, _disable_padding, _get_unk_token_id -from model2vec.train.dataset import ColumnRows, PairDataset, TextDataset, has_only_strings, has_transform +from model2vec.train.dataset import ( + ColumnRows, + PairDataset, + TextDataset, + TokenBatch, + has_only_strings, + has_transform, +) from model2vec.train.trainer import MetricsFn, default_metrics, resolve_device, run_training_loop from model2vec.train.utils import ( MAX_VALIDATION_SIZE, @@ -46,7 +52,7 @@ def __init__( weights: torch.Tensor | None = None, freeze: bool = False, normalize: bool = True, - freeze_weights: bool = False, + freeze_weights: bool | None = None, max_length: int | None = DEFAULT_MAX_LENGTH, ) -> None: """Initialize a trainable StaticModel from a StaticModel. @@ -59,10 +65,13 @@ def __init__( :param out_dim: The output dimension of the head. :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 weights: The token weights of the model. If None and `freeze_weights` is not False, the model has + no token weights and takes the unweighted mean of the token embeddings. If None and `freeze_weights` + is False, the token weights are initialized to 1. :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 freeze_weights: Whether to freeze the token weights. If None, the model's own weights are trained, + and a model without weights gets none. If False, a model without weights learns weights that start at 1. :param max_length: The default maximum sequence length (in tokens) used to tokenize inputs. Matches `StaticModel.max_length`, defaulting to 512. """ @@ -91,9 +100,9 @@ def __init__( self.token_mapping = torch.arange(len(vectors), dtype=torch.int64) self.token_mapping = nn.Parameter(self.token_mapping, requires_grad=False) self.freeze = freeze - self.embeddings = nn.Embedding.from_pretrained(vectors.clone(), freeze=self.freeze, padding_idx=pad_id) - self.head = self.construct_head() self._weights = weights + self.embeddings = self.construct_embeddings() + self.head = self.construct_head() self.w = self.construct_weights() # Truncation happens here through `max_length`; a StaticModel's tokenizer carries its own setting. self.tokenizer = copy.deepcopy(tokenizer) @@ -120,13 +129,24 @@ def _to_targets(self, labels: Any) -> torch.Tensor: """Turn a batch of labels, such as vectors, into a float tensor of targets.""" return torch.as_tensor(labels, dtype=torch.float32) - def construct_weights(self) -> nn.Parameter: - """Construct the weights for the model.""" - if self._weights is not None: - w = self._weights + @property + def _has_weights(self) -> bool: + """Whether the model has token weights: either given ones, or ones that are learned.""" + return self._weights is not None or self.freeze_weights is False + + def construct_embeddings(self) -> nn.EmbeddingBag: + """Construct the embeddings, which sum the weighted tokens if the model has weights, and average them otherwise.""" + mode = "sum" if self._has_weights else "mean" + return nn.EmbeddingBag.from_pretrained(self.vectors.clone(), freeze=self.freeze, mode=mode) + + def construct_weights(self) -> nn.Parameter | None: + """Construct the token weights for the model, or None if the model has no weights.""" + if not self._has_weights: + return None + if self._weights is None: + w = torch.ones(len(self.token_mapping)) else: - w = torch.ones(len(self.token_mapping)).float() - w[self.pad_id] = 0 + w = self._weights.clone().float() return nn.Parameter(w, requires_grad=not self.freeze_weights) def construct_head(self) -> nn.Sequential: @@ -162,9 +182,7 @@ def construct_head(self) -> nn.Sequential: def _initialize(self) -> None: """Initialize the classifier for training.""" self.head = self.construct_head() - self.embeddings = nn.Embedding.from_pretrained( - self.vectors.clone(), freeze=self.freeze, padding_idx=self.pad_id - ) + self.embeddings = self.construct_embeddings() self.w = self.construct_weights() self.train() @@ -182,7 +200,7 @@ def from_pretrained( out_dim: int = 2, freeze: bool = False, normalize: bool = True, - freeze_weights: bool = False, + freeze_weights: bool | None = None, **kwargs: Any, ) -> T: """Load the model from a pretrained model2vec model. @@ -198,7 +216,8 @@ def from_pretrained( :param out_dim: The output dimension of the head. :param freeze: Whether to freeze the embeddings. :param normalize: Whether to normalize the embeddings. - :param freeze_weights: Whether to freeze the learned token weights. + :param freeze_weights: Whether to freeze the token weights. If None, the model's own weights are trained, + and a model without weights gets none. If False, a model without weights learns weights that start at 1. :param **kwargs: Additional keyword arguments passed to the constructor. :return: The initialized model. """ @@ -228,7 +247,7 @@ def from_static_model( out_dim: int = 2, freeze: bool = False, normalize: bool = True, - freeze_weights: bool = False, + freeze_weights: bool | None = None, **kwargs: Any, ) -> T: """Load the model from a static model. @@ -242,7 +261,8 @@ def from_static_model( :param out_dim: The output dimension of the head. :param freeze: Whether to freeze the embeddings. :param normalize: Whether to normalize the embeddings. - :param freeze_weights: Whether to freeze the learned token weights. + :param freeze_weights: Whether to freeze the token weights. If None, the model's own weights are trained, + and a model without weights gets none. If False, a model without weights learns weights that start at 1. :param **kwargs: Additional keyword arguments passed to the constructor. :return: The initialized model. """ @@ -257,43 +277,39 @@ def from_static_model( **kwargs, ) - def _apply_token_dropout(self, keep_mask: torch.Tensor) -> torch.Tensor: - """Randomly zero out a fraction of the kept tokens, leaving at least one token per sample. + def _apply_token_dropout(self, batch: TokenBatch) -> TokenBatch: + """Randomly drop a fraction of the tokens, leaving at least one token per non-empty sequence. - :param keep_mask: A 2D float tensor (batch, seq_len), 1 for real tokens and 0 for padding. - :return: `keep_mask` with a random subset of real tokens additionally zeroed out. + :param batch: The batch of token ids. + :return: The batch with a random subset of its tokens removed. """ if not self.training or self.token_dropout <= 0: - return keep_mask - survives = torch.rand_like(keep_mask) >= self.token_dropout - dropped_mask = keep_mask * survives - needs_rescue = (dropped_mask.sum(dim=1) == 0) & (keep_mask.sum(dim=1) > 0) - if needs_rescue.any(): - rescue_idx = keep_mask.argmax(dim=1) - dropped_mask[needs_rescue, rescue_idx[needs_rescue]] = 1.0 - return dropped_mask - - def _encode(self, input_ids: torch.Tensor) -> torch.Tensor: + return batch + lengths = batch.lengths + sequence_of_token = torch.repeat_interleave(lengths) + keep = torch.rand(len(batch.ids), device=batch.ids.device) >= self.token_dropout + kept_lengths = torch.bincount(sequence_of_token[keep], minlength=len(batch)) + needs_rescue = (kept_lengths == 0) & (lengths > 0) + keep[batch.offsets[needs_rescue]] = True + kept_lengths = kept_lengths + needs_rescue + return TokenBatch.from_lengths(batch.ids[keep], kept_lengths) + + def _encode(self, batch: TokenBatch) -> torch.Tensor: """A forward pass and mean pooling. This function is analogous to `StaticModel.encode`, but reimplemented to allow gradients to pass through. - :param input_ids: A 2D tensor of input ids. All input ids are have to be within bounds. - :return: The mean over the input ids, weighted by token weights. + :param batch: The batch of token ids. All token ids have to be within bounds. + :return: The mean over the token ids, weighted by token weights if the model has them. """ - zeros = (input_ids != self.pad_id).float() - zeros = self._apply_token_dropout(zeros) - length = zeros.sum(1).clamp(min=1) - input_ids_embeddings = self.token_mapping[input_ids] - embedded = self.embeddings(input_ids_embeddings) - - w = self.w[input_ids] - w = w * zeros - # Weigh each token - embedded = torch.bmm(w[:, None, :], embedded).squeeze(1) - # Mean pooling by dividing by the length - embedded = embedded / length[:, None] + batch = self._apply_token_dropout(batch) + embedding_ids = self.token_mapping[batch.ids] + if self.w is None: + embedded = self.embeddings(embedding_ids, batch.offsets) + else: + embedded = self.embeddings(embedding_ids, batch.offsets, per_sample_weights=self.w[batch.ids]) + embedded = embedded / batch.lengths.clamp(min=1)[:, None] if self.normalize: return nn.functional.normalize(embedded) @@ -313,20 +329,19 @@ def encode(self, X: list[str], batch_size: int = 1024, show_progress_bar: bool = return np.concatenate(pred, axis=0) - def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + def forward(self, batch: TokenBatch) -> torch.Tensor: """Forward pass through the mean, and a classifier layer after.""" - return self.head(self._encode(input_ids)) + return self.head(self._encode(batch)) - def tokenize(self, texts: list[str]) -> torch.Tensor: - """Tokenize a bunch of strings into a single padded 2D tensor. + def tokenize(self, texts: list[str]) -> TokenBatch: + """Tokenize a bunch of strings into a single batch of token ids. Note that this is not used during training. :param texts: The texts to tokenize. - :return: A 2D padded tensor + :return: The batch of token ids. """ - encoded_ids: list[torch.Tensor] = [torch.LongTensor(token_ids) for token_ids in self._tokenize_ids(texts)] - return pad_sequence(encoded_ids, batch_first=True, padding_value=self.pad_id) + return TokenBatch.from_token_ids(self._tokenize_ids(texts)) @property def device(self) -> torch.device: @@ -336,13 +351,12 @@ def device(self) -> torch.device: def to_static_model(self) -> StaticModel: """Convert the model to a static model.""" with torch.no_grad(): - emb = self.embeddings.weight - emb = emb.cpu().numpy() - w = self.w.cpu().numpy() + emb = self.embeddings.weight.cpu().numpy() + w = None if self.w is None else self.w.cpu().numpy() - # If the weights and emb are the same length, the model was not quantized before training. - if len(w) == len(emb): - emb = emb * w[:, None] + # If the token mapping and emb are the same length, the model was not quantized before training. + if len(self.token_mapping) == len(emb): + emb = emb if w is None else emb * w[:, None] return StaticModel( vectors=emb, weights=None, @@ -509,7 +523,7 @@ def _text_dataset(self, rows: ColumnRows, indices: np.ndarray | None = None) -> :param indices: The indices of the rows that belong to the dataset. If None, all rows belong to it. :return: The dataset. """ - return TextDataset(rows, self._tokenize_ids, self._to_targets, indices, pad_id=self.pad_id) + return TextDataset(rows, self._tokenize_ids, self._to_targets, indices) def _create_datasets( self, diff --git a/model2vec/train/classifier.py b/model2vec/train/classifier.py index 4a927d2..6bd2a3c 100644 --- a/model2vec/train/classifier.py +++ b/model2vec/train/classifier.py @@ -89,7 +89,7 @@ def __init__( weights: torch.Tensor | None = None, freeze: bool = False, normalize: bool = True, - freeze_weights: bool = False, + freeze_weights: bool | None = None, max_length: int | None = DEFAULT_MAX_LENGTH, ) -> None: """Initialize a standard classifier model.""" @@ -126,7 +126,7 @@ def from_pretrained( out_dim: int = 2, freeze: bool = False, normalize: bool = True, - freeze_weights: bool = False, + freeze_weights: bool | None = None, **kwargs: Any, ) -> T: """Load a classifier from a pretrained model2vec model. @@ -142,7 +142,8 @@ def from_pretrained( :param out_dim: The number of classes. This is reset when calling `fit`. :param freeze: Whether to freeze the embeddings. :param normalize: Whether to normalize the embeddings. - :param freeze_weights: Whether to freeze the learned token weights. + :param freeze_weights: Whether to freeze the token weights. If None, the model's own weights are trained, + and a model without weights gets none. If False, a model without weights learns weights that start at 1. :param **kwargs: Additional keyword arguments passed to the constructor. :return: The initialized classifier. """ @@ -172,7 +173,7 @@ def from_static_model( out_dim: int = 2, freeze: bool = False, normalize: bool = True, - freeze_weights: bool = False, + freeze_weights: bool | None = None, **kwargs: Any, ) -> T: """Load a classifier from a static model. @@ -186,7 +187,8 @@ def from_static_model( :param out_dim: The number of classes. This is reset when calling `fit`. :param freeze: Whether to freeze the embeddings. :param normalize: Whether to normalize the embeddings. - :param freeze_weights: Whether to freeze the learned token weights. + :param freeze_weights: Whether to freeze the token weights. If None, the model's own weights are trained, + and a model without weights gets none. If False, a model without weights learns weights that start at 1. :param **kwargs: Additional keyword arguments passed to the constructor. :return: The initialized classifier. """ diff --git a/model2vec/train/dataset.py b/model2vec/train/dataset.py index c5640ec..f493f55 100644 --- a/model2vec/train/dataset.py +++ b/model2vec/train/dataset.py @@ -4,6 +4,8 @@ from abc import ABC, abstractmethod from collections import Counter, defaultdict from collections.abc import Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain from typing import Any import numpy as np @@ -12,7 +14,6 @@ import torch from datasets import Column from datasets import Dataset as HFDataset -from torch.nn.utils.rnn import pad_sequence from torch.utils.data import BatchSampler, DataLoader, Dataset, RandomSampler, SequentialSampler logger = logging.getLogger(__name__) @@ -215,17 +216,56 @@ def __getitem__(self, indices: list[int]) -> dict[str, Any]: } +@dataclass(frozen=True) +class TokenBatch: + """A batch of token id sequences, stored as one flat tensor of ids with the offset and length of each sequence.""" + + ids: torch.Tensor + offsets: torch.Tensor + lengths: torch.Tensor + + @classmethod + def from_lengths(cls, ids: torch.Tensor, lengths: torch.Tensor) -> TokenBatch: + """Create a batch from a flat tensor of ids and the length of each sequence.""" + return cls(ids, lengths.cumsum(0) - lengths, lengths) + + @classmethod + def from_token_ids(cls, token_ids: Sequence[Sequence[int]]) -> TokenBatch: + """Flatten lists of token ids into a batch.""" + lengths = np.fromiter(map(len, token_ids), dtype=np.int64, count=len(token_ids)) + ids = np.fromiter(chain.from_iterable(token_ids), dtype=np.int64, count=int(lengths.sum())) + return cls.from_lengths(torch.from_numpy(ids), torch.from_numpy(lengths)) + + def __len__(self) -> int: + """Return the number of sequences.""" + return len(self.lengths) + + def to(self, device: torch.device | str) -> TokenBatch: + """Move the batch to a device.""" + return TokenBatch(self.ids.to(device), self.offsets.to(device), self.lengths.to(device)) + + +@dataclass(frozen=True) +class PairBatch: + """A batch of pairs: the first texts followed by the second texts, and an id per text that is shared by identical texts.""" + + tokens: TokenBatch + text_ids: torch.Tensor + + def to(self, device: torch.device | str) -> PairBatch: + """Move the batch to a device.""" + return PairBatch(self.tokens.to(device), self.text_ids.to(device)) + + class _Batches(Dataset, ABC): - def __init__(self, rows: ColumnRows, indices: np.ndarray | None, pad_id: int) -> None: + def __init__(self, rows: ColumnRows, indices: np.ndarray | None) -> None: """A dataset that fetches rows and turns them into items per batch. :param rows: The rows to draw items from. :param indices: The indices of the rows that belong to this dataset. If None, all rows belong to it. - :param pad_id: The id used to pad batches. Must match the `pad_id` of the model being trained. """ self.rows = rows self.indices = np.arange(len(rows)) if indices is None else indices - self.pad_id = pad_id def __len__(self) -> int: """Return the length of the dataset.""" @@ -244,7 +284,7 @@ def _to_items(self, rows: Mapping[str, Any]) -> list[Any]: """Turn a batch of rows into items.""" @abstractmethod - def collate_fn(self, batch: list[Any]) -> tuple[torch.Tensor, torch.Tensor]: + def collate_fn(self, batch: list[Any]) -> tuple[TokenBatch | PairBatch, torch.Tensor]: """Collate a batch of items into model inputs and targets.""" def _drop_last(self, batch_size: int) -> bool: @@ -268,7 +308,6 @@ def __init__( tokenize: Callable[[list[str]], list[list[int]]], to_targets: Callable[[Any], torch.Tensor], indices: np.ndarray | None = None, - pad_id: int = 0, ) -> None: """A dataset of labeled texts, which are tokenized per batch. @@ -276,9 +315,8 @@ def __init__( :param tokenize: Turns a batch of texts into lists of token ids. :param to_targets: Turns a batch of labels into a tensor of targets. :param indices: The indices of the rows that belong to this dataset. If None, all rows belong to it. - :param pad_id: The id used to pad batches. Must match the `pad_id` of the model being trained. """ - super().__init__(rows, indices, pad_id) + super().__init__(rows, indices) self.tokenize = tokenize self.to_targets = to_targets @@ -286,14 +324,10 @@ def _to_items(self, rows: Mapping[str, Any]) -> list[tuple[list[int], torch.Tens """Tokenize the texts and turn the labels into targets.""" return list(zip(self.tokenize(rows[TEXT_COLUMN]), self.to_targets(rows[LABEL_COLUMN]))) - def collate_fn(self, batch: list[tuple[list[int], torch.Tensor]]) -> tuple[torch.Tensor, torch.Tensor]: + def collate_fn(self, batch: list[tuple[list[int], torch.Tensor]]) -> tuple[TokenBatch, torch.Tensor]: """Collate function.""" texts, targets = zip(*batch) - - tensors: list[torch.Tensor] = [torch.LongTensor(x) for x in texts] - padded = pad_sequence(tensors, batch_first=True, padding_value=self.pad_id) - - return padded, torch.stack(targets) + return TokenBatch.from_token_ids(texts), torch.stack(targets) class PairDataset(_Batches): @@ -302,36 +336,33 @@ def __init__( rows: ColumnRows, tokenize: Callable[[list[str]], list[list[int]]], indices: np.ndarray | None = None, - pad_id: int = 0, ) -> None: """A dataset of aligned text pairs, which are tokenized per batch. :param rows: The pairs, in a `text_a` and a `text_b` column. :param tokenize: Turns a batch of texts into lists of token ids. :param indices: The indices of the rows that belong to this dataset. If None, all rows belong to it. - :param pad_id: The id used to pad batches. Must match the `pad_id` of the model being trained. """ - super().__init__(rows, indices, pad_id) + super().__init__(rows, indices) self.tokenize = tokenize def _to_items(self, rows: Mapping[str, Any]) -> list[tuple[list[int], list[int]]]: """Tokenize both halves of each pair.""" return list(zip(self.tokenize(rows[TEXT_A_COLUMN]), self.tokenize(rows[TEXT_B_COLUMN]))) - def collate_fn(self, batch: list[tuple[list[int], list[int]]]) -> tuple[torch.Tensor, torch.Tensor]: + def collate_fn(self, batch: list[tuple[list[int], list[int]]]) -> tuple[PairBatch, 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. The targets are the index of each - pair's second text within the batch. + The first texts and the second texts are put into a single batch, first texts first. The targets + are the index of each pair's second text within the batch. """ texts_a, texts_b = zip(*batch) + texts = (*texts_a, *texts_b) + ids_by_text: dict[tuple[int, ...], int] = {} + text_ids = [ids_by_text.setdefault(tuple(text), len(ids_by_text)) for text in texts] - 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.arange(len(texts_a)) + pair_batch = PairBatch(TokenBatch.from_token_ids(texts), torch.tensor(text_ids, dtype=torch.int64)) + return pair_batch, torch.arange(len(texts_a)) def _drop_last(self, batch_size: int) -> bool: """Drop a final batch with a single pair, unless it is the only pair.""" diff --git a/model2vec/train/pairs.py b/model2vec/train/pairs.py index 0f3e17c..2c5be08 100644 --- a/model2vec/train/pairs.py +++ b/model2vec/train/pairs.py @@ -11,7 +11,7 @@ from model2vec.model import DEFAULT_MAX_LENGTH, PathLike, StaticModel from model2vec.train.base import BaseFinetuneable, _load_static_model, _static_model_arguments -from model2vec.train.dataset import ColumnRows, PairDataset +from model2vec.train.dataset import ColumnRows, PairBatch, PairDataset from model2vec.train.utils import DEFAULT_RANDOM_SEED, MAX_VALIDATION_SIZE, seed_everything, split_indices from model2vec.types import _UNSET, _UnsetType @@ -71,7 +71,7 @@ def __init__( weights: torch.Tensor | None = None, freeze: bool = False, normalize: bool = True, - freeze_weights: bool = False, + freeze_weights: bool | None = None, max_length: int | None = DEFAULT_MAX_LENGTH, ) -> None: """Initialize a model that is trained to embed pairs of texts close together. @@ -84,10 +84,13 @@ def __init__( :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 weights: The token weights of the model. If None and `freeze_weights` is not False, the model has + no token weights and takes the unweighted mean of the token embeddings. If None and `freeze_weights` + is False, the token weights are initialized to 1. :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 freeze_weights: Whether to freeze the token weights. If None, the model's own weights are trained, + and a model without weights gets none. If False, a model without weights learns weights that start at 1. :param max_length: The default maximum sequence length (in tokens) used to tokenize inputs. """ super().__init__( @@ -119,7 +122,7 @@ def from_pretrained( out_dim: int | None = None, freeze: bool = False, normalize: bool = True, - freeze_weights: bool = False, + freeze_weights: bool | None = None, **kwargs: Any, ) -> T: """Load the model from a pretrained model2vec model. @@ -136,7 +139,8 @@ def from_pretrained( :param out_dim: The output embedding dimension. If None, defaults to the input embedding dimension. :param freeze: Whether to freeze the embeddings. :param normalize: Whether to normalize the embeddings. - :param freeze_weights: Whether to freeze the learned token weights. + :param freeze_weights: Whether to freeze the token weights. If None, the model's own weights are trained, + and a model without weights gets none. If False, a model without weights learns weights that start at 1. :param **kwargs: Additional keyword arguments passed to the constructor. :return: The initialized model. """ @@ -166,7 +170,7 @@ def from_static_model( out_dim: int | None = None, freeze: bool = False, normalize: bool = True, - freeze_weights: bool = False, + freeze_weights: bool | None = None, **kwargs: Any, ) -> T: """Load the model from a static model. @@ -181,7 +185,8 @@ def from_static_model( :param out_dim: The output embedding dimension. If None, defaults to the input embedding dimension. :param freeze: Whether to freeze the embeddings. :param normalize: Whether to normalize the embeddings. - :param freeze_weights: Whether to freeze the learned token weights. + :param freeze_weights: Whether to freeze the token weights. If None, the model's own weights are trained, + and a model without weights gets none. If False, a model without weights learns weights that start at 1. :param **kwargs: Additional keyword arguments passed to the constructor. :return: The initialized model. """ @@ -197,18 +202,16 @@ def from_static_model( ) def forward( # type: ignore[override] - self, input_ids: torch.Tensor + self, batch: PairBatch ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """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. + :param batch: The pair batch, holding the first texts followed by the second texts. :return: The head outputs for the first and second set of texts, followed by an id for each first text and each second text. Identical texts have the same id. """ - out_a = self.head(self._encode(input_ids[0])) - out_b = self.head(self._encode(input_ids[1])) - ids_a = torch.unique(input_ids[0], dim=0, return_inverse=True)[1] - ids_b = torch.unique(input_ids[1], dim=0, return_inverse=True)[1] + out_a, out_b = self.head(self._encode(batch.tokens)).chunk(2) + ids_a, ids_b = batch.text_ids.chunk(2) return out_a, out_b, ids_a, ids_b @staticmethod @@ -233,7 +236,7 @@ def _pair_dataset(self, rows: ColumnRows, indices: np.ndarray | None = None) -> :param indices: The indices of the rows that belong to the dataset. If None, all rows belong to it. :return: The dataset. """ - return PairDataset(rows, self._tokenize_ids, indices, pad_id=self.pad_id) + return PairDataset(rows, self._tokenize_ids, indices) def _create_pair_datasets( self, diff --git a/model2vec/train/similarity.py b/model2vec/train/similarity.py index 9ae9261..aca7a81 100644 --- a/model2vec/train/similarity.py +++ b/model2vec/train/similarity.py @@ -79,7 +79,7 @@ def __init__( weights: torch.Tensor | None = None, freeze: bool = False, normalize: bool = True, - freeze_weights: bool = False, + freeze_weights: bool | None = None, max_length: int | None = DEFAULT_MAX_LENGTH, ) -> None: """Initialize a standard similarity model.""" @@ -112,7 +112,7 @@ def from_pretrained( out_dim: int = 2, freeze: bool = False, normalize: bool = True, - freeze_weights: bool = False, + freeze_weights: bool | None = None, **kwargs: Any, ) -> T: """Load the model from a pretrained model2vec model. @@ -128,7 +128,8 @@ def from_pretrained( :param out_dim: The output dimension of the head. This is reset when calling `fit`. :param freeze: Whether to freeze the embeddings. :param normalize: Whether to normalize the embeddings. - :param freeze_weights: Whether to freeze the learned token weights. + :param freeze_weights: Whether to freeze the token weights. If None, the model's own weights are trained, + and a model without weights gets none. If False, a model without weights learns weights that start at 1. :param **kwargs: Additional keyword arguments passed to the constructor. :return: The initialized model. """ @@ -158,7 +159,7 @@ def from_static_model( out_dim: int = 2, freeze: bool = False, normalize: bool = True, - freeze_weights: bool = False, + freeze_weights: bool | None = None, **kwargs: Any, ) -> T: """Load the model from a static model. @@ -172,7 +173,8 @@ def from_static_model( :param out_dim: The output dimension of the head. This is reset when calling `fit`. :param freeze: Whether to freeze the embeddings. :param normalize: Whether to normalize the embeddings. - :param freeze_weights: Whether to freeze the learned token weights. + :param freeze_weights: Whether to freeze the token weights. If None, the model's own weights are trained, + and a model without weights gets none. If False, a model without weights learns weights that start at 1. :param **kwargs: Additional keyword arguments passed to the constructor. :return: The initialized model. """ diff --git a/model2vec/train/trainer.py b/model2vec/train/trainer.py index 6fe0493..dea7192 100644 --- a/model2vec/train/trainer.py +++ b/model2vec/train/trainer.py @@ -79,7 +79,7 @@ def _run_validation( total_samples = 0 for x, y in val_loader: x, y = x.to(device), y.to(device) - batch_size = x.shape[0] + batch_size = len(y) head_out = model(x) loss = loss_function(head_out, y) for key, value in compute_metrics(head_out, y, loss).items(): diff --git a/tests/test_trainable.py b/tests/test_trainable.py index 38874c6..b9692db 100644 --- a/tests/test_trainable.py +++ b/tests/test_trainable.py @@ -25,6 +25,7 @@ ColumnRows, PairDataset, TextDataset, + TokenBatch, _column_strata, iter_column, ) @@ -45,7 +46,7 @@ def test_init_predict(n_layers: int, mock_vectors: np.ndarray, mock_tokenizer: T vectors_torched = torch.from_numpy(mock_vectors) s = StaticModelForClassification(vectors=vectors_torched, tokenizer=mock_tokenizer, n_layers=n_layers) assert s.vectors.shape == mock_vectors.shape - assert s.w.shape[0] == mock_vectors.shape[0] + assert s.w is None assert list(s.classes) == s.classes_ assert list(s.classes) == ["0", "1"] @@ -63,7 +64,7 @@ def test_init_base_class(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> vectors=vectors_torched, tokenizer=mock_tokenizer, hidden_dim=256, out_dim=3, n_layers=0, pad_id=0 ) assert s.vectors.shape == mock_vectors.shape - assert s.w.shape[0] == mock_vectors.shape[0] + assert s.w is None head = s.construct_head() assert head[0].in_features == mock_vectors.shape[1] @@ -95,13 +96,13 @@ def test_init_base_from_model(mock_vectors: np.ndarray, mock_tokenizer: Tokenize model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer) s = BaseFinetuneable.from_static_model(model=model) assert s.vectors.shape == mock_vectors.shape - assert s.w.shape[0] == mock_vectors.shape[0] + assert s.w is None with TemporaryDirectory() as temp_dir: model.save_pretrained(temp_dir) s = BaseFinetuneable.from_pretrained(path=temp_dir) assert s.vectors.shape == mock_vectors.shape - assert s.w.shape[0] == mock_vectors.shape[0] + assert s.w is None def test_init_classifier_from_model(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: @@ -109,13 +110,13 @@ def test_init_classifier_from_model(mock_vectors: np.ndarray, mock_tokenizer: To model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer) s = StaticModelForClassification.from_static_model(model=model) assert s.vectors.shape == mock_vectors.shape - assert s.w.shape[0] == mock_vectors.shape[0] + assert s.w is None with TemporaryDirectory() as temp_dir: model.save_pretrained(temp_dir) s = StaticModelForClassification.from_pretrained(path=temp_dir) assert s.vectors.shape == mock_vectors.shape - assert s.w.shape[0] == mock_vectors.shape[0] + assert s.w is None FINETUNEABLE_CLASSES = [ @@ -160,7 +161,7 @@ def test_from_pretrained_explicit_arguments( cls: type[BaseFinetuneable], mock_vectors: np.ndarray, mock_tokenizer: Tokenizer ) -> None: """Test that from_pretrained passes each argument on to the model.""" - model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer) + model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer, weights=np.ones(len(mock_vectors))) with TemporaryDirectory() as temp_dir: model.save_pretrained(temp_dir) s = cls.from_pretrained( @@ -256,11 +257,28 @@ def test_init_classifier_from_model_w(mock_vectors: np.ndarray, mock_tokenizer: s = StaticModelForClassification.from_static_model(model=model) assert s._weights is not None assert torch.all(s._weights == torch.ones(len(mock_vectors))) + assert s.freeze_weights is None + assert s.w is not None and s.w.requires_grad w = s.construct_weights() assert w.shape[0] == mock_vectors.shape[0] assert torch.all(w == torch.ones(len(mock_vectors))) +def test_unfrozen_weights_start_at_one(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Without given weights, unfrozen token weights are initialized to 1, which encodes like the plain mean.""" + vectors = torch.from_numpy(mock_vectors).float() + unweighted = StaticModelForClassification(vectors=vectors, tokenizer=mock_tokenizer) + weighted = StaticModelForClassification(vectors=vectors, tokenizer=mock_tokenizer, freeze_weights=False) + assert unweighted.w is None + assert weighted.w is not None + assert weighted.w.requires_grad + assert torch.equal(weighted.w, torch.ones(len(mock_vectors))) + + batch = TokenBatch.from_token_ids([[1, 2, 3], [4], []]) + with torch.no_grad(): + assert torch.allclose(weighted._encode(batch), unweighted._encode(batch)) + + def test_pad_token(mock_tokenizer: Tokenizer) -> None: """Test initializion from a static model.""" tokenizer_model = TokenizerModel.from_tokenizer(mock_tokenizer) @@ -269,42 +287,65 @@ def test_pad_token(mock_tokenizer: Tokenizer) -> None: vectors = np.random.RandomState().randn(6, 10) model = StaticModel(vectors=vectors, tokenizer=tokenizer) s = StaticModelForClassification.from_static_model(model=model, pad_token="[HELLO]") - assert s.w.shape[0] == vectors.shape[0] + assert s.w is None assert s.pad_id == 5 with pytest.raises(KeyError): StaticModelForClassification.from_static_model(model=model, pad_token="[BRR]") +def _sequences(batch: TokenBatch) -> list[list[int]]: + return [ids.tolist() for ids in batch.ids.split(batch.lengths.tolist())] + + def test_encode(mock_trained_pipeline: StaticModelForClassification) -> None: """Test the encode function.""" - result = mock_trained_pipeline._encode(torch.tensor([[0, 1], [1, 0]]).long()) + result = mock_trained_pipeline._encode(TokenBatch.from_token_ids([[0, 1], [1, 0]])) assert result.shape == (2, 12) assert torch.allclose(result[0], result[1]) +def test_encode_with_weights(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """With token weights, a text is encoded as the weighted sum of its tokens divided by its length.""" + weights = torch.rand(len(mock_vectors)) + vectors = torch.from_numpy(mock_vectors).float() + s = StaticModelForClassification(vectors=vectors, tokenizer=mock_tokenizer, weights=weights, normalize=False) + assert s.w is not None + with torch.no_grad(): + result = s._encode(TokenBatch.from_token_ids([[1, 2], [3], []])) + expected_first = (vectors[1] * weights[1] + vectors[2] * weights[2]) / 2 + assert torch.allclose(result[0], expected_first) + assert torch.allclose(result[1], vectors[3] * weights[3]) + assert torch.equal(result[2], torch.zeros(vectors.shape[1])) + + def test_tokenize(mock_trained_pipeline: StaticModelForClassification) -> None: """Test the encode function.""" result = mock_trained_pipeline.tokenize(["dog dog", "cat"]) - assert result.shape == torch.Size([2, 2]) - assert result[1, 1] == 0 + assert len(result) == 2 + assert result.lengths.tolist() == [2, 1] def test_device(mock_trained_pipeline: StaticModelForClassification) -> None: """Get the device.""" assert mock_trained_pipeline.device == torch.device(type="cpu") # type: ignore # False positive - assert mock_trained_pipeline.device == mock_trained_pipeline.w.device + assert mock_trained_pipeline.device == mock_trained_pipeline.token_mapping.device -def test_conversion(mock_trained_pipeline: StaticModelForClassification) -> None: - """Test the conversion to numpy.""" - staticmodel = mock_trained_pipeline.to_static_model() +@pytest.mark.parametrize("with_weights", [False, True]) +def test_conversion(with_weights: bool, mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Test the conversion to a static model, with and without token weights.""" + weights = torch.rand(len(mock_vectors)) if with_weights else None + s = StaticModelForClassification( + vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer, weights=weights + ) + staticmodel = s.to_static_model() with torch.no_grad(): - result_1 = mock_trained_pipeline._encode(torch.tensor([[1, 2], [2, 1]]).long()).numpy() + result_1 = s._encode(TokenBatch.from_token_ids([[1, 2], [2, 1]])).numpy() result_2 = staticmodel.embedding[[[1, 2], [2, 1]]].mean(0) result_2 /= np.linalg.norm(result_2, axis=1, keepdims=True) - assert np.allclose(result_1, result_2) + assert np.allclose(result_1, result_2, atol=1e-6) def test_token_dropout_default_is_zero(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: @@ -318,8 +359,8 @@ def test_apply_token_dropout_noop_in_eval(mock_vectors: np.ndarray, mock_tokeniz s = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) s.token_dropout = 0.9 s.eval() - mask = torch.ones(4, 5) - assert torch.equal(s._apply_token_dropout(mask), mask) + batch = TokenBatch.from_token_ids([[1, 2, 3, 4, 5]] * 4) + assert s._apply_token_dropout(batch) is batch def test_apply_token_dropout_noop_when_rate_is_zero(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: @@ -327,8 +368,8 @@ def test_apply_token_dropout_noop_when_rate_is_zero(mock_vectors: np.ndarray, mo s = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) s.train() s.token_dropout = 0.0 - mask = torch.ones(4, 5) - assert torch.equal(s._apply_token_dropout(mask), mask) + batch = TokenBatch.from_token_ids([[1, 2, 3, 4, 5]] * 4) + assert s._apply_token_dropout(batch) is batch def test_apply_token_dropout_never_empties_a_nonempty_row(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: @@ -336,12 +377,14 @@ def test_apply_token_dropout_never_empties_a_nonempty_row(mock_vectors: np.ndarr s = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) s.train() s.token_dropout = 0.99 - mask = torch.tensor([[1.0, 1.0, 1.0, 1.0], [1.0, 0.0, 0.0, 0.0]]) + batch = TokenBatch.from_token_ids([[1, 2, 3, 4], [], [5]]) torch.manual_seed(0) for _ in range(20): - out = s._apply_token_dropout(mask) - assert (out.sum(dim=1) >= 1).all() + out = s._apply_token_dropout(batch) + sequences = _sequences(out) + assert len(sequences[0]) == 1 and sequences[0][0] in [1, 2, 3, 4] + assert sequences[1:] == [[], [5]] def test_apply_token_dropout_leaves_fully_padded_rows_untouched( @@ -351,12 +394,12 @@ def test_apply_token_dropout_leaves_fully_padded_rows_untouched( s = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) s.train() s.token_dropout = 0.99 - mask = torch.zeros(1, 4) + batch = TokenBatch.from_token_ids([[], []]) torch.manual_seed(0) for _ in range(20): - out = s._apply_token_dropout(mask) - assert out.sum() == 0 + out = s._apply_token_dropout(batch) + assert _sequences(out) == [[], []] def test_apply_token_dropout_drops_some_tokens(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: @@ -364,11 +407,12 @@ def test_apply_token_dropout_drops_some_tokens(mock_vectors: np.ndarray, mock_to s = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) s.train() s.token_dropout = 0.5 - mask = torch.ones(1, 1000) + batch = TokenBatch.from_token_ids([list(range(1000))]) torch.manual_seed(0) - out = s._apply_token_dropout(mask) - assert 0 < out.sum().item() < mask.sum().item() + out = s._apply_token_dropout(batch) + assert 0 < len(out.ids) < len(batch.ids) + assert out.offsets.tolist() == [0] def test_fit_invalid_token_dropout_raises() -> None: @@ -434,15 +478,15 @@ def test_column_rows() -> None: ColumnRows(text=["a"], label=torch.arange(2)) -def test_training_batch_padding_is_masked(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: - """Training batches should pad with the model's pad id, so padding stays masked and out of the mean.""" - s = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer, pad_id=1) +def test_training_batch_matches_tokenize(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: + """Training batches hold the same tokens as `tokenize`, and a text encodes the same regardless of its batch.""" + s = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer) texts = ["word2", "word2 word3"] dataset = s._text_dataset(ColumnRows(text=texts, label=["0", "1"])) batch, _ = next(iter(dataset.to_dataloader(shuffle=False, batch_size=2))) - assert torch.equal(batch, s.tokenize(texts)) + assert _sequences(batch) == _sequences(s.tokenize(texts)) with torch.no_grad(): assert torch.allclose(s._encode(batch)[0], s._encode(s.tokenize(texts[:1]))[0]) @@ -454,7 +498,7 @@ def test_unknown_tokens_are_dropped(mock_vectors: np.ndarray, mock_tokenizer: To texts = ["word1 unknownword", "unknownword word2 otherunknown"] expected = static.tokenize(texts) - assert [row[row != s.pad_id].tolist() for row in s.tokenize(texts)] == expected + assert _sequences(s.tokenize(texts)) == expected assert s._tokenize_ids(texts) == expected @@ -472,7 +516,7 @@ def test_tokenize_without_unk_token(mock_vectors: np.ndarray) -> None: texts = ["word1 word2", "word3 word1 word2"] expected = static.tokenize(texts) - assert [row[row != s.pad_id].tolist() for row in s.tokenize(texts)] == expected + assert _sequences(s.tokenize(texts)) == expected assert s._tokenize_ids(texts) == expected @@ -483,14 +527,14 @@ def test_max_length_is_not_capped_by_the_static_model(mock_vectors: np.ndarray, assert len(static.tokenize(texts)[0]) == 2 s = StaticModelForClassification.from_static_model(model=static, max_length=4) - assert s.tokenize(texts).shape[1] == 4 + assert s.tokenize(texts).lengths.tolist() == [4] assert [len(row) for row in s._tokenize_ids(texts)] == [4] assert len(s.to_static_model().tokenize(texts)[0]) == 4 # The static model keeps its own setting, and the trainer keeps its own once the static model changes. assert len(static.tokenize(texts)[0]) == 2 static.max_length = None - assert s.tokenize(texts).shape[1] == 4 + assert s.tokenize(texts).lengths.tolist() == [4] @pytest.mark.parametrize("cls", FINETUNEABLE_CLASSES) @@ -596,13 +640,12 @@ def test_pairdataset_init_incorrect() -> None: def test_pairdataset_collate() -> None: - """Batches should stack the two padded halves into a single (2, batch, seq_len) tensor.""" - dataset = PairDataset(ColumnRows(text_a=[[1], [1, 2]], text_b=[[1, 2, 3], [1]]), _pretokenized, pad_id=0) + """Batches hold the first texts followed by the second texts, with a shared id for identical texts.""" + dataset = PairDataset(ColumnRows(text_a=[[1], [1, 2]], text_b=[[1, 2, 3], [1]]), _pretokenized) batch, y = next(iter(dataset.to_dataloader(shuffle=False, batch_size=2))) - assert batch.shape == (2, 2, 3) assert torch.equal(y, torch.tensor([0, 1])) - 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]])) + assert _sequences(batch.tokens) == [[1], [1, 2], [1, 2, 3], [1]] + assert batch.text_ids.tolist() == [0, 1, 2, 0] def _distinct(n: int) -> torch.Tensor: From aab4791ab23e7a9ab7c56faedd633e9a71e95881 Mon Sep 17 00:00:00 2001 From: stephantul Date: Tue, 6 Oct 2026 15:13:41 +0200 Subject: [PATCH 2/3] feat: speed up training --- model2vec/train/base.py | 21 ++++----- model2vec/train/classifier.py | 6 +-- model2vec/train/pairs.py | 8 ++-- model2vec/train/similarity.py | 6 +-- model2vec/train/utils.py | 16 ------- tests/test_trainable.py | 86 ++++++++++------------------------- 6 files changed, 42 insertions(+), 101 deletions(-) diff --git a/model2vec/train/base.py b/model2vec/train/base.py index eaa181e..0efc657 100644 --- a/model2vec/train/base.py +++ b/model2vec/train/base.py @@ -26,7 +26,6 @@ from model2vec.train.trainer import MetricsFn, default_metrics, resolve_device, run_training_loop from model2vec.train.utils import ( MAX_VALIDATION_SIZE, - get_probable_pad_token_id, split_indices, to_pipeline, ) @@ -47,7 +46,7 @@ def __init__( hidden_dim: int = 256, n_layers: int = 0, out_dim: int = 2, - pad_id: int = 0, + pad_id: int | None = None, token_mapping: list[int] | None = None, weights: torch.Tensor | None = None, freeze: bool = False, @@ -63,7 +62,7 @@ def __init__( :param n_layers: The number of layers in the head. If this is 0 and `out_dim` equals the embedding dimension, the model has no head and the embeddings are used as is. :param out_dim: The output dimension of the head. - :param pad_id: The padding id. This is set to 0 in almost all model2vec models + :param pad_id: Deprecated and ignored. :param token_mapping: The token mapping. If None, the token mapping is set to the range of the number of vectors. :param weights: The token weights of the model. If None and `freeze_weights` is not False, the model has no token weights and takes the unweighted mean of the token embeddings. If None and `freeze_weights` @@ -76,7 +75,8 @@ def __init__( Matches `StaticModel.max_length`, defaulting to 512. """ super().__init__() - self.pad_id = pad_id + if pad_id is not None: + logger.warning("The 'pad_id' argument is deprecated and ignored.") self.out_dim = out_dim self.embed_dim = vectors.shape[1] self.hidden_dim = hidden_dim @@ -208,7 +208,7 @@ def from_pretrained( :param path: The path to the folder containing the model, or a repository on the Hugging Face Hub. :param token: The token to use to download the model from the hub. :param model_name: Deprecated alias for `path`. - :param pad_token: The token to use for padding. If None, it is inferred from the tokenizer. + :param pad_token: Deprecated and ignored. :param max_length: The default maximum sequence length to use for tokenization. If not passed, the static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of layers in the head. @@ -253,7 +253,7 @@ def from_static_model( """Load the model from a static model. :param model: The static model to load from. - :param pad_token: The token to use for padding. If None, it is inferred from the tokenizer. + :param pad_token: Deprecated and ignored. :param max_length: The default maximum sequence length to use for tokenization. If not passed, the static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of layers in the head. @@ -586,21 +586,18 @@ def _static_model_arguments( """Derive the constructor arguments of a finetuneable model from a static model. :param model: The static model to derive the arguments from. - :param pad_token: The token to use for padding. If None, it is inferred from the tokenizer. + :param pad_token: Deprecated and ignored. :param max_length: The default maximum sequence length to use for tokenization. If unset, the static model's `max_length` is used. If None, no truncation is done. :return: The constructor arguments. """ + if pad_token is not None: + logger.warning("The 'pad_token' argument is deprecated and ignored.") model.embedding = np.nan_to_num(model.embedding) weights = torch.from_numpy(model.weights) if model.weights is not None else None token_mapping = model.token_mapping.tolist() if model.token_mapping is not None else None - if pad_token is not None: - pad_id = model.tokenizer.get_vocab()[pad_token] - else: - pad_id = get_probable_pad_token_id(model.tokenizer) return { "vectors": torch.from_numpy(model.embedding), - "pad_id": pad_id, "tokenizer": model.tokenizer, "token_mapping": token_mapping, "weights": weights, diff --git a/model2vec/train/classifier.py b/model2vec/train/classifier.py index 6bd2a3c..111cdb5 100644 --- a/model2vec/train/classifier.py +++ b/model2vec/train/classifier.py @@ -84,7 +84,7 @@ def __init__( n_layers: int = 1, hidden_dim: int = 512, out_dim: int = 2, - pad_id: int = 0, + pad_id: int | None = None, token_mapping: list[int] | None = None, weights: torch.Tensor | None = None, freeze: bool = False, @@ -134,7 +134,7 @@ def from_pretrained( :param path: The path to the folder containing the model, or a repository on the Hugging Face Hub. :param token: The token to use to download the model from the hub. :param model_name: Deprecated alias for `path`. - :param pad_token: The token to use for padding. If None, it is inferred from the tokenizer. + :param pad_token: Deprecated and ignored. :param max_length: The default maximum sequence length to use for tokenization. If not passed, the static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of hidden layers in the head. @@ -179,7 +179,7 @@ def from_static_model( """Load a classifier from a static model. :param model: The static model to load from. - :param pad_token: The token to use for padding. If None, it is inferred from the tokenizer. + :param pad_token: Deprecated and ignored. :param max_length: The default maximum sequence length to use for tokenization. If not passed, the static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of hidden layers in the head. diff --git a/model2vec/train/pairs.py b/model2vec/train/pairs.py index 2c5be08..62b89a7 100644 --- a/model2vec/train/pairs.py +++ b/model2vec/train/pairs.py @@ -66,7 +66,7 @@ def __init__( n_layers: int = 1, hidden_dim: int = 512, out_dim: int | None = None, - pad_id: int = 0, + pad_id: int | None = None, token_mapping: list[int] | None = None, weights: torch.Tensor | None = None, freeze: bool = False, @@ -82,7 +82,7 @@ def __init__( dimension, the model has no head, and the embeddings are used as is. :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 pad_id: Deprecated and ignored. :param token_mapping: The token mapping. If None, the token mapping is set to the range of the number of vectors. :param weights: The token weights of the model. If None and `freeze_weights` is not False, the model has no token weights and takes the unweighted mean of the token embeddings. If None and `freeze_weights` @@ -130,7 +130,7 @@ def from_pretrained( :param path: The path to the folder containing the model, or a repository on the Hugging Face Hub. :param token: The token to use to download the model from the hub. :param model_name: Deprecated alias for `path`. - :param pad_token: The token to use for padding. If None, it is inferred from the tokenizer. + :param pad_token: Deprecated and ignored. :param max_length: The default maximum sequence length to use for tokenization. If not passed, the static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of layers in the head. If this is 0 and `out_dim` equals the embedding @@ -176,7 +176,7 @@ def from_static_model( """Load the model from a static model. :param model: The static model to load from. - :param pad_token: The token to use for padding. If None, it is inferred from the tokenizer. + :param pad_token: Deprecated and ignored. :param max_length: The default maximum sequence length to use for tokenization. If not passed, the static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of layers in the head. If this is 0 and `out_dim` equals the embedding diff --git a/model2vec/train/similarity.py b/model2vec/train/similarity.py index aca7a81..74e853f 100644 --- a/model2vec/train/similarity.py +++ b/model2vec/train/similarity.py @@ -74,7 +74,7 @@ def __init__( n_layers: int = 1, hidden_dim: int = 512, out_dim: int = 2, - pad_id: int = 0, + pad_id: int | None = None, token_mapping: list[int] | None = None, weights: torch.Tensor | None = None, freeze: bool = False, @@ -120,7 +120,7 @@ def from_pretrained( :param path: The path to the folder containing the model, or a repository on the Hugging Face Hub. :param token: The token to use to download the model from the hub. :param model_name: Deprecated alias for `path`. - :param pad_token: The token to use for padding. If None, it is inferred from the tokenizer. + :param pad_token: Deprecated and ignored. :param max_length: The default maximum sequence length to use for tokenization. If not passed, the static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of hidden layers in the head. @@ -165,7 +165,7 @@ def from_static_model( """Load the model from a static model. :param model: The static model to load from. - :param pad_token: The token to use for padding. If None, it is inferred from the tokenizer. + :param pad_token: Deprecated and ignored. :param max_length: The default maximum sequence length to use for tokenization. If not passed, the static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of hidden layers in the head. diff --git a/model2vec/train/utils.py b/model2vec/train/utils.py index ad38c83..562112e 100644 --- a/model2vec/train/utils.py +++ b/model2vec/train/utils.py @@ -8,7 +8,6 @@ import numpy as np import torch -from tokenizers import Tokenizer from torch import nn from model2vec.inference import StaticModelPipeline @@ -24,21 +23,6 @@ DEFAULT_RANDOM_SEED = 42 MAX_VALIDATION_SIZE = 10_000 -_KNOWN_PAD_TOKENS = ("[PAD]", "") - - -def get_probable_pad_token_id(tokenizer: Tokenizer) -> int: - """Get a probable pad token by using the padding module and falling back to guessing.""" - if tokenizer.padding is not None: - return tokenizer.padding["pad_id"] - vocab = tokenizer.get_vocab() - for token in _KNOWN_PAD_TOKENS: - token_id = vocab.get(token) - if token_id is not None: - return token_id - - logger.warning("No known pad token found, using 0 as default") - return 0 def to_pipeline(model: "BaseFinetuneable | StaticModelForClassification") -> StaticModelPipeline: diff --git a/tests/test_trainable.py b/tests/test_trainable.py index b9692db..8a3f244 100644 --- a/tests/test_trainable.py +++ b/tests/test_trainable.py @@ -9,7 +9,6 @@ import pytest import torch from datasets import Dataset, DatasetDict, Features, Sequence, Value -from skeletoken import TokenizerModel from tokenizers import Tokenizer from tokenizers.models import BPE from tokenizers.pre_tokenizers import Whitespace @@ -34,7 +33,6 @@ from model2vec.train.similarity import StaticModelForSimilarity, _vector_dim from model2vec.train.trainer import _resolve_max_epochs, resolve_device, run_training_loop from model2vec.train.utils import ( - get_probable_pad_token_id, seed_everything, split_indices, ) @@ -60,9 +58,7 @@ def test_init_predict(n_layers: int, mock_vectors: np.ndarray, mock_tokenizer: T def test_init_base_class(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: """Test successful initialization of the base class.""" vectors_torched = torch.from_numpy(mock_vectors) - s = BaseFinetuneable( - vectors=vectors_torched, tokenizer=mock_tokenizer, hidden_dim=256, out_dim=3, n_layers=0, pad_id=0 - ) + s = BaseFinetuneable(vectors=vectors_torched, tokenizer=mock_tokenizer, hidden_dim=256, out_dim=3, n_layers=0) assert s.vectors.shape == mock_vectors.shape assert s.w is None @@ -279,19 +275,29 @@ def test_unfrozen_weights_start_at_one(mock_vectors: np.ndarray, mock_tokenizer: assert torch.allclose(weighted._encode(batch), unweighted._encode(batch)) -def test_pad_token(mock_tokenizer: Tokenizer) -> None: - """Test initializion from a static model.""" - tokenizer_model = TokenizerModel.from_tokenizer(mock_tokenizer) - tokenizer_model.pad_token = "[HELLO]" - tokenizer = tokenizer_model.to_tokenizer() - vectors = np.random.RandomState().randn(6, 10) - model = StaticModel(vectors=vectors, tokenizer=tokenizer) - s = StaticModelForClassification.from_static_model(model=model, pad_token="[HELLO]") - assert s.w is None - assert s.pad_id == 5 +@pytest.mark.parametrize("cls", FINETUNEABLE_CLASSES) +def test_pad_arguments_are_deprecated( + cls: type[BaseFinetuneable], + mock_vectors: np.ndarray, + mock_tokenizer: Tokenizer, + caplog: pytest.LogCaptureFixture, +) -> None: + """Passing `pad_token` or `pad_id` logs a deprecation warning and has no effect.""" + model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer) + with caplog.at_level(logging.WARNING, logger="model2vec.train.base"): + s = cls.from_static_model(model=model, pad_token="[BRR]") + assert "The 'pad_token' argument is deprecated and ignored." in caplog.text + assert not hasattr(s, "pad_id") - with pytest.raises(KeyError): - StaticModelForClassification.from_static_model(model=model, pad_token="[BRR]") + caplog.clear() + with caplog.at_level(logging.WARNING, logger="model2vec.train.base"): + cls(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer, pad_id=5) + assert "The 'pad_id' argument is deprecated and ignored." in caplog.text + + caplog.clear() + with caplog.at_level(logging.WARNING, logger="model2vec.train.base"): + cls.from_static_model(model=model) + assert "deprecated" not in caplog.text def _sequences(batch: TokenBatch) -> list[list[int]]: @@ -973,52 +979,6 @@ def test_evaluate(mock_trained_pipeline: StaticModelForClassification) -> None: mock_trained_pipeline.evaluate(["dog cat", "dog"], [1, 1]) # type: ignore -def test_get_probable_pad_token_id(mock_tokenizer: Tokenizer, caplog: pytest.LogCaptureFixture) -> None: - """Test loading from a static model with a pad token.""" - tokenizer_model = TokenizerModel.from_tokenizer(mock_tokenizer) - t = tokenizer_model.to_tokenizer() - token_id = get_probable_pad_token_id(t) - assert token_id == 0 - - # Adds new token - tokenizer_model.pad_token = "haha" - t = tokenizer_model.to_tokenizer() - token_id = get_probable_pad_token_id(t) - assert token_id == 5 - - tokenizer_model.pad_token = "word1" - t = tokenizer_model.to_tokenizer() - token_id = get_probable_pad_token_id(t) - assert token_id == 1 - - # Remove padding token - tokenizer_model.pad_token = None - t = tokenizer_model.to_tokenizer() - token_id = get_probable_pad_token_id(t) - assert token_id == tokenizer_model.vocabulary["[PAD]"] - - tokenizer_model = tokenizer_model.remove_token_from_vocabulary("[PAD]") - t = tokenizer_model.to_tokenizer() - with caplog.at_level(logging.WARNING, logger="model2vec.train.utils"): - token_id = get_probable_pad_token_id(t) - assert token_id == 0 - assert "No known pad token found, using 0 as default" in caplog.text - - -def test_get_probable_pad_token_id_through_static_model(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None: - """Test that a non-standard pad token survives StaticModel construction.""" - tokenizer_model = TokenizerModel.from_tokenizer(mock_tokenizer) - tokenizer_model.pad_token = "word1" - pad_id = tokenizer_model.pad_token_id - assert pad_id != 0 - - t = tokenizer_model.to_tokenizer() - model = StaticModel(vectors=mock_vectors, tokenizer=t) - - assert model.tokenizer.padding is not None - assert get_probable_pad_token_id(model.tokenizer) == pad_id - - def test_resolve_class_weight(mock_trained_pipeline: StaticModelForClassification) -> None: """Test what the class weights are.""" w_dict = dict(zip(mock_trained_pipeline.classes, [0.5, 3])) From 0ad18b230cbf3194d76e68135acd5a97be586311 Mon Sep 17 00:00:00 2001 From: stephantul Date: Tue, 6 Oct 2026 15:21:39 +0200 Subject: [PATCH 3/3] remove pad token --- model2vec/train/base.py | 18 ++---------------- model2vec/train/classifier.py | 9 +-------- model2vec/train/pairs.py | 10 +--------- model2vec/train/similarity.py | 9 +-------- tests/test_trainable.py | 27 --------------------------- 5 files changed, 5 insertions(+), 68 deletions(-) diff --git a/model2vec/train/base.py b/model2vec/train/base.py index 0efc657..06f75e4 100644 --- a/model2vec/train/base.py +++ b/model2vec/train/base.py @@ -46,7 +46,6 @@ def __init__( hidden_dim: int = 256, n_layers: int = 0, out_dim: int = 2, - pad_id: int | None = None, token_mapping: list[int] | None = None, weights: torch.Tensor | None = None, freeze: bool = False, @@ -62,7 +61,6 @@ def __init__( :param n_layers: The number of layers in the head. If this is 0 and `out_dim` equals the embedding dimension, the model has no head and the embeddings are used as is. :param out_dim: The output dimension of the head. - :param pad_id: Deprecated and ignored. :param token_mapping: The token mapping. If None, the token mapping is set to the range of the number of vectors. :param weights: The token weights of the model. If None and `freeze_weights` is not False, the model has no token weights and takes the unweighted mean of the token embeddings. If None and `freeze_weights` @@ -75,8 +73,6 @@ def __init__( Matches `StaticModel.max_length`, defaulting to 512. """ super().__init__() - if pad_id is not None: - logger.warning("The 'pad_id' argument is deprecated and ignored.") self.out_dim = out_dim self.embed_dim = vectors.shape[1] self.hidden_dim = hidden_dim @@ -193,7 +189,6 @@ def from_pretrained( *, token: str | None = None, model_name: PathLike | None = None, - pad_token: str | None = None, max_length: int | None | _UnsetType = _UNSET, n_layers: int = 0, hidden_dim: int = 256, @@ -208,7 +203,6 @@ def from_pretrained( :param path: The path to the folder containing the model, or a repository on the Hugging Face Hub. :param token: The token to use to download the model from the hub. :param model_name: Deprecated alias for `path`. - :param pad_token: Deprecated and ignored. :param max_length: The default maximum sequence length to use for tokenization. If not passed, the static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of layers in the head. @@ -224,7 +218,6 @@ def from_pretrained( model = _load_static_model(path, token=token, model_name=model_name) return cls.from_static_model( model=model, - pad_token=pad_token, max_length=max_length, n_layers=n_layers, hidden_dim=hidden_dim, @@ -240,7 +233,6 @@ def from_static_model( cls: type[T], *, model: StaticModel, - pad_token: str | None = None, max_length: int | None | _UnsetType = _UNSET, n_layers: int = 0, hidden_dim: int = 256, @@ -253,7 +245,6 @@ def from_static_model( """Load the model from a static model. :param model: The static model to load from. - :param pad_token: Deprecated and ignored. :param max_length: The default maximum sequence length to use for tokenization. If not passed, the static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of layers in the head. @@ -267,7 +258,7 @@ def from_static_model( :return: The initialized model. """ return cls( - **_static_model_arguments(model, pad_token=pad_token, max_length=max_length), + **_static_model_arguments(model, max_length=max_length), n_layers=n_layers, hidden_dim=hidden_dim, out_dim=out_dim, @@ -580,19 +571,14 @@ def _load_static_model(path: PathLike, *, token: str | None, model_name: PathLik return StaticModel.from_pretrained(path, token=token) -def _static_model_arguments( - model: StaticModel, *, pad_token: str | None, max_length: int | None | _UnsetType -) -> dict[str, Any]: +def _static_model_arguments(model: StaticModel, *, max_length: int | None | _UnsetType) -> dict[str, Any]: """Derive the constructor arguments of a finetuneable model from a static model. :param model: The static model to derive the arguments from. - :param pad_token: Deprecated and ignored. :param max_length: The default maximum sequence length to use for tokenization. If unset, the static model's `max_length` is used. If None, no truncation is done. :return: The constructor arguments. """ - if pad_token is not None: - logger.warning("The 'pad_token' argument is deprecated and ignored.") model.embedding = np.nan_to_num(model.embedding) weights = torch.from_numpy(model.weights) if model.weights is not None else None token_mapping = model.token_mapping.tolist() if model.token_mapping is not None else None diff --git a/model2vec/train/classifier.py b/model2vec/train/classifier.py index 111cdb5..fb0d3f6 100644 --- a/model2vec/train/classifier.py +++ b/model2vec/train/classifier.py @@ -84,7 +84,6 @@ def __init__( n_layers: int = 1, hidden_dim: int = 512, out_dim: int = 2, - pad_id: int | None = None, token_mapping: list[int] | None = None, weights: torch.Tensor | None = None, freeze: bool = False, @@ -100,7 +99,6 @@ def __init__( super().__init__( vectors=vectors, out_dim=out_dim, - pad_id=pad_id, tokenizer=tokenizer, token_mapping=token_mapping, weights=weights, @@ -119,7 +117,6 @@ def from_pretrained( *, token: str | None = None, model_name: PathLike | None = None, - pad_token: str | None = None, max_length: int | None | _UnsetType = _UNSET, n_layers: int = 1, hidden_dim: int = 512, @@ -134,7 +131,6 @@ def from_pretrained( :param path: The path to the folder containing the model, or a repository on the Hugging Face Hub. :param token: The token to use to download the model from the hub. :param model_name: Deprecated alias for `path`. - :param pad_token: Deprecated and ignored. :param max_length: The default maximum sequence length to use for tokenization. If not passed, the static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of hidden layers in the head. @@ -150,7 +146,6 @@ def from_pretrained( model = _load_static_model(path, token=token, model_name=model_name) return cls.from_static_model( model=model, - pad_token=pad_token, max_length=max_length, n_layers=n_layers, hidden_dim=hidden_dim, @@ -166,7 +161,6 @@ def from_static_model( cls: type[T], *, model: StaticModel, - pad_token: str | None = None, max_length: int | None | _UnsetType = _UNSET, n_layers: int = 1, hidden_dim: int = 512, @@ -179,7 +173,6 @@ def from_static_model( """Load a classifier from a static model. :param model: The static model to load from. - :param pad_token: Deprecated and ignored. :param max_length: The default maximum sequence length to use for tokenization. If not passed, the static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of hidden layers in the head. @@ -193,7 +186,7 @@ def from_static_model( :return: The initialized classifier. """ return cls( - **_static_model_arguments(model, pad_token=pad_token, max_length=max_length), + **_static_model_arguments(model, max_length=max_length), n_layers=n_layers, hidden_dim=hidden_dim, out_dim=out_dim, diff --git a/model2vec/train/pairs.py b/model2vec/train/pairs.py index 62b89a7..e1f46fd 100644 --- a/model2vec/train/pairs.py +++ b/model2vec/train/pairs.py @@ -66,7 +66,6 @@ def __init__( n_layers: int = 1, hidden_dim: int = 512, out_dim: int | None = None, - pad_id: int | None = None, token_mapping: list[int] | None = None, weights: torch.Tensor | None = None, freeze: bool = False, @@ -82,7 +81,6 @@ def __init__( dimension, the model has no head, and the embeddings are used as is. :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: Deprecated and ignored. :param token_mapping: The token mapping. If None, the token mapping is set to the range of the number of vectors. :param weights: The token weights of the model. If None and `freeze_weights` is not False, the model has no token weights and takes the unweighted mean of the token embeddings. If None and `freeze_weights` @@ -96,7 +94,6 @@ def __init__( 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, @@ -115,7 +112,6 @@ def from_pretrained( *, token: str | None = None, model_name: PathLike | None = None, - pad_token: str | None = None, max_length: int | None | _UnsetType = _UNSET, n_layers: int = 1, hidden_dim: int = 512, @@ -130,7 +126,6 @@ def from_pretrained( :param path: The path to the folder containing the model, or a repository on the Hugging Face Hub. :param token: The token to use to download the model from the hub. :param model_name: Deprecated alias for `path`. - :param pad_token: Deprecated and ignored. :param max_length: The default maximum sequence length to use for tokenization. If not passed, the static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of layers in the head. If this is 0 and `out_dim` equals the embedding @@ -147,7 +142,6 @@ def from_pretrained( model = _load_static_model(path, token=token, model_name=model_name) return cls.from_static_model( model=model, - pad_token=pad_token, max_length=max_length, n_layers=n_layers, hidden_dim=hidden_dim, @@ -163,7 +157,6 @@ def from_static_model( cls: type[T], *, model: StaticModel, - pad_token: str | None = None, max_length: int | None | _UnsetType = _UNSET, n_layers: int = 1, hidden_dim: int = 512, @@ -176,7 +169,6 @@ def from_static_model( """Load the model from a static model. :param model: The static model to load from. - :param pad_token: Deprecated and ignored. :param max_length: The default maximum sequence length to use for tokenization. If not passed, the static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of layers in the head. If this is 0 and `out_dim` equals the embedding @@ -191,7 +183,7 @@ def from_static_model( :return: The initialized model. """ return cls( - **_static_model_arguments(model, pad_token=pad_token, max_length=max_length), + **_static_model_arguments(model, max_length=max_length), n_layers=n_layers, hidden_dim=hidden_dim, out_dim=out_dim, diff --git a/model2vec/train/similarity.py b/model2vec/train/similarity.py index 74e853f..86568a0 100644 --- a/model2vec/train/similarity.py +++ b/model2vec/train/similarity.py @@ -74,7 +74,6 @@ def __init__( n_layers: int = 1, hidden_dim: int = 512, out_dim: int = 2, - pad_id: int | None = None, token_mapping: list[int] | None = None, weights: torch.Tensor | None = None, freeze: bool = False, @@ -86,7 +85,6 @@ def __init__( super().__init__( vectors=vectors, out_dim=out_dim, - pad_id=pad_id, tokenizer=tokenizer, token_mapping=token_mapping, weights=weights, @@ -105,7 +103,6 @@ def from_pretrained( *, token: str | None = None, model_name: PathLike | None = None, - pad_token: str | None = None, max_length: int | None | _UnsetType = _UNSET, n_layers: int = 1, hidden_dim: int = 512, @@ -120,7 +117,6 @@ def from_pretrained( :param path: The path to the folder containing the model, or a repository on the Hugging Face Hub. :param token: The token to use to download the model from the hub. :param model_name: Deprecated alias for `path`. - :param pad_token: Deprecated and ignored. :param max_length: The default maximum sequence length to use for tokenization. If not passed, the static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of hidden layers in the head. @@ -136,7 +132,6 @@ def from_pretrained( model = _load_static_model(path, token=token, model_name=model_name) return cls.from_static_model( model=model, - pad_token=pad_token, max_length=max_length, n_layers=n_layers, hidden_dim=hidden_dim, @@ -152,7 +147,6 @@ def from_static_model( cls: type[T], *, model: StaticModel, - pad_token: str | None = None, max_length: int | None | _UnsetType = _UNSET, n_layers: int = 1, hidden_dim: int = 512, @@ -165,7 +159,6 @@ def from_static_model( """Load the model from a static model. :param model: The static model to load from. - :param pad_token: Deprecated and ignored. :param max_length: The default maximum sequence length to use for tokenization. If not passed, the static model's `max_length` is used. Pass None to disable truncation. :param n_layers: The number of hidden layers in the head. @@ -179,7 +172,7 @@ def from_static_model( :return: The initialized model. """ return cls( - **_static_model_arguments(model, pad_token=pad_token, max_length=max_length), + **_static_model_arguments(model, max_length=max_length), n_layers=n_layers, hidden_dim=hidden_dim, out_dim=out_dim, diff --git a/tests/test_trainable.py b/tests/test_trainable.py index 8a3f244..ca1cc35 100644 --- a/tests/test_trainable.py +++ b/tests/test_trainable.py @@ -129,10 +129,8 @@ def test_init_classifier_from_model(mock_vectors: np.ndarray, mock_tokenizer: To "token", "model_name", "model", - "pad_token", "vectors", "tokenizer", - "pad_id", "token_mapping", "weights", "max_length", @@ -275,31 +273,6 @@ def test_unfrozen_weights_start_at_one(mock_vectors: np.ndarray, mock_tokenizer: assert torch.allclose(weighted._encode(batch), unweighted._encode(batch)) -@pytest.mark.parametrize("cls", FINETUNEABLE_CLASSES) -def test_pad_arguments_are_deprecated( - cls: type[BaseFinetuneable], - mock_vectors: np.ndarray, - mock_tokenizer: Tokenizer, - caplog: pytest.LogCaptureFixture, -) -> None: - """Passing `pad_token` or `pad_id` logs a deprecation warning and has no effect.""" - model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer) - with caplog.at_level(logging.WARNING, logger="model2vec.train.base"): - s = cls.from_static_model(model=model, pad_token="[BRR]") - assert "The 'pad_token' argument is deprecated and ignored." in caplog.text - assert not hasattr(s, "pad_id") - - caplog.clear() - with caplog.at_level(logging.WARNING, logger="model2vec.train.base"): - cls(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer, pad_id=5) - assert "The 'pad_id' argument is deprecated and ignored." in caplog.text - - caplog.clear() - with caplog.at_level(logging.WARNING, logger="model2vec.train.base"): - cls.from_static_model(model=model) - assert "deprecated" not in caplog.text - - def _sequences(batch: TokenBatch) -> list[list[int]]: return [ids.tolist() for ids in batch.ids.split(batch.lengths.tolist())]