From 70e49097f17a16c8b7c6239c29dba393c3e0fc11 Mon Sep 17 00:00:00 2001 From: j45856021-dev <259477747+j45856021-dev@users.noreply.github.com> Date: Sat, 10 Oct 2026 21:24:35 -0700 Subject: [PATCH 1/2] fix: preserve ColBERT tokenizer truncation direction Keep the resolved direction when reserving marker-token space for document and query tokenizers. Cover the shared ColBERT-family loader with local-tokenizer regressions and right-truncation controls. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- fastembed/late_interaction/colbert.py | 8 +- tests/test_late_interaction_embeddings.py | 100 ++++++++++++++++++++++ 2 files changed, 105 insertions(+), 3 deletions(-) diff --git a/fastembed/late_interaction/colbert.py b/fastembed/late_interaction/colbert.py index ecbf8972..48379173 100644 --- a/fastembed/late_interaction/colbert.py +++ b/fastembed/late_interaction/colbert.py @@ -232,7 +232,8 @@ def _load_tokenizer(self, model_dir: Path) -> None: """Load the document and the query tokenizers, and apply colbert's own configuration. Both tokenizers are always (re)loaded from disk here, so the truncation length is - derived from the on-disk value and never decremented twice. + derived from the on-disk value and never decremented twice. Both preserve the + truncation direction resolved by `load_tokenizer`. """ tokenizer, special_token_to_id = load_tokenizer(model_dir=model_dir) query_tokenizer, _ = load_tokenizer(model_dir=model_dir) @@ -247,8 +248,9 @@ def _load_tokenizer(self, model_dir: Path) -> None: } # ensure not to overflow after adding document-marker max_length = tokenizer.truncation["max_length"] - self.RESERVED_MARKER_TOKENS - tokenizer.enable_truncation(max_length=max_length) - query_tokenizer.enable_truncation(max_length=max_length) + truncation_direction = tokenizer.truncation["direction"] + tokenizer.enable_truncation(max_length=max_length, direction=truncation_direction) + query_tokenizer.enable_truncation(max_length=max_length, direction=truncation_direction) query_tokenizer.enable_padding( pad_token=self.MASK_TOKEN, pad_id=mask_token_id, diff --git a/tests/test_late_interaction_embeddings.py b/tests/test_late_interaction_embeddings.py index 134c8d1d..cee4ccdf 100644 --- a/tests/test_late_interaction_embeddings.py +++ b/tests/test_late_interaction_embeddings.py @@ -1,8 +1,11 @@ +import json import os +import string from contextlib import contextmanager import pytest import numpy as np +from tokenizers import Tokenizer, models, pre_tokenizers, processors from fastembed.late_interaction.late_interaction_text_embedding import ( LateInteractionTextEmbedding, @@ -412,3 +415,100 @@ def test_token_count_without_query_expansion(model_cache, model_name: str) -> No long_text = "token " * 400 assert model.token_count(long_text, include_extension=True) == 300 assert model.token_count(long_text, is_doc=False, include_extension=True) == 300 + + +@pytest.mark.parametrize( + "model_name,reserved_marker_tokens", + [ + ("colbert-ir/colbertv2.0", 1), + ("jinaai/jina-colbert-v2", 1), + ("lightonai/LateOn", 0), + ("lightonai/mLateOn", 0), + ], +) +@pytest.mark.parametrize( + "serialized_direction,truncation_side,expected_direction", + [ + (None, None, "right"), + ("left", None, "left"), + (None, "left", "left"), + ("left", "right", "right"), + ], +) +@pytest.mark.parametrize("is_doc", [True, False]) +def test_colbert_tokenizer_preserves_truncation_direction( + tmp_path, + model_name, + reserved_marker_tokens, + serialized_direction, + truncation_side, + expected_direction, + is_doc, +) -> None: + """Keep the configured end of long inputs when reserving the ColBERT marker.""" + words = [f"word{i}" for i in range(50)] + special_tokens = [ + "[PAD]", + "[unused0]", + "[unused1]", + "[UNK]", + "[CLS]", + "[SEP]", + "[MASK]", + "", + ] + vocab = { + token: index + for index, token in enumerate(special_tokens + words + list(string.punctuation)) + } + tokenizer = Tokenizer(models.WordLevel(vocab, unk_token="[UNK]")) + tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() + tokenizer.add_special_tokens(special_tokens) + tokenizer.post_processor = processors.TemplateProcessing( + single="[CLS] $A [SEP]", + special_tokens=[("[CLS]", vocab["[CLS]"]), ("[SEP]", vocab["[SEP]"])], + ) + if serialized_direction is not None: + tokenizer.enable_truncation(max_length=40, direction=serialized_direction) + tokenizer.enable_padding(pad_id=vocab["[PAD]"], pad_token="[PAD]") + tokenizer.save(str(tmp_path / "tokenizer.json")) + (tmp_path / "tokenizer_config.json").write_text( + json.dumps( + { + "model_max_length": 40, + "pad_token": "[PAD]", + "truncation_side": truncation_side, + } + ), + encoding="utf-8", + ) + model = LateInteractionTextEmbedding( + model_name, + lazy_load=True, + local_files_only=True, + specific_model_path=str(tmp_path), + cache_dir=str(tmp_path / "cache"), + ) + text = " ".join(words) + max_length = 40 - reserved_marker_tokens + assert model.token_count(text, is_doc=is_doc) == max_length + kept_words = ( + words[-(max_length - 2) :] if expected_direction == "left" else words[: max_length - 2] + ) + expected_tokens = ["[CLS]", *kept_words, "[SEP]"] + + for reload_tokenizer in (False, True): + if reload_tokenizer: + model.model._load_tokenizer(tmp_path) + encoded = model.model.tokenize([text], is_doc=is_doc)[0] + assert encoded.tokens == expected_tokens + assert encoded.ids == [vocab[token] for token in expected_tokens] + assert encoded.attention_mask == [1] * max_length + + short = model.model.tokenize(["word0 word1"], is_doc=is_doc)[0] + assert [token for token, mask in zip(short.tokens, short.attention_mask) if mask] == [ + "[CLS]", + "word0", + "word1", + "[SEP]", + ] From f1716c4ff375fe8277d4cae0739414c7c30a746b Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Wed, 28 Oct 2026 15:14:51 +0700 Subject: [PATCH 2/2] tests: shrink ColBERT truncation direction test --- tests/test_late_interaction_embeddings.py | 107 +++------------------- 1 file changed, 13 insertions(+), 94 deletions(-) diff --git a/tests/test_late_interaction_embeddings.py b/tests/test_late_interaction_embeddings.py index cee4ccdf..3bd30fe8 100644 --- a/tests/test_late_interaction_embeddings.py +++ b/tests/test_late_interaction_embeddings.py @@ -1,11 +1,10 @@ import json import os -import string from contextlib import contextmanager import pytest import numpy as np -from tokenizers import Tokenizer, models, pre_tokenizers, processors +from tokenizers import Tokenizer, models from fastembed.late_interaction.late_interaction_text_embedding import ( LateInteractionTextEmbedding, @@ -417,98 +416,18 @@ def test_token_count_without_query_expansion(model_cache, model_name: str) -> No assert model.token_count(long_text, is_doc=False, include_extension=True) == 300 -@pytest.mark.parametrize( - "model_name,reserved_marker_tokens", - [ - ("colbert-ir/colbertv2.0", 1), - ("jinaai/jina-colbert-v2", 1), - ("lightonai/LateOn", 0), - ("lightonai/mLateOn", 0), - ], -) -@pytest.mark.parametrize( - "serialized_direction,truncation_side,expected_direction", - [ - (None, None, "right"), - ("left", None, "left"), - (None, "left", "left"), - ("left", "right", "right"), - ], -) -@pytest.mark.parametrize("is_doc", [True, False]) -def test_colbert_tokenizer_preserves_truncation_direction( - tmp_path, - model_name, - reserved_marker_tokens, - serialized_direction, - truncation_side, - expected_direction, - is_doc, -) -> None: - """Keep the configured end of long inputs when reserving the ColBERT marker.""" - words = [f"word{i}" for i in range(50)] - special_tokens = [ - "[PAD]", - "[unused0]", - "[unused1]", - "[UNK]", - "[CLS]", - "[SEP]", - "[MASK]", - "", - ] - vocab = { - token: index - for index, token in enumerate(special_tokens + words + list(string.punctuation)) - } +def test_colbert_keeps_truncation_direction(tmp_path) -> None: + vocab = {"[PAD]": 0, "[UNK]": 1, "[MASK]": 2} tokenizer = Tokenizer(models.WordLevel(vocab, unk_token="[UNK]")) - tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() - tokenizer.add_special_tokens(special_tokens) - tokenizer.post_processor = processors.TemplateProcessing( - single="[CLS] $A [SEP]", - special_tokens=[("[CLS]", vocab["[CLS]"]), ("[SEP]", vocab["[SEP]"])], - ) - if serialized_direction is not None: - tokenizer.enable_truncation(max_length=40, direction=serialized_direction) - tokenizer.enable_padding(pad_id=vocab["[PAD]"], pad_token="[PAD]") + tokenizer.add_special_tokens(list(vocab)) tokenizer.save(str(tmp_path / "tokenizer.json")) - (tmp_path / "tokenizer_config.json").write_text( - json.dumps( - { - "model_max_length": 40, - "pad_token": "[PAD]", - "truncation_side": truncation_side, - } - ), - encoding="utf-8", - ) + tokenizer_config = {"model_max_length": 8, "pad_token": "[PAD]", "truncation_side": "left"} + (tmp_path / "tokenizer_config.json").write_text(json.dumps(tokenizer_config)) + model = LateInteractionTextEmbedding( - model_name, - lazy_load=True, - local_files_only=True, - specific_model_path=str(tmp_path), - cache_dir=str(tmp_path / "cache"), - ) - text = " ".join(words) - max_length = 40 - reserved_marker_tokens - assert model.token_count(text, is_doc=is_doc) == max_length - kept_words = ( - words[-(max_length - 2) :] if expected_direction == "left" else words[: max_length - 2] - ) - expected_tokens = ["[CLS]", *kept_words, "[SEP]"] - - for reload_tokenizer in (False, True): - if reload_tokenizer: - model.model._load_tokenizer(tmp_path) - encoded = model.model.tokenize([text], is_doc=is_doc)[0] - assert encoded.tokens == expected_tokens - assert encoded.ids == [vocab[token] for token in expected_tokens] - assert encoded.attention_mask == [1] * max_length - - short = model.model.tokenize(["word0 word1"], is_doc=is_doc)[0] - assert [token for token, mask in zip(short.tokens, short.attention_mask) if mask] == [ - "[CLS]", - "word0", - "word1", - "[SEP]", - ] + "colbert-ir/colbertv2.0", lazy_load=True, specific_model_path=str(tmp_path) + ).model + model._load_tokenizer(tmp_path) + + assert model.tokenizer.truncation["direction"] == "left" + assert model.query_tokenizer.truncation["direction"] == "left"