fix: support explicit max_length override during tokenizer loading

fastembed/common/preprocessor_utils.py

@@ -71,7 +71,9 @@ def _resolve_max_context(tokenizer_config: dict[str, Any], model_dir: Path) -> i
     return min(candidates)
 
 
-def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
+def load_tokenizer(
+    model_dir: Path, max_length: int | None = None
+) -> tuple[Tokenizer, dict[str, int]]:
     tokenizer_path = model_dir / "tokenizer.json"
     if not tokenizer_path.exists():
         raise ValueError(f"Could not find tokenizer.json in {model_dir}")
@@ -90,7 +92,15 @@ def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
     with open(str(tokenizer_config_path)) as tokenizer_config_file:
         tokenizer_config = json.load(tokenizer_config_file)
 
-    max_context = _resolve_max_context(tokenizer_config, model_dir)
+    if max_length is not None:
+        valid_max_length = _valid_context(max_length)
+        if valid_max_length is None:
+            raise ValueError(
+                f"max_length must be a positive integer <= {sys.maxsize}, got {max_length}"
+            )
+        max_context = valid_max_length
+    else:
+        max_context = _resolve_max_context(tokenizer_config, model_dir)
 
     tokens_map = load_special_tokens(model_dir)
 

fastembed/text/custom_text_embedding.py

@@ -65,6 +65,7 @@ def _get_worker_class(cls) -> Type["TextEmbeddingWorker[NumpyArray]"]:
 
     def _get_worker_init_kwargs(self) -> dict[str, Any]:
         return {
+            **super()._get_worker_init_kwargs(),
             "model_description": self.model_description,
             "postprocessing_config": self.POSTPROCESSING_MAPPING[self.model_description.model],
         }

fastembed/text/onnx_embedding.py

@@ -270,6 +270,7 @@ def __init__(
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
+        self.max_length = kwargs.get("max_length")
         self._extra_session_options = self._select_exposed_session_options(kwargs)
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
@@ -336,6 +337,12 @@ def embed(
     def _get_worker_class(cls) -> Type["TextEmbeddingWorker[NumpyArray]"]:
         return OnnxTextEmbeddingWorker
 
+    def _get_worker_init_kwargs(self) -> dict[str, Any]:
+        return {
+            **super()._get_worker_init_kwargs(),
+            "max_length": self.max_length,
+        }
+
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], **kwargs: Any
     ) -> dict[str, NumpyArray]:
@@ -366,6 +373,7 @@ def load_onnx_model(self) -> None:
             cuda=self.cuda,
             device_id=self.device_id,
             extra_session_options=self._extra_session_options,
+            max_length=self.max_length,
         )
 
     def token_count(

fastembed/text/onnx_text_model.py

@@ -55,6 +55,7 @@ def _load_onnx_model(
         cuda: bool | Device = Device.AUTO,
         device_id: int | None = None,
         extra_session_options: dict[str, Any] | None = None,
+        max_length: int | None = None,
     ) -> None:
         super()._load_onnx_model(
             model_dir=model_dir,
@@ -65,7 +66,9 @@ def _load_onnx_model(
             device_id=device_id,
             extra_session_options=extra_session_options,
         )
-        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
+        self.tokenizer, self.special_token_to_id = load_tokenizer(
+            model_dir=model_dir, max_length=max_length
+        )
 
     def load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")

tests/test_preprocessor_utils.py

@@ -331,3 +331,21 @@ def test_list_valued_map_entries_are_registered(make_model_dir, additional) -> N
     _, specials = load_tokenizer(model_dir)
 
     assert "<|list_str|>" in specials
+
+
+def test_explicit_max_length_override(make_model_dir) -> None:
+    """User-provided max_length takes precedence over serialized tokenizer configs."""
+    model_dir = make_model_dir(
+        tokenizer_config={"model_max_length": 512, "max_length": 128},
+    )
+
+    tokenizer, _ = load_tokenizer(model_dir, max_length=256)
+    assert tokenizer.truncation["max_length"] == 256
+
+
+def test_invalid_explicit_max_length_raises(make_model_dir) -> None:
+    model_dir = make_model_dir()
+
+    with pytest.raises(ValueError, match="max_length must be a positive integer"):
+        load_tokenizer(model_dir, max_length=0)
+

tests/test_text_onnx_embeddings.py

@@ -331,3 +331,14 @@ def test_mixed_length_batch_with_fixed_padding(model_cache, model_name: str, dim
 
         embeddings = np.stack(list(model.embed(["hello world", "retrieval " * 200])), axis=0)
         assert embeddings.shape == (2, dim)
+
+
+def test_onnx_text_embedding_worker_init_kwargs_preserves_max_length(model_cache) -> None:
+    with model_cache("BAAI/bge-small-en-v1.5"):
+        embedding = TextEmbedding(
+            model_name="BAAI/bge-small-en-v1.5",
+            lazy_load=True,
+            max_length=256,
+        )
+        assert embedding.model._get_worker_init_kwargs().get("max_length") == 256
+