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/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 ff527fc6f..8d7dd1225 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 @@ -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( 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) + diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index 4ace42ec9..502056903 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/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 +