From 8119e4c67d0a373a8f0e400675f6b0cf49aabcd1 Mon Sep 17 00:00:00 2001 From: stephantul Date: Mon, 5 Oct 2026 21:13:15 +0200 Subject: [PATCH 1/2] feat: make max length argument explicit --- model2vec/train/base.py | 23 +++++++++++++---------- model2vec/train/classifier.py | 13 +++++++------ model2vec/train/pairs.py | 13 +++++++------ model2vec/train/similarity.py | 13 +++++++------ tests/test_trainable.py | 11 +++++++++++ 5 files changed, 45 insertions(+), 28 deletions(-) diff --git a/model2vec/train/base.py b/model2vec/train/base.py index bc47e27..2a487e9 100644 --- a/model2vec/train/base.py +++ b/model2vec/train/base.py @@ -24,6 +24,7 @@ split_indices, to_pipeline, ) +from model2vec.types import _UNSET, _UnsetType logger = logging.getLogger(__name__) @@ -175,7 +176,7 @@ def from_pretrained( token: str | None = None, model_name: PathLike | None = None, pad_token: str | None = None, - max_length: int | None = None, + max_length: int | None | _UnsetType = _UNSET, n_layers: int = 0, hidden_dim: int = 256, out_dim: int = 2, @@ -190,8 +191,8 @@ def from_pretrained( :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 max_length: The default maximum sequence length to use for tokenization. If None, the - static model's `max_length` is used. + :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. :param hidden_dim: The hidden dimension of the head. :param out_dim: The output dimension of the head. @@ -221,7 +222,7 @@ def from_static_model( *, model: StaticModel, pad_token: str | None = None, - max_length: int | None = None, + max_length: int | None | _UnsetType = _UNSET, n_layers: int = 0, hidden_dim: int = 256, out_dim: int = 2, @@ -234,8 +235,8 @@ def from_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 max_length: The default maximum sequence length to use for tokenization. If None, the - static model's `max_length` is used. + :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. :param hidden_dim: The hidden dimension of the head. :param out_dim: The output dimension of the head. @@ -565,13 +566,15 @@ 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) -> dict[str, Any]: +def _static_model_arguments( + model: StaticModel, *, pad_token: str | None, 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: The token to use for padding. If None, it is inferred from the tokenizer. - :param max_length: The default maximum sequence length to use for tokenization. If None, the - static model's `max_length` is used. + :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. """ model.embedding = np.nan_to_num(model.embedding) @@ -587,5 +590,5 @@ def _static_model_arguments(model: StaticModel, *, pad_token: str | None, max_le "tokenizer": model.tokenizer, "token_mapping": token_mapping, "weights": weights, - "max_length": model.max_length if max_length is None else max_length, + "max_length": model.max_length if isinstance(max_length, _UnsetType) else max_length, } diff --git a/model2vec/train/classifier.py b/model2vec/train/classifier.py index 3c163ae..4a927d2 100644 --- a/model2vec/train/classifier.py +++ b/model2vec/train/classifier.py @@ -18,6 +18,7 @@ from model2vec.train.base import BaseFinetuneable, _load_static_model, _static_model_arguments from model2vec.train.dataset import read_label_column from model2vec.train.utils import DEFAULT_RANDOM_SEED, seed_everything +from model2vec.types import _UNSET, _UnsetType logger = logging.getLogger(__name__) @@ -119,7 +120,7 @@ def from_pretrained( token: str | None = None, model_name: PathLike | None = None, pad_token: str | None = None, - max_length: int | None = None, + max_length: int | None | _UnsetType = _UNSET, n_layers: int = 1, hidden_dim: int = 512, out_dim: int = 2, @@ -134,8 +135,8 @@ def from_pretrained( :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 max_length: The default maximum sequence length to use for tokenization. If None, the - static model's `max_length` is used. + :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. :param hidden_dim: The hidden dimension of the head. :param out_dim: The number of classes. This is reset when calling `fit`. @@ -165,7 +166,7 @@ def from_static_model( *, model: StaticModel, pad_token: str | None = None, - max_length: int | None = None, + max_length: int | None | _UnsetType = _UNSET, n_layers: int = 1, hidden_dim: int = 512, out_dim: int = 2, @@ -178,8 +179,8 @@ def from_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 max_length: The default maximum sequence length to use for tokenization. If None, the - static model's `max_length` is used. + :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. :param hidden_dim: The hidden dimension of the head. :param out_dim: The number of classes. This is reset when calling `fit`. diff --git a/model2vec/train/pairs.py b/model2vec/train/pairs.py index a9e792c..0f3e17c 100644 --- a/model2vec/train/pairs.py +++ b/model2vec/train/pairs.py @@ -13,6 +13,7 @@ from model2vec.train.base import BaseFinetuneable, _load_static_model, _static_model_arguments from model2vec.train.dataset import ColumnRows, PairDataset from model2vec.train.utils import DEFAULT_RANDOM_SEED, MAX_VALIDATION_SIZE, seed_everything, split_indices +from model2vec.types import _UNSET, _UnsetType logger = logging.getLogger(__name__) @@ -112,7 +113,7 @@ def from_pretrained( token: str | None = None, model_name: PathLike | None = None, pad_token: str | None = None, - max_length: int | None = None, + max_length: int | None | _UnsetType = _UNSET, n_layers: int = 1, hidden_dim: int = 512, out_dim: int | None = None, @@ -127,8 +128,8 @@ def from_pretrained( :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 max_length: The default maximum sequence length to use for tokenization. If None, the - static model's `max_length` is used. + :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 dimension, the model has no head, and the embeddings are used as is. :param hidden_dim: The hidden dimension of the head. @@ -159,7 +160,7 @@ def from_static_model( *, model: StaticModel, pad_token: str | None = None, - max_length: int | None = None, + max_length: int | None | _UnsetType = _UNSET, n_layers: int = 1, hidden_dim: int = 512, out_dim: int | None = None, @@ -172,8 +173,8 @@ def from_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 max_length: The default maximum sequence length to use for tokenization. If None, the - static model's `max_length` is used. + :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 dimension, the model has no head, and the embeddings are used as is. :param hidden_dim: The hidden dimension of the head. diff --git a/model2vec/train/similarity.py b/model2vec/train/similarity.py index 4928b07..9ae9261 100644 --- a/model2vec/train/similarity.py +++ b/model2vec/train/similarity.py @@ -14,6 +14,7 @@ from model2vec.train.base import BaseFinetuneable, _load_static_model, _static_model_arguments from model2vec.train.dataset import get_vector_dims_from_column from model2vec.train.utils import DEFAULT_RANDOM_SEED, seed_everything +from model2vec.types import _UNSET, _UnsetType logger = logging.getLogger(__name__) @@ -105,7 +106,7 @@ def from_pretrained( token: str | None = None, model_name: PathLike | None = None, pad_token: str | None = None, - max_length: int | None = None, + max_length: int | None | _UnsetType = _UNSET, n_layers: int = 1, hidden_dim: int = 512, out_dim: int = 2, @@ -120,8 +121,8 @@ def from_pretrained( :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 max_length: The default maximum sequence length to use for tokenization. If None, the - static model's `max_length` is used. + :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. :param hidden_dim: The hidden dimension of the head. :param out_dim: The output dimension of the head. This is reset when calling `fit`. @@ -151,7 +152,7 @@ def from_static_model( *, model: StaticModel, pad_token: str | None = None, - max_length: int | None = None, + max_length: int | None | _UnsetType = _UNSET, n_layers: int = 1, hidden_dim: int = 512, out_dim: int = 2, @@ -164,8 +165,8 @@ def from_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 max_length: The default maximum sequence length to use for tokenization. If None, the - static model's `max_length` is used. + :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. :param hidden_dim: The hidden dimension of the head. :param out_dim: The output dimension of the head. This is reset when calling `fit`. diff --git a/tests/test_trainable.py b/tests/test_trainable.py index ddd99a6..4d33359 100644 --- a/tests/test_trainable.py +++ b/tests/test_trainable.py @@ -493,6 +493,17 @@ def test_max_length_is_not_capped_by_the_static_model(mock_vectors: np.ndarray, assert s.tokenize(texts).shape[1] == 4 +@pytest.mark.parametrize("cls", FINETUNEABLE_CLASSES) +def test_from_static_model_max_length_resolution( + cls: type[BaseFinetuneable], mock_vectors: np.ndarray, mock_tokenizer: Tokenizer +) -> None: + """Unset inherits the static model's `max_length`, None disables truncation, and an int is used as is.""" + static = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer, max_length=2) + assert cls.from_static_model(model=static).max_length == 2 + assert cls.from_static_model(model=static, max_length=None).max_length is None + assert cls.from_static_model(model=static, max_length=5).max_length == 5 + + def test_predict(mock_trained_pipeline: StaticModelForClassification) -> None: """Test the predict function.""" result = mock_trained_pipeline.predict(["dog cat", "dog"]).tolist() From 1a46e7084e0db7238b6602f38245f0809fc5f0ae Mon Sep 17 00:00:00 2001 From: stephantul Date: Mon, 5 Oct 2026 21:20:46 +0200 Subject: [PATCH 2/2] fix test --- tests/test_trainable.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/tests/test_trainable.py b/tests/test_trainable.py index 4d33359..38874c6 100644 --- a/tests/test_trainable.py +++ b/tests/test_trainable.py @@ -494,14 +494,25 @@ def test_max_length_is_not_capped_by_the_static_model(mock_vectors: np.ndarray, @pytest.mark.parametrize("cls", FINETUNEABLE_CLASSES) +@pytest.mark.parametrize( + "kwargs, expected_max_length, expected_tokens", + [({}, 2, 2), ({"max_length": None}, None, 6), ({"max_length": 5}, 5, 5)], +) def test_from_static_model_max_length_resolution( - cls: type[BaseFinetuneable], mock_vectors: np.ndarray, mock_tokenizer: Tokenizer + cls: type[BaseFinetuneable], + kwargs: dict[str, Any], + expected_max_length: int | None, + expected_tokens: int, + mock_vectors: np.ndarray, + mock_tokenizer: Tokenizer, ) -> None: """Unset inherits the static model's `max_length`, None disables truncation, and an int is used as is.""" static = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer, max_length=2) - assert cls.from_static_model(model=static).max_length == 2 - assert cls.from_static_model(model=static, max_length=None).max_length is None - assert cls.from_static_model(model=static, max_length=5).max_length == 5 + texts = ["word1 word2 word3 word1 word2 word3"] + s = cls.from_static_model(model=static, **kwargs) + assert s.max_length == expected_max_length + assert [len(row) for row in s._tokenize_ids(texts)] == [expected_tokens] + assert len(s.to_static_model().tokenize(texts)[0]) == expected_tokens def test_predict(mock_trained_pipeline: StaticModelForClassification) -> None: