fix: support mixed-length batches for models with fixed padding (#703)

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):

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)