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

fastembed/common/preprocessor_utils.py

@@ -51,8 +51,19 @@ 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)
     if not tokenizer.padding:
+        pad_to_multiple_of = (
+            tokenizer_config.get("pad_to_multiple_of")
+            if tokenizer_config.get("pad_to_multiple_of") is not None
+            else (
+                config.get("pad_to_multiple_of")
+                if config.get("pad_to_multiple_of") is not None
+                else tokenizer_config.get("pad_to_multiple_of")
+            )
+        )
         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["pad_token"],
+            pad_to_multiple_of=pad_to_multiple_of,
         )
 
     for token in tokens_map.values():

tests/test_common.py

@@ -59,3 +59,28 @@ 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_with_pad_to_multiple_of(tmp_path):
+    import json
+    from fastembed.common.preprocessor_utils import load_tokenizer
+
+    model_dir = tmp_path / "model"
+    model_dir.mkdir()
+
+    (model_dir / "config.json").write_text(json.dumps({"pad_token_id": 0}))
+    (model_dir / "tokenizer_config.json").write_text(
+        json.dumps({"model_max_length": 128, "pad_token": "<pad>", "pad_to_multiple_of": 8})
+    )
+    (model_dir / "special_tokens_map.json").write_text(json.dumps({}))
+
+    from tokenizers import Tokenizer
+    from tokenizers.models import BPE
+
+    tokenizer = Tokenizer(BPE())
+    tokenizer.save(str(model_dir / "tokenizer.json"))
+
+    loaded_tokenizer, _ = load_tokenizer(model_dir)
+
+    assert loaded_tokenizer.padding is not None
+    assert loaded_tokenizer.padding.get("pad_to_multiple_of") == 8