diff --git a/fastembed/common/preprocessor_utils.py b/fastembed/common/preprocessor_utils.py index 3b702f79..dedb7e0a 100644 --- a/fastembed/common/preprocessor_utils.py +++ b/fastembed/common/preprocessor_utils.py @@ -1,6 +1,6 @@ import json -from typing import Any from pathlib import Path +from typing import Any from tokenizers import AddedToken, Tokenizer @@ -19,6 +19,25 @@ def load_special_tokens(model_dir: Path) -> dict[str, Any]: def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]: + """ + Load and configure a tokenizer from a model directory. + + Configures truncation to the model context length, converts any fixed-length + padding to dynamic batch padding (avoiding ragged batch failures), preserves + serialized padding direction and token IDs, and optionally applies padding + multiples (e.g. pad_to_multiple_of) when configured. + + Args: + model_dir: Directory path containing tokenizer configuration files + (config.json, tokenizer.json, tokenizer_config.json, special_tokens_map.json). + + Returns: + A tuple of (configured Tokenizer instance, mapping of special token strings to token IDs). + + Raises: + ValueError: If required configuration files are missing or if pad_to_multiple_of + is not a positive integer. + """ config_path = model_dir / "config.json" if not config_path.exists(): raise ValueError(f"Could not find config.json in {model_dir}") @@ -50,10 +69,52 @@ def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]: tokenizer = Tokenizer.from_file(str(tokenizer_path)) tokenizer.enable_truncation(max_length=max_context) + + pad_to_multiple_of = tokenizer_config.get("pad_to_multiple_of") + if pad_to_multiple_of is None: + pad_to_multiple_of = config.get("pad_to_multiple_of") + + if pad_to_multiple_of is not None and ( + not isinstance(pad_to_multiple_of, int) + or isinstance(pad_to_multiple_of, bool) + or pad_to_multiple_of <= 0 + ): + raise ValueError("pad_to_multiple_of must be a positive integer") + if not tokenizer.padding: tokenizer.enable_padding( - pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"] + pad_id=config.get("pad_token_id", 0), + pad_token=tokenizer_config.get("pad_token", "[PAD]"), + pad_to_multiple_of=pad_to_multiple_of, + ) + else: + padding_params = tokenizer.padding + target_pad_to_multiple_of = ( + pad_to_multiple_of + if pad_to_multiple_of is not None + else padding_params.get("pad_to_multiple_of") ) + if target_pad_to_multiple_of is not None and ( + not isinstance(target_pad_to_multiple_of, int) + or isinstance(target_pad_to_multiple_of, bool) + or target_pad_to_multiple_of <= 0 + ): + raise ValueError("pad_to_multiple_of must be a positive integer") + + if padding_params.get("length") is not None or ( + pad_to_multiple_of is not None + and padding_params.get("pad_to_multiple_of") != target_pad_to_multiple_of + ): + tokenizer.enable_padding( + direction=padding_params.get("direction", "right"), + pad_id=padding_params.get("pad_id", config.get("pad_token_id", 0)), + pad_type_id=padding_params.get("pad_type_id", 0), + pad_token=padding_params.get( + "pad_token", tokenizer_config.get("pad_token", "[PAD]") + ), + pad_to_multiple_of=target_pad_to_multiple_of, + length=None, + ) for token in tokens_map.values(): if isinstance(token, str): diff --git a/tests/test_common.py b/tests/test_common.py index f7cae5ba..744219fa 100644 --- a/tests/test_common.py +++ b/tests/test_common.py @@ -1,11 +1,11 @@ import numpy as np from fastembed import ( - TextEmbedding, - SparseTextEmbedding, ImageEmbedding, LateInteractionMultimodalEmbedding, LateInteractionTextEmbedding, + SparseTextEmbedding, + TextEmbedding, ) from fastembed.common.utils import last_token_pooling @@ -23,13 +23,13 @@ def test_text_list_supported_models(): description = supported_models[0] assert isinstance(description, dict) - assert "model" in description and description["model"] + assert description.get("model") if model_type != SparseTextEmbedding: - assert "dim" in description and description["dim"] - assert "license" in description and description["license"] - assert "size_in_GB" in description and description["size_in_GB"] - assert "model_file" in description and description["model_file"] - assert "sources" in description and description["sources"] + assert description.get("dim") + assert description.get("license") + assert description.get("size_in_GB") + assert description.get("model_file") + assert description.get("sources") assert "hf" in description["sources"] or "url" in description["sources"] @@ -59,3 +59,148 @@ def test_last_token_pooling_with_left_padding(): pooled = last_token_pooling(token_embeddings, attention_mask) assert np.allclose(pooled, [[2.0, 2.0], [6.0, 6.0]]) + + +def test_load_tokenizer_fixed_length_padding_converted_to_dynamic(tmp_path): + """ + Verify that models with serialized fixed-length padding (e.g. gte-base with length=128) + have their padding relaxed to dynamic batch padding (length=None) to support mixed-length batches. + """ + import json + + from tokenizers import Tokenizer, models + + from fastembed.common.preprocessor_utils import load_tokenizer + + config = {"pad_token_id": 0} + with open(tmp_path / "config.json", "w") as f: + json.dump(config, f) + + tokenizer_config = { + "model_max_length": 512, + "pad_token": "[PAD]", + } + with open(tmp_path / "tokenizer_config.json", "w") as f: + json.dump(tokenizer_config, f) + + with open(tmp_path / "special_tokens_map.json", "w") as f: + json.dump({"pad_token": "[PAD]"}, f) + + # Tokenizer initialized with fixed-length padding (e.g. gte-base with length=128) + tokenizer = Tokenizer(models.BPE()) + tokenizer.add_special_tokens(["[PAD]"]) + tokenizer.enable_padding(length=128, pad_id=0, pad_token="[PAD]", direction="right") + tokenizer.save(str(tmp_path / "tokenizer.json")) + + loaded_tokenizer, _ = load_tokenizer(tmp_path) + + # Fixed length must be relaxed to None (dynamic batch padding) to prevent ragged arrays + assert loaded_tokenizer.padding is not None + assert loaded_tokenizer.padding["length"] is None + assert loaded_tokenizer.padding["direction"] == "right" + assert loaded_tokenizer.padding["pad_id"] == 0 + assert loaded_tokenizer.padding["pad_token"] == "[PAD]" + assert loaded_tokenizer.truncation["max_length"] == 512 + + +def test_load_tokenizer_preserves_left_padding(tmp_path): + """ + Verify that tokenizers with serialized left padding (e.g. ColModernVBERT) + preserve their padding direction and pad_token_id without being overridden. + """ + import json + + from tokenizers import Tokenizer, models + + from fastembed.common.preprocessor_utils import load_tokenizer + + config = {"pad_token_id": 50283} + with open(tmp_path / "config.json", "w") as f: + json.dump(config, f) + + tokenizer_config = { + "model_max_length": 8192, + "pad_token": "[PAD]", + } + with open(tmp_path / "tokenizer_config.json", "w") as f: + json.dump(tokenizer_config, f) + + with open(tmp_path / "special_tokens_map.json", "w") as f: + json.dump({"pad_token": "[PAD]"}, f) + + # Tokenizer with dynamic left padding (e.g. ColModernVBERT) + tokenizer = Tokenizer(models.BPE()) + tokenizer.add_special_tokens(["[PAD]"]) + tokenizer.enable_padding(length=None, pad_id=50283, pad_token="[PAD]", direction="left") + tokenizer.save(str(tmp_path / "tokenizer.json")) + + loaded_tokenizer, _ = load_tokenizer(tmp_path) + + # Preserves left padding direction and pad_id + assert loaded_tokenizer.padding is not None + assert loaded_tokenizer.padding["length"] is None + assert loaded_tokenizer.padding["direction"] == "left" + assert loaded_tokenizer.padding["pad_id"] == 50283 + + +def test_load_tokenizer_pad_to_multiple_of_and_validation(tmp_path): + """ + Verify pad_to_multiple_of configuration precedence (tokenizer_config > config fallback), + application to unpadded and serialized tokenizers, and validation of positive integer values. + """ + import json + + import pytest + from tokenizers import Tokenizer, models + + from fastembed.common.preprocessor_utils import load_tokenizer + + # 1. Test pad_to_multiple_of via config.json fallback + dir_fallback = tmp_path / "fallback" + dir_fallback.mkdir() + (dir_fallback / "config.json").write_text( + json.dumps({"pad_token_id": 0, "pad_to_multiple_of": 16}) + ) + (dir_fallback / "tokenizer_config.json").write_text( + json.dumps({"model_max_length": 128, "pad_token": "[PAD]"}) + ) + (dir_fallback / "special_tokens_map.json").write_text(json.dumps({"pad_token": "[PAD]"})) + tok = Tokenizer(models.BPE()) + tok.add_special_tokens(["[PAD]"]) + tok.save(str(dir_fallback / "tokenizer.json")) + + loaded_tok, _ = load_tokenizer(dir_fallback) + assert loaded_tok.padding is not None + assert loaded_tok.padding.get("pad_to_multiple_of") == 16 + + # 2. Test tokenizer_config.json precedence over config.json + dir_prec = tmp_path / "prec" + dir_prec.mkdir() + (dir_prec / "config.json").write_text( + json.dumps({"pad_token_id": 0, "pad_to_multiple_of": 16}) + ) + (dir_prec / "tokenizer_config.json").write_text( + json.dumps({"model_max_length": 128, "pad_token": "[PAD]", "pad_to_multiple_of": 8}) + ) + (dir_prec / "special_tokens_map.json").write_text(json.dumps({"pad_token": "[PAD]"})) + tok.save(str(dir_prec / "tokenizer.json")) + + loaded_tok, _ = load_tokenizer(dir_prec) + assert loaded_tok.padding is not None + assert loaded_tok.padding.get("pad_to_multiple_of") == 8 + + # 3. Test validation error on invalid pad_to_multiple_of (<= 0 or non-integer) + for invalid_val in [0, -1, "8", False]: + dir_invalid = tmp_path / f"invalid_{invalid_val}" + dir_invalid.mkdir() + (dir_invalid / "config.json").write_text(json.dumps({"pad_token_id": 0})) + (dir_invalid / "tokenizer_config.json").write_text( + json.dumps( + {"model_max_length": 128, "pad_token": "[PAD]", "pad_to_multiple_of": invalid_val} + ) + ) + (dir_invalid / "special_tokens_map.json").write_text(json.dumps({"pad_token": "[PAD]"})) + tok.save(str(dir_invalid / "tokenizer.json")) + + with pytest.raises(ValueError, match="pad_to_multiple_of must be a positive integer"): + load_tokenizer(dir_invalid)