feat: make max length argument explicit

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,
     }

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`.

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.

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`.

tests/test_trainable.py

@@ -493,6 +493,28 @@ 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)
+@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],
+    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)
+    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:
     """Test the predict function."""
     result = mock_trained_pipeline.predict(["dog cat", "dog"]).tolist()