diff --git a/fastembed/common/preprocessor_utils.py b/fastembed/common/preprocessor_utils.py index 934b234a..2fbd9986 100644 --- a/fastembed/common/preprocessor_utils.py +++ b/fastembed/common/preprocessor_utils.py @@ -98,6 +98,7 @@ def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]: pad_id=padding.get("pad_id", config.get("pad_token_id", 0)), pad_type_id=padding.get("pad_type_id", 0), pad_token=pad_token, + pad_to_multiple_of=padding.get("pad_to_multiple_of"), length=None, ) diff --git a/tests/test_preprocessor_utils.py b/tests/test_preprocessor_utils.py index 5b899529..23cfc58c 100644 --- a/tests/test_preprocessor_utils.py +++ b/tests/test_preprocessor_utils.py @@ -138,6 +138,21 @@ def test_serialized_left_padding_is_preserved(make_model_dir) -> None: assert encoded[1].attention_mask[0] == 0 +def test_serialized_pad_to_multiple_of_is_preserved(make_model_dir) -> None: + """Everything the tokenizer declared is kept; only the fixed length is overridden.""" + model_dir = make_model_dir( + padding={"length": 128, "pad_id": 0, "pad_token": "[PAD]", "pad_to_multiple_of": 8}, + ) + + tokenizer, _ = load_tokenizer(model_dir) + + assert tokenizer.padding["length"] is None + assert tokenizer.padding["pad_to_multiple_of"] == 8 + + encoded = tokenizer.encode_batch(["hello world", "hello"]) + assert len(encoded[0].ids) % 8 == 0 + + def test_serialized_pad_id_takes_precedence_over_config(make_model_dir) -> None: mask_pad_id = 103 # [MASK] in the bert-base vocab, any id other than the config's works model_dir = make_model_dir(