From 8409d6e4695b7004c58a7aa9286539b201e041e7 Mon Sep 17 00:00:00 2001 From: Hulk Date: Sun, 4 Oct 2026 12:06:30 +0800 Subject: [PATCH 1/2] fix: preserve serialized tokenizer truncation direction --- fastembed/common/preprocessor_utils.py | 7 +- tests/test_tokenizer_truncation.py | 111 +++++++++++++++++++++++++ 2 files changed, 117 insertions(+), 1 deletion(-) create mode 100644 tests/test_tokenizer_truncation.py diff --git a/fastembed/common/preprocessor_utils.py b/fastembed/common/preprocessor_utils.py index c08658db..d6029db5 100644 --- a/fastembed/common/preprocessor_utils.py +++ b/fastembed/common/preprocessor_utils.py @@ -95,7 +95,12 @@ def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]: tokens_map = load_special_tokens(model_dir) tokenizer = Tokenizer.from_file(str(tokenizer_path)) - tokenizer.enable_truncation(max_length=max_context) + # Updating the context limit must not reset a saved left-truncating tokenizer + # to the tokenizers default of keeping the start of the input. + truncation = tokenizer.truncation or {} + tokenizer.enable_truncation( + max_length=max_context, direction=truncation.get("direction", "right") + ) # Registered before the padding is resolved: the map may name a pad token that # tokenizer.json does not carry, and it only gets an id once it is added. diff --git a/tests/test_tokenizer_truncation.py b/tests/test_tokenizer_truncation.py new file mode 100644 index 00000000..f09e672d --- /dev/null +++ b/tests/test_tokenizer_truncation.py @@ -0,0 +1,111 @@ +"""Offline regression coverage for serialized tokenizer truncation direction.""" + +import json +from pathlib import Path + +import numpy as np +import pytest +from tokenizers import Tokenizer, models, pre_tokenizers, processors + +from fastembed.common.preprocessor_utils import load_tokenizer +from fastembed.text.onnx_text_model import OnnxTextModel + + +def make_tokenizer_dir( + path: Path, + direction: str | None, + context: int, + serialized_context: int = 5, + max_length: int | None = None, + stride: int = 0, + padding_direction: str = "right", +) -> Path: + """Save a tiny real tokenizer without downloading a model or starting ONNX.""" + tokens = ["[UNK]", "[PAD]", "[CLS]", "[SEP]", "one", "two", "three", "four", "five"] + tokenizer = Tokenizer(models.WordLevel(dict(zip(tokens, range(len(tokens)))), "[UNK]")) + tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() + tokenizer.post_processor = processors.TemplateProcessing( + single="[CLS] $A [SEP]", special_tokens=[("[CLS]", 2), ("[SEP]", 3)] + ) + if direction is not None: + tokenizer.enable_truncation( + max_length=serialized_context, direction=direction, stride=stride + ) + tokenizer.enable_padding(pad_id=1, pad_token="[PAD]", direction=padding_direction) + tokenizer.save(str(path / "tokenizer.json")) + (path / "tokenizer_config.json").write_text( + json.dumps({"model_max_length": context, "max_length": max_length, "pad_token": "[PAD]"}) + ) + return path + + +@pytest.mark.parametrize("direction", ["left", "right", None]) +@pytest.mark.parametrize("context", [4, 5, 6]) +def test_load_tokenizer_preserves_truncation_direction(tmp_path, direction, context): + """Changing the context limit must not change which end of the input survives.""" + model_dir = make_tokenizer_dir(tmp_path, direction, context) + tokenizer, _ = load_tokenizer(model_dir) + words = ["one", "two", "three", "four", "five"] + kept_words = words[-(context - 2) :] if direction == "left" else words[: context - 2] + + assert tokenizer.truncation["max_length"] == context + assert tokenizer.truncation["direction"] == (direction or "right") + assert tokenizer.encode(" ".join(words)).tokens == ["[CLS]", *kept_words, "[SEP]"] + + +@pytest.mark.parametrize("direction", ["left", "right"]) +def test_text_model_uses_preserved_direction_in_mixed_length_batch(tmp_path, direction): + """The real text-model tokenizer retains the declared text and still pads batches.""" + model_dir = make_tokenizer_dir(tmp_path, direction, context=5) + model = OnnxTextModel() + model._load_tokenizer(model_dir) + + encoded = model.tokenize(["one two three four five", "one"]) + expected_long = ["three", "four", "five"] if direction == "left" else ["one", "two", "three"] + assert encoded[0].tokens == ["[CLS]", *expected_long, "[SEP]"] + assert encoded[1].tokens == ["[CLS]", "one", "[SEP]", "[PAD]", "[PAD]"] + assert encoded[1].attention_mask == [1, 1, 1, 0, 0] + assert np.array([encoding.ids for encoding in encoded]).shape == (2, 5) + + +@pytest.mark.parametrize("context,max_length,expected_length", [(8, 5, 5), (4, 6, 4)]) +def test_saved_direction_keeps_the_stricter_config_limit( + tmp_path, context, max_length, expected_length +): + """Direction preservation does not replace the loader's context-limit resolution.""" + model_dir = make_tokenizer_dir( + tmp_path, "left", context, serialized_context=8, max_length=max_length + ) + tokenizer, _ = load_tokenizer(model_dir) + words = ["one", "two", "three", "four", "five"] + + assert tokenizer.truncation["max_length"] == expected_length + assert tokenizer.encode(" ".join(words)).tokens == [ + "[CLS]", + *words[-(expected_length - 2) :], + "[SEP]", + ] + + +def test_saved_direction_does_not_restore_stride(tmp_path): + """The loader still disables overflow stride when applying a smaller context.""" + model_dir = make_tokenizer_dir(tmp_path, "left", context=4, serialized_context=8, stride=3) + + tokenizer, _ = load_tokenizer(model_dir) + + assert tokenizer.truncation["direction"] == "left" + assert tokenizer.truncation["stride"] == 0 + assert tokenizer.truncation["strategy"] == "longest_first" + assert tokenizer.encode("one two three four five").tokens == ["[CLS]", "four", "five", "[SEP]"] + + +def test_left_truncation_and_left_padding_remain_independent(tmp_path): + """Padding keeps short items aligned without changing the long item's retained end.""" + model_dir = make_tokenizer_dir(tmp_path, "left", context=5, padding_direction="left") + tokenizer, _ = load_tokenizer(model_dir) + + encoded = tokenizer.encode_batch(["one two three four five", "one"]) + + assert encoded[0].tokens == ["[CLS]", "three", "four", "five", "[SEP]"] + assert encoded[1].tokens == ["[PAD]", "[PAD]", "[CLS]", "one", "[SEP]"] + assert encoded[1].attention_mask == [0, 0, 1, 1, 1] From 90889caa81b588e476a4adb3368e31616b06418c Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Sun, 4 Oct 2026 23:54:51 +0700 Subject: [PATCH 2/2] fix: read truncation_side from tokenizer_config.json before tokenizer.json --- fastembed/common/preprocessor_utils.py | 7 +- tests/test_preprocessor_utils.py | 26 ++++++ tests/test_tokenizer_truncation.py | 111 ------------------------- 3 files changed, 30 insertions(+), 114 deletions(-) delete mode 100644 tests/test_tokenizer_truncation.py diff --git a/fastembed/common/preprocessor_utils.py b/fastembed/common/preprocessor_utils.py index d6029db5..fcd8acaa 100644 --- a/fastembed/common/preprocessor_utils.py +++ b/fastembed/common/preprocessor_utils.py @@ -95,11 +95,12 @@ def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]: tokens_map = load_special_tokens(model_dir) tokenizer = Tokenizer.from_file(str(tokenizer_path)) - # Updating the context limit must not reset a saved left-truncating tokenizer - # to the tokenizers default of keeping the start of the input. + # enable_truncation resets the direction to right unless it is passed. The direction is + # resolved as in transformers: tokenizer_config.json, then tokenizer.json, then right. truncation = tokenizer.truncation or {} tokenizer.enable_truncation( - max_length=max_context, direction=truncation.get("direction", "right") + max_length=max_context, + direction=tokenizer_config.get("truncation_side") or truncation.get("direction", "right"), ) # Registered before the padding is resolved: the map may name a pad token that diff --git a/tests/test_preprocessor_utils.py b/tests/test_preprocessor_utils.py index 0b928f50..b9f718a8 100644 --- a/tests/test_preprocessor_utils.py +++ b/tests/test_preprocessor_utils.py @@ -210,6 +210,32 @@ def test_max_context_resolution(make_model_dir, model_max_length, max_length, ex assert tokenizer.truncation["max_length"] == expected +@pytest.mark.parametrize( + "serialized_direction,truncation_side,expected", + [ + (None, None, "right"), + ("left", None, "left"), + (None, "left", "left"), + ("left", "right", "right"), # tokenizer_config.json wins, as in transformers + ], +) +def test_truncation_direction_resolution( + make_model_dir, serialized_direction, truncation_side, expected +) -> None: + model_dir = make_model_dir(tokenizer_config={"truncation_side": truncation_side}) + tokenizer_path = model_dir / "tokenizer.json" + serialized = Tokenizer.from_file(str(tokenizer_path)) + if serialized_direction is None: + serialized.no_truncation() + else: + serialized.enable_truncation(max_length=512, direction=serialized_direction) + serialized.save(str(tokenizer_path)) + + tokenizer, _ = load_tokenizer(model_dir) + + assert tokenizer.truncation["direction"] == expected + + @pytest.mark.parametrize( "model_max_length,max_length", [ diff --git a/tests/test_tokenizer_truncation.py b/tests/test_tokenizer_truncation.py deleted file mode 100644 index f09e672d..00000000 --- a/tests/test_tokenizer_truncation.py +++ /dev/null @@ -1,111 +0,0 @@ -"""Offline regression coverage for serialized tokenizer truncation direction.""" - -import json -from pathlib import Path - -import numpy as np -import pytest -from tokenizers import Tokenizer, models, pre_tokenizers, processors - -from fastembed.common.preprocessor_utils import load_tokenizer -from fastembed.text.onnx_text_model import OnnxTextModel - - -def make_tokenizer_dir( - path: Path, - direction: str | None, - context: int, - serialized_context: int = 5, - max_length: int | None = None, - stride: int = 0, - padding_direction: str = "right", -) -> Path: - """Save a tiny real tokenizer without downloading a model or starting ONNX.""" - tokens = ["[UNK]", "[PAD]", "[CLS]", "[SEP]", "one", "two", "three", "four", "five"] - tokenizer = Tokenizer(models.WordLevel(dict(zip(tokens, range(len(tokens)))), "[UNK]")) - tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() - tokenizer.post_processor = processors.TemplateProcessing( - single="[CLS] $A [SEP]", special_tokens=[("[CLS]", 2), ("[SEP]", 3)] - ) - if direction is not None: - tokenizer.enable_truncation( - max_length=serialized_context, direction=direction, stride=stride - ) - tokenizer.enable_padding(pad_id=1, pad_token="[PAD]", direction=padding_direction) - tokenizer.save(str(path / "tokenizer.json")) - (path / "tokenizer_config.json").write_text( - json.dumps({"model_max_length": context, "max_length": max_length, "pad_token": "[PAD]"}) - ) - return path - - -@pytest.mark.parametrize("direction", ["left", "right", None]) -@pytest.mark.parametrize("context", [4, 5, 6]) -def test_load_tokenizer_preserves_truncation_direction(tmp_path, direction, context): - """Changing the context limit must not change which end of the input survives.""" - model_dir = make_tokenizer_dir(tmp_path, direction, context) - tokenizer, _ = load_tokenizer(model_dir) - words = ["one", "two", "three", "four", "five"] - kept_words = words[-(context - 2) :] if direction == "left" else words[: context - 2] - - assert tokenizer.truncation["max_length"] == context - assert tokenizer.truncation["direction"] == (direction or "right") - assert tokenizer.encode(" ".join(words)).tokens == ["[CLS]", *kept_words, "[SEP]"] - - -@pytest.mark.parametrize("direction", ["left", "right"]) -def test_text_model_uses_preserved_direction_in_mixed_length_batch(tmp_path, direction): - """The real text-model tokenizer retains the declared text and still pads batches.""" - model_dir = make_tokenizer_dir(tmp_path, direction, context=5) - model = OnnxTextModel() - model._load_tokenizer(model_dir) - - encoded = model.tokenize(["one two three four five", "one"]) - expected_long = ["three", "four", "five"] if direction == "left" else ["one", "two", "three"] - assert encoded[0].tokens == ["[CLS]", *expected_long, "[SEP]"] - assert encoded[1].tokens == ["[CLS]", "one", "[SEP]", "[PAD]", "[PAD]"] - assert encoded[1].attention_mask == [1, 1, 1, 0, 0] - assert np.array([encoding.ids for encoding in encoded]).shape == (2, 5) - - -@pytest.mark.parametrize("context,max_length,expected_length", [(8, 5, 5), (4, 6, 4)]) -def test_saved_direction_keeps_the_stricter_config_limit( - tmp_path, context, max_length, expected_length -): - """Direction preservation does not replace the loader's context-limit resolution.""" - model_dir = make_tokenizer_dir( - tmp_path, "left", context, serialized_context=8, max_length=max_length - ) - tokenizer, _ = load_tokenizer(model_dir) - words = ["one", "two", "three", "four", "five"] - - assert tokenizer.truncation["max_length"] == expected_length - assert tokenizer.encode(" ".join(words)).tokens == [ - "[CLS]", - *words[-(expected_length - 2) :], - "[SEP]", - ] - - -def test_saved_direction_does_not_restore_stride(tmp_path): - """The loader still disables overflow stride when applying a smaller context.""" - model_dir = make_tokenizer_dir(tmp_path, "left", context=4, serialized_context=8, stride=3) - - tokenizer, _ = load_tokenizer(model_dir) - - assert tokenizer.truncation["direction"] == "left" - assert tokenizer.truncation["stride"] == 0 - assert tokenizer.truncation["strategy"] == "longest_first" - assert tokenizer.encode("one two three four five").tokens == ["[CLS]", "four", "five", "[SEP]"] - - -def test_left_truncation_and_left_padding_remain_independent(tmp_path): - """Padding keeps short items aligned without changing the long item's retained end.""" - model_dir = make_tokenizer_dir(tmp_path, "left", context=5, padding_direction="left") - tokenizer, _ = load_tokenizer(model_dir) - - encoded = tokenizer.encode_batch(["one two three four five", "one"]) - - assert encoded[0].tokens == ["[CLS]", "three", "four", "five", "[SEP]"] - assert encoded[1].tokens == ["[PAD]", "[PAD]", "[CLS]", "one", "[SEP]"] - assert encoded[1].attention_mask == [0, 0, 1, 1, 1]