From 85fcf05f248bea94c6b93053bc3d82983e86ceb4 Mon Sep 17 00:00:00 2001 From: Mohammed Alshyakh Date: Wed, 23 Sep 2026 02:53:22 +0200 Subject: [PATCH 1/3] fix: support explicit max_length override during tokenizer loading --- fastembed/common/preprocessor_utils.py | 14 ++++++++++++-- fastembed/text/onnx_embedding.py | 2 ++ fastembed/text/onnx_text_model.py | 5 ++++- tests/test_preprocessor_utils.py | 18 ++++++++++++++++++ 4 files changed, 36 insertions(+), 3 deletions(-) diff --git a/fastembed/common/preprocessor_utils.py b/fastembed/common/preprocessor_utils.py index c08658db9..264da8528 100644 --- a/fastembed/common/preprocessor_utils.py +++ b/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) diff --git a/fastembed/text/onnx_embedding.py b/fastembed/text/onnx_embedding.py index ff527fc6f..7872c1486 100644 --- a/fastembed/text/onnx_embedding.py +++ b/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 @@ -366,6 +367,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( diff --git a/fastembed/text/onnx_text_model.py b/fastembed/text/onnx_text_model.py index feeb70d55..d8b1fab28 100644 --- a/fastembed/text/onnx_text_model.py +++ b/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") diff --git a/tests/test_preprocessor_utils.py b/tests/test_preprocessor_utils.py index 0b928f504..e4cf38cb4 100644 --- a/tests/test_preprocessor_utils.py +++ b/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) + From 4fc73915ab3cea8d3ce09f746b8d7a6435902307 Mon Sep 17 00:00:00 2001 From: Mohammed Alshyakh Date: Wed, 23 Sep 2026 03:06:27 +0200 Subject: [PATCH 2/3] fix: propagate max_length to parallel worker init kwargs --- fastembed/text/custom_text_embedding.py | 1 + fastembed/text/onnx_embedding.py | 6 ++++++ tests/test_text_onnx_embeddings.py | 7 +++++++ 3 files changed, 14 insertions(+) diff --git a/fastembed/text/custom_text_embedding.py b/fastembed/text/custom_text_embedding.py index c3d5cf5f0..e56668ed5 100644 --- a/fastembed/text/custom_text_embedding.py +++ b/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], } diff --git a/fastembed/text/onnx_embedding.py b/fastembed/text/onnx_embedding.py index 7872c1486..8d7dd1225 100644 --- a/fastembed/text/onnx_embedding.py +++ b/fastembed/text/onnx_embedding.py @@ -337,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]: diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index 4ace42ec9..767bfe24d 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/tests/test_text_onnx_embeddings.py @@ -331,3 +331,10 @@ 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() -> None: + embedding = OnnxTextEmbedding.__new__(OnnxTextEmbedding) + embedding.max_length = 256 + assert embedding._get_worker_init_kwargs().get("max_length") == 256 + From 289ae10fe2f48886bdd04c287fbb4a2f65af6b70 Mon Sep 17 00:00:00 2001 From: Mohammed Alshyakh Date: Wed, 23 Sep 2026 03:15:25 +0200 Subject: [PATCH 3/3] test: exercise public constructor in worker max_length test --- tests/test_text_onnx_embeddings.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index 767bfe24d..502056903 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/tests/test_text_onnx_embeddings.py @@ -333,8 +333,12 @@ def test_mixed_length_batch_with_fixed_padding(model_cache, model_name: str, dim assert embeddings.shape == (2, dim) -def test_onnx_text_embedding_worker_init_kwargs_preserves_max_length() -> None: - embedding = OnnxTextEmbedding.__new__(OnnxTextEmbedding) - embedding.max_length = 256 - assert embedding._get_worker_init_kwargs().get("max_length") == 256 +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