fix: preserve ColBERT tokenizer truncation direction

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,

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]",
+        "<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]",
+        ]