From 1440fef40c86d5ff852f3f78aff3c176f0379afc Mon Sep 17 00:00:00 2001 From: Oliver Mannion <125105+tekumara@users.noreply.github.com> Date: Mon, 1 Jun 2026 11:45:44 +1000 Subject: [PATCH 1/8] Add LateOn late-interaction model --- .../late_interaction_text_embedding.py | 3 +- fastembed/late_interaction/lateon.py | 128 ++++++++++++++++++ tests/test_late_interaction_embeddings.py | 21 +++ 3 files changed, 151 insertions(+), 1 deletion(-) create mode 100644 fastembed/late_interaction/lateon.py diff --git a/fastembed/late_interaction/late_interaction_text_embedding.py b/fastembed/late_interaction/late_interaction_text_embedding.py index 30c8b70d6..7ccecab35 100644 --- a/fastembed/late_interaction/late_interaction_text_embedding.py +++ b/fastembed/late_interaction/late_interaction_text_embedding.py @@ -6,13 +6,14 @@ from fastembed.common import OnnxProvider from fastembed.late_interaction.colbert import Colbert from fastembed.late_interaction.jina_colbert import JinaColbert +from fastembed.late_interaction.lateon import LateOn from fastembed.late_interaction.late_interaction_embedding_base import ( LateInteractionTextEmbeddingBase, ) class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase): - EMBEDDINGS_REGISTRY: list[Type[LateInteractionTextEmbeddingBase]] = [Colbert, JinaColbert] + EMBEDDINGS_REGISTRY: list[Type[LateInteractionTextEmbeddingBase]] = [Colbert, JinaColbert, LateOn] @classmethod def list_supported_models(cls) -> list[dict[str, Any]]: diff --git a/fastembed/late_interaction/lateon.py b/fastembed/late_interaction/lateon.py new file mode 100644 index 000000000..7c5816bc0 --- /dev/null +++ b/fastembed/late_interaction/lateon.py @@ -0,0 +1,128 @@ +import string +from typing import Any, Iterable, Type + +import numpy as np + +from fastembed.common.model_description import DenseModelDescription, ModelSource +from fastembed.common.onnx_model import OnnxOutputContext +from fastembed.common.types import NumpyArray +from fastembed.common.preprocessor_utils import load_tokenizer +from fastembed.common.utils import iter_batch +from fastembed.late_interaction.colbert import Colbert, ColbertEmbeddingWorker + + +supported_lateon_models: list[DenseModelDescription] = [ + DenseModelDescription( + model="lightonai/LateOn", + dim=128, + description=( + "PyLate/ColBERT late-interaction English model based on ModernBERT, " + "300 document tokens, 32 query tokens, 2025 year" + ), + license="apache-2.0", + size_in_GB=0.616, + sources=ModelSource(hf="lightonai/LateOn"), + model_file="model.onnx", + additional_files=["onnx_config.json"], + ), +] + + +class LateOn(Colbert): + QUERY_MARKER_TOKEN_ID = 50368 + DOCUMENT_MARKER_TOKEN_ID = 50369 + QUERY_LENGTH = 32 + DOCUMENT_LENGTH = 300 + MASK_TOKEN = "[MASK]" + + @classmethod + def _get_worker_class(cls) -> Type[ColbertEmbeddingWorker]: + return LateOnEmbeddingWorker + + @classmethod + def _list_supported_models(cls) -> list[DenseModelDescription]: + """Lists the supported LateOn models.""" + return supported_lateon_models + + def load_onnx_model(self) -> None: + self._load_onnx_model( + model_dir=self._model_dir, + model_file=self.model_description.model_file, + threads=self.threads, + providers=self.providers, + cuda=self.cuda, + device_id=self.device_id, + extra_session_options=self._extra_session_options, + ) + self.query_tokenizer, _ = load_tokenizer(model_dir=self._model_dir) + + assert self.tokenizer is not None + self.mask_token_id = self.special_token_to_id[self.MASK_TOKEN] + self.pad_token_id = self.mask_token_id + self.skip_list = { + self.tokenizer.encode(symbol, add_special_tokens=False).ids[0] + for symbol in string.punctuation + } + # LateOn's PyLate config uses document_length/query_length including the inserted + # [D]/[Q] prefix token. Configure the tokenizer for the pre-prefix lengths. + self.tokenizer.enable_truncation(max_length=self.DOCUMENT_LENGTH - 1) + self.query_tokenizer.enable_truncation(max_length=self.QUERY_LENGTH - 1) + + def _post_process_onnx_output( + self, output: OnnxOutputContext, is_doc: bool = True, **kwargs: Any + ) -> Iterable[NumpyArray]: + if is_doc: + yield from super()._post_process_onnx_output(output, is_doc=is_doc, **kwargs) + return + + if output.attention_mask is None: + raise ValueError("attention_mask must be provided for query post-processing") + + for embedding, attention_mask in zip(output.model_output, output.attention_mask): + # LateOn was exported with do_query_expansion=false, so query embeddings should + # only include non-padding query tokens instead of ColBERT mask-token expansion. + embedding = embedding[attention_mask == 1] + norm = np.linalg.norm(embedding, ord=2, axis=1, keepdims=True) + norm_clamped = np.maximum(norm, 1e-12) + yield embedding / norm_clamped + + def token_count( + self, + texts: str | Iterable[str], + batch_size: int = 1024, + is_doc: bool = True, + include_extension: bool = False, + **kwargs: Any, + ) -> int: + if is_doc: + return super().token_count( + texts, + batch_size=batch_size, + is_doc=is_doc, + include_extension=include_extension, + **kwargs, + ) + + if not hasattr(self, "model") or self.model is None: + self.load_onnx_model() + + token_num = 0 + texts = [texts] if isinstance(texts, str) else texts + assert self.query_tokenizer is not None + for batch in iter_batch(texts, batch_size): + for tokens in self.query_tokenizer.encode_batch(batch): + token_num += sum(tokens.attention_mask) + if include_extension: + token_num += len(batch) # add one [Q] prefix token per query + + return token_num + + +class LateOnEmbeddingWorker(ColbertEmbeddingWorker): + def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> LateOn: + return LateOn( + model_name=model_name, + cache_dir=cache_dir, + threads=1, + **kwargs, + ) diff --git a/tests/test_late_interaction_embeddings.py b/tests/test_late_interaction_embeddings.py index 0e433fce8..e004bc363 100644 --- a/tests/test_late_interaction_embeddings.py +++ b/tests/test_late_interaction_embeddings.py @@ -38,6 +38,15 @@ [0.0766, 0.0452, -0.2343, -0.183, 0.0058], ] ), + "lightonai/LateOn": np.array( + [ + [0.00039, 0.00651, 0.0146, 0.00346, 0.00244], + [-0.0029, 0.00423, 0.00042, 0.02236, 0.00981], + [-0.0287, 0.01159, 0.02401, -0.00312, -0.04338], + [-0.04709, 0.00209, 0.02174, -0.00381, -0.00608], + [-0.02461, -0.02876, 0.03014, -0.0035, -0.00431], + ] + ), } CANONICAL_QUERY_VALUES = { @@ -149,6 +158,15 @@ [0.0204, -0.0856, -0.0386, -0.1232, -0.0332], ] ), + "lightonai/LateOn": np.array( + [ + [0.00202, -0.02634, 0.00685, 0.00993, 0.03093], + [-0.02321, -0.0226, 0.00356, 0.02836, 0.01729], + [-0.01066, 0.00595, 0.02884, 0.00267, -0.10405], + [-0.10359, -0.06927, 0.03218, 0.05037, -0.03338], + [-0.02992, -0.03874, 0.10582, 0.06303, 0.05831], + ] + ), } _MODELS_TO_CACHE = ("answerdotai/answerai-colbert-small-v1",) @@ -298,6 +316,9 @@ def test_get_embedding_size(): model_name = "answerdotai/answerai-ColBERT-small-v1" assert LateInteractionTextEmbedding.get_embedding_size(model_name) == 96 + model_name = "lightonai/LateOn" + assert LateInteractionTextEmbedding.get_embedding_size(model_name) == 128 + def test_embedding_size(): is_ci = os.getenv("CI") From 1a9b62001691327a08588d935bb2c3f55003d70e Mon Sep 17 00:00:00 2001 From: Oliver Mannion <125105+tekumara@users.noreply.github.com> Date: Mon, 1 Jun 2026 12:50:59 +1000 Subject: [PATCH 2/8] ruff format --- .../late_interaction/late_interaction_text_embedding.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/fastembed/late_interaction/late_interaction_text_embedding.py b/fastembed/late_interaction/late_interaction_text_embedding.py index 7ccecab35..c57b044ab 100644 --- a/fastembed/late_interaction/late_interaction_text_embedding.py +++ b/fastembed/late_interaction/late_interaction_text_embedding.py @@ -13,7 +13,11 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase): - EMBEDDINGS_REGISTRY: list[Type[LateInteractionTextEmbeddingBase]] = [Colbert, JinaColbert, LateOn] + EMBEDDINGS_REGISTRY: list[Type[LateInteractionTextEmbeddingBase]] = [ + Colbert, + JinaColbert, + LateOn, + ] @classmethod def list_supported_models(cls) -> list[dict[str, Any]]: From b0daa7c2e659180322b8a5b78c93008c808d5417 Mon Sep 17 00:00:00 2001 From: Oliver Mannion <125105+tekumara@users.noreply.github.com> Date: Mon, 1 Jun 2026 12:52:22 +1000 Subject: [PATCH 3/8] compare against pylate --- tests/test_lateon_pylate_reference.py | 136 ++++++++++++++++++++++++++ 1 file changed, 136 insertions(+) create mode 100644 tests/test_lateon_pylate_reference.py diff --git a/tests/test_lateon_pylate_reference.py b/tests/test_lateon_pylate_reference.py new file mode 100644 index 000000000..26b0aebf6 --- /dev/null +++ b/tests/test_lateon_pylate_reference.py @@ -0,0 +1,136 @@ +from typing import Any, Iterable + +import numpy as np +import pytest + +from fastembed import LateInteractionTextEmbedding +from fastembed.common.types import NumpyArray + + +LATEON_MODEL_NAME = "lightonai/LateOn" +REFERENCE_TEXTS = [ + "Hello World", + "Late interaction models compare token embeddings!", +] +RETRIEVAL_QUERIES = [ + "Which animal purrs?", + "What is the capital of France?", +] +RETRIEVAL_DOCUMENT_IDS = ["python", "cat", "paris"] +RETRIEVAL_DOCUMENTS = [ + "Python is a programming language used for machine learning.", + "Cats are small animals that often purr when they are happy.", + "Paris is the capital and largest city of France.", +] + + +def _as_numpy_arrays(embeddings: Iterable[NumpyArray]) -> list[np.ndarray]: + return [np.asarray(embedding, dtype=np.float32) for embedding in embeddings] + + +def _pylate_model(): + pylate_models = pytest.importorskip( + "pylate.models", reason="PyLate is required for the LateOn reference test" + ) + pytest.importorskip("torch", reason="PyLate reference inference requires PyTorch") + return pylate_models.ColBERT(model_name_or_path=LATEON_MODEL_NAME) + + +def _pylate_reference_embeddings(texts: list[str], is_query: bool) -> list[np.ndarray]: + model = _pylate_model() + embeddings = model.encode( + texts, + batch_size=2, + is_query=is_query, + show_progress_bar=False, + convert_to_numpy=True, + normalize_embeddings=True, + device="cpu", + ) + return _as_numpy_arrays(embeddings) + + +def _rerank( + queries_embeddings: list[np.ndarray], documents_embeddings: list[np.ndarray] +) -> list[list[Any]]: + pylate_rank = pytest.importorskip( + "pylate.rank", reason="PyLate is required for the LateOn retrieval reference test" + ) + return pylate_rank.rerank( + documents_ids=[RETRIEVAL_DOCUMENT_IDS] * len(queries_embeddings), + queries_embeddings=queries_embeddings, + documents_embeddings=[documents_embeddings] * len(queries_embeddings), + device="cpu", + ) + + +def _result_id(result: Any) -> str: + return result["id"] if isinstance(result, dict) else result.id + + +def _result_score(result: Any) -> float: + return result["score"] if isinstance(result, dict) else result.score + + +@pytest.mark.parametrize("is_query", [False, True]) +def test_lateon_matches_pylate_reference(is_query: bool) -> None: + pylate_embeddings = _pylate_reference_embeddings(REFERENCE_TEXTS, is_query=is_query) + + fastembed_model = LateInteractionTextEmbedding(LATEON_MODEL_NAME, threads=1) + fastembed_embeddings = _as_numpy_arrays( + fastembed_model.query_embed(REFERENCE_TEXTS) + if is_query + else fastembed_model.embed(REFERENCE_TEXTS, batch_size=2) + ) + + assert len(fastembed_embeddings) == len(pylate_embeddings) + for fastembed_embedding, pylate_embedding in zip(fastembed_embeddings, pylate_embeddings): + assert fastembed_embedding.shape == pylate_embedding.shape + assert np.allclose(fastembed_embedding, pylate_embedding, rtol=1e-3, atol=1e-4) + + +def test_lateon_retrieval_matches_pylate_reference() -> None: + pylate_model = _pylate_model() + pylate_query_embeddings = _as_numpy_arrays( + pylate_model.encode( + RETRIEVAL_QUERIES, + batch_size=2, + is_query=True, + show_progress_bar=False, + convert_to_numpy=True, + normalize_embeddings=True, + device="cpu", + ) + ) + pylate_document_embeddings = _as_numpy_arrays( + pylate_model.encode( + RETRIEVAL_DOCUMENTS, + batch_size=2, + is_query=False, + show_progress_bar=False, + convert_to_numpy=True, + normalize_embeddings=True, + device="cpu", + ) + ) + + fastembed_model = LateInteractionTextEmbedding(LATEON_MODEL_NAME, threads=1) + fastembed_query_embeddings = _as_numpy_arrays(fastembed_model.query_embed(RETRIEVAL_QUERIES)) + fastembed_document_embeddings = _as_numpy_arrays( + fastembed_model.embed(RETRIEVAL_DOCUMENTS, batch_size=2) + ) + + pylate_results = _rerank(pylate_query_embeddings, pylate_document_embeddings) + fastembed_results = _rerank(fastembed_query_embeddings, fastembed_document_embeddings) + + assert len(fastembed_results) == len(pylate_results) + for fastembed_query_results, pylate_query_results in zip(fastembed_results, pylate_results): + assert [_result_id(result) for result in fastembed_query_results] == [ + _result_id(result) for result in pylate_query_results + ] + assert np.allclose( + [_result_score(result) for result in fastembed_query_results], + [_result_score(result) for result in pylate_query_results], + rtol=1e-3, + atol=1e-3, + ) From 4013088801c12bd59a720a401a619192c33a4d2d Mon Sep 17 00:00:00 2001 From: Oliver Mannion <125105+tekumara@users.noreply.github.com> Date: Mon, 1 Jun 2026 12:52:43 +1000 Subject: [PATCH 4/8] remove pylate comparision --- tests/test_lateon_pylate_reference.py | 136 -------------------------- 1 file changed, 136 deletions(-) delete mode 100644 tests/test_lateon_pylate_reference.py diff --git a/tests/test_lateon_pylate_reference.py b/tests/test_lateon_pylate_reference.py deleted file mode 100644 index 26b0aebf6..000000000 --- a/tests/test_lateon_pylate_reference.py +++ /dev/null @@ -1,136 +0,0 @@ -from typing import Any, Iterable - -import numpy as np -import pytest - -from fastembed import LateInteractionTextEmbedding -from fastembed.common.types import NumpyArray - - -LATEON_MODEL_NAME = "lightonai/LateOn" -REFERENCE_TEXTS = [ - "Hello World", - "Late interaction models compare token embeddings!", -] -RETRIEVAL_QUERIES = [ - "Which animal purrs?", - "What is the capital of France?", -] -RETRIEVAL_DOCUMENT_IDS = ["python", "cat", "paris"] -RETRIEVAL_DOCUMENTS = [ - "Python is a programming language used for machine learning.", - "Cats are small animals that often purr when they are happy.", - "Paris is the capital and largest city of France.", -] - - -def _as_numpy_arrays(embeddings: Iterable[NumpyArray]) -> list[np.ndarray]: - return [np.asarray(embedding, dtype=np.float32) for embedding in embeddings] - - -def _pylate_model(): - pylate_models = pytest.importorskip( - "pylate.models", reason="PyLate is required for the LateOn reference test" - ) - pytest.importorskip("torch", reason="PyLate reference inference requires PyTorch") - return pylate_models.ColBERT(model_name_or_path=LATEON_MODEL_NAME) - - -def _pylate_reference_embeddings(texts: list[str], is_query: bool) -> list[np.ndarray]: - model = _pylate_model() - embeddings = model.encode( - texts, - batch_size=2, - is_query=is_query, - show_progress_bar=False, - convert_to_numpy=True, - normalize_embeddings=True, - device="cpu", - ) - return _as_numpy_arrays(embeddings) - - -def _rerank( - queries_embeddings: list[np.ndarray], documents_embeddings: list[np.ndarray] -) -> list[list[Any]]: - pylate_rank = pytest.importorskip( - "pylate.rank", reason="PyLate is required for the LateOn retrieval reference test" - ) - return pylate_rank.rerank( - documents_ids=[RETRIEVAL_DOCUMENT_IDS] * len(queries_embeddings), - queries_embeddings=queries_embeddings, - documents_embeddings=[documents_embeddings] * len(queries_embeddings), - device="cpu", - ) - - -def _result_id(result: Any) -> str: - return result["id"] if isinstance(result, dict) else result.id - - -def _result_score(result: Any) -> float: - return result["score"] if isinstance(result, dict) else result.score - - -@pytest.mark.parametrize("is_query", [False, True]) -def test_lateon_matches_pylate_reference(is_query: bool) -> None: - pylate_embeddings = _pylate_reference_embeddings(REFERENCE_TEXTS, is_query=is_query) - - fastembed_model = LateInteractionTextEmbedding(LATEON_MODEL_NAME, threads=1) - fastembed_embeddings = _as_numpy_arrays( - fastembed_model.query_embed(REFERENCE_TEXTS) - if is_query - else fastembed_model.embed(REFERENCE_TEXTS, batch_size=2) - ) - - assert len(fastembed_embeddings) == len(pylate_embeddings) - for fastembed_embedding, pylate_embedding in zip(fastembed_embeddings, pylate_embeddings): - assert fastembed_embedding.shape == pylate_embedding.shape - assert np.allclose(fastembed_embedding, pylate_embedding, rtol=1e-3, atol=1e-4) - - -def test_lateon_retrieval_matches_pylate_reference() -> None: - pylate_model = _pylate_model() - pylate_query_embeddings = _as_numpy_arrays( - pylate_model.encode( - RETRIEVAL_QUERIES, - batch_size=2, - is_query=True, - show_progress_bar=False, - convert_to_numpy=True, - normalize_embeddings=True, - device="cpu", - ) - ) - pylate_document_embeddings = _as_numpy_arrays( - pylate_model.encode( - RETRIEVAL_DOCUMENTS, - batch_size=2, - is_query=False, - show_progress_bar=False, - convert_to_numpy=True, - normalize_embeddings=True, - device="cpu", - ) - ) - - fastembed_model = LateInteractionTextEmbedding(LATEON_MODEL_NAME, threads=1) - fastembed_query_embeddings = _as_numpy_arrays(fastembed_model.query_embed(RETRIEVAL_QUERIES)) - fastembed_document_embeddings = _as_numpy_arrays( - fastembed_model.embed(RETRIEVAL_DOCUMENTS, batch_size=2) - ) - - pylate_results = _rerank(pylate_query_embeddings, pylate_document_embeddings) - fastembed_results = _rerank(fastembed_query_embeddings, fastembed_document_embeddings) - - assert len(fastembed_results) == len(pylate_results) - for fastembed_query_results, pylate_query_results in zip(fastembed_results, pylate_results): - assert [_result_id(result) for result in fastembed_query_results] == [ - _result_id(result) for result in pylate_query_results - ] - assert np.allclose( - [_result_score(result) for result in fastembed_query_results], - [_result_score(result) for result in pylate_query_results], - rtol=1e-3, - atol=1e-3, - ) From 4d86139177fc7ec4cc408e0543f0d8ca41623333 Mon Sep 17 00:00:00 2001 From: Oliver Mannion <125105+tekumara@users.noreply.github.com> Date: Mon, 1 Jun 2026 13:04:24 +1000 Subject: [PATCH 5/8] pylate canonical values script --- pylate_canonical_values.txt | 27 +++++ scripts/generate_lateon_canonical_values.py | 124 ++++++++++++++++++++ 2 files changed, 151 insertions(+) create mode 100644 pylate_canonical_values.txt create mode 100755 scripts/generate_lateon_canonical_values.py diff --git a/pylate_canonical_values.txt b/pylate_canonical_values.txt new file mode 100644 index 000000000..e50ddc4dd --- /dev/null +++ b/pylate_canonical_values.txt @@ -0,0 +1,27 @@ +# Generated from the PyLate reference implementation, not FastEmbed. +# +# Reference code: +# from pylate import models +# model = models.ColBERT(model_name_or_path='lightonai/LateOn') +# model.encode([text], is_query=False/True, convert_to_numpy=True, normalize_embeddings=True) +# text = 'Hello World' +# document_shape = (5, 128) +# query_shape = (5, 128) + +# CANONICAL_COLUMN_VALUES entry + "lightonai/LateOn": np.array( + [[ 0.00039, 0.00651, 0.0146 , 0.00346, 0.00244], + [-0.0029 , 0.00423, 0.00042, 0.02236, 0.00981], + [-0.0287 , 0.01159, 0.02401, -0.00312, -0.04338], + [-0.04709, 0.00209, 0.02174, -0.00381, -0.00608], + [-0.02461, -0.02876, 0.03014, -0.0035 , -0.00431]] + ), + +# CANONICAL_QUERY_VALUES entry + "lightonai/LateOn": np.array( + [[ 0.00202, -0.02634, 0.00685, 0.00993, 0.03093], + [-0.02321, -0.0226 , 0.00356, 0.02836, 0.01729], + [-0.01066, 0.00595, 0.02884, 0.00267, -0.10405], + [-0.10359, -0.06927, 0.03217, 0.05037, -0.03338], + [-0.02992, -0.03874, 0.10582, 0.06303, 0.05831]] + ), diff --git a/scripts/generate_lateon_canonical_values.py b/scripts/generate_lateon_canonical_values.py new file mode 100755 index 000000000..2fb137568 --- /dev/null +++ b/scripts/generate_lateon_canonical_values.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +"""Generate LateOn canonical test values from the PyLate reference implementation. + +This script prints the abridged document and query vectors used by +``tests/test_late_interaction_embeddings.py``. It intentionally uses PyLate, +not FastEmbed, so the generated values come from the original reference model. + +Example: + python scripts/generate_lateon_canonical_values.py +""" + +from __future__ import annotations + +import argparse +from typing import Sequence + +import numpy as np + + +DEFAULT_MODEL = "lightonai/LateOn" +DEFAULT_TEXT = "Hello World" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Generate canonical LateOn test vectors with PyLate." + ) + parser.add_argument("--model", default=DEFAULT_MODEL, help="PyLate/HF model name") + parser.add_argument("--text", default=DEFAULT_TEXT, help="Text to encode") + parser.add_argument( + "--rows", + type=int, + default=5, + help="Number of token rows to print from each embedding", + ) + parser.add_argument( + "--dims", + type=int, + default=5, + help="Number of dimensions to print from each token embedding", + ) + parser.add_argument( + "--precision", + type=int, + default=5, + help="Decimal precision for printed values", + ) + parser.add_argument( + "--device", + default="cpu", + help="Device passed to PyLate encode, e.g. cpu or cuda", + ) + return parser.parse_args() + + +def load_pylate_model(model_name: str): + try: + from pylate import models + except ImportError as exc: + raise SystemExit( + "PyLate is required to generate reference values. " + "Install it with `pip install pylate`." + ) from exc + + return models.ColBERT(model_name_or_path=model_name) + + +def encode_reference(model, texts: Sequence[str], *, is_query: bool, device: str) -> np.ndarray: + embeddings = model.encode( + list(texts), + batch_size=1, + is_query=is_query, + show_progress_bar=False, + convert_to_numpy=True, + normalize_embeddings=True, + device=device, + ) + return np.asarray(embeddings[0], dtype=np.float32) + + +def format_dict_entry(model_name: str, values: np.ndarray, *, precision: int) -> str: + array = np.array2string( + values, + precision=precision, + separator=", ", + suppress_small=False, + ) + # Indent nested array lines to match the style in tests/test_late_interaction_embeddings.py. + array = "\n".join(f" {line}" for line in array.splitlines()) + return f' "{model_name}": np.array(\n{array}\n ),' + + +def main() -> None: + args = parse_args() + model = load_pylate_model(args.model) + + document_embedding = encode_reference(model, [args.text], is_query=False, device=args.device) + query_embedding = encode_reference(model, [args.text], is_query=True, device=args.device) + + document_values = document_embedding[: args.rows, : args.dims] + query_values = query_embedding[: args.rows, : args.dims] + + print("# Generated from the PyLate reference implementation, not FastEmbed.") + print("#") + print("# Reference code:") + print("# from pylate import models") + print(f"# model = models.ColBERT(model_name_or_path={args.model!r})") + print( + "# model.encode([text], is_query=False/True, " + "convert_to_numpy=True, normalize_embeddings=True)" + ) + print(f"# text = {args.text!r}") + print(f"# document_shape = {tuple(document_embedding.shape)}") + print(f"# query_shape = {tuple(query_embedding.shape)}") + print() + print("# CANONICAL_COLUMN_VALUES entry") + print(format_dict_entry(args.model, document_values, precision=args.precision)) + print() + print("# CANONICAL_QUERY_VALUES entry") + print(format_dict_entry(args.model, query_values, precision=args.precision)) + + +if __name__ == "__main__": + main() From f57557cf2290e01146b61990125dacbf21c4f1d4 Mon Sep 17 00:00:00 2001 From: Oliver Mannion <125105+tekumara@users.noreply.github.com> Date: Mon, 1 Jun 2026 13:04:42 +1000 Subject: [PATCH 6/8] remove pylate canonical values script --- pylate_canonical_values.txt | 27 ----- scripts/generate_lateon_canonical_values.py | 124 -------------------- 2 files changed, 151 deletions(-) delete mode 100644 pylate_canonical_values.txt delete mode 100755 scripts/generate_lateon_canonical_values.py diff --git a/pylate_canonical_values.txt b/pylate_canonical_values.txt deleted file mode 100644 index e50ddc4dd..000000000 --- a/pylate_canonical_values.txt +++ /dev/null @@ -1,27 +0,0 @@ -# Generated from the PyLate reference implementation, not FastEmbed. -# -# Reference code: -# from pylate import models -# model = models.ColBERT(model_name_or_path='lightonai/LateOn') -# model.encode([text], is_query=False/True, convert_to_numpy=True, normalize_embeddings=True) -# text = 'Hello World' -# document_shape = (5, 128) -# query_shape = (5, 128) - -# CANONICAL_COLUMN_VALUES entry - "lightonai/LateOn": np.array( - [[ 0.00039, 0.00651, 0.0146 , 0.00346, 0.00244], - [-0.0029 , 0.00423, 0.00042, 0.02236, 0.00981], - [-0.0287 , 0.01159, 0.02401, -0.00312, -0.04338], - [-0.04709, 0.00209, 0.02174, -0.00381, -0.00608], - [-0.02461, -0.02876, 0.03014, -0.0035 , -0.00431]] - ), - -# CANONICAL_QUERY_VALUES entry - "lightonai/LateOn": np.array( - [[ 0.00202, -0.02634, 0.00685, 0.00993, 0.03093], - [-0.02321, -0.0226 , 0.00356, 0.02836, 0.01729], - [-0.01066, 0.00595, 0.02884, 0.00267, -0.10405], - [-0.10359, -0.06927, 0.03217, 0.05037, -0.03338], - [-0.02992, -0.03874, 0.10582, 0.06303, 0.05831]] - ), diff --git a/scripts/generate_lateon_canonical_values.py b/scripts/generate_lateon_canonical_values.py deleted file mode 100755 index 2fb137568..000000000 --- a/scripts/generate_lateon_canonical_values.py +++ /dev/null @@ -1,124 +0,0 @@ -#!/usr/bin/env python3 -"""Generate LateOn canonical test values from the PyLate reference implementation. - -This script prints the abridged document and query vectors used by -``tests/test_late_interaction_embeddings.py``. It intentionally uses PyLate, -not FastEmbed, so the generated values come from the original reference model. - -Example: - python scripts/generate_lateon_canonical_values.py -""" - -from __future__ import annotations - -import argparse -from typing import Sequence - -import numpy as np - - -DEFAULT_MODEL = "lightonai/LateOn" -DEFAULT_TEXT = "Hello World" - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Generate canonical LateOn test vectors with PyLate." - ) - parser.add_argument("--model", default=DEFAULT_MODEL, help="PyLate/HF model name") - parser.add_argument("--text", default=DEFAULT_TEXT, help="Text to encode") - parser.add_argument( - "--rows", - type=int, - default=5, - help="Number of token rows to print from each embedding", - ) - parser.add_argument( - "--dims", - type=int, - default=5, - help="Number of dimensions to print from each token embedding", - ) - parser.add_argument( - "--precision", - type=int, - default=5, - help="Decimal precision for printed values", - ) - parser.add_argument( - "--device", - default="cpu", - help="Device passed to PyLate encode, e.g. cpu or cuda", - ) - return parser.parse_args() - - -def load_pylate_model(model_name: str): - try: - from pylate import models - except ImportError as exc: - raise SystemExit( - "PyLate is required to generate reference values. " - "Install it with `pip install pylate`." - ) from exc - - return models.ColBERT(model_name_or_path=model_name) - - -def encode_reference(model, texts: Sequence[str], *, is_query: bool, device: str) -> np.ndarray: - embeddings = model.encode( - list(texts), - batch_size=1, - is_query=is_query, - show_progress_bar=False, - convert_to_numpy=True, - normalize_embeddings=True, - device=device, - ) - return np.asarray(embeddings[0], dtype=np.float32) - - -def format_dict_entry(model_name: str, values: np.ndarray, *, precision: int) -> str: - array = np.array2string( - values, - precision=precision, - separator=", ", - suppress_small=False, - ) - # Indent nested array lines to match the style in tests/test_late_interaction_embeddings.py. - array = "\n".join(f" {line}" for line in array.splitlines()) - return f' "{model_name}": np.array(\n{array}\n ),' - - -def main() -> None: - args = parse_args() - model = load_pylate_model(args.model) - - document_embedding = encode_reference(model, [args.text], is_query=False, device=args.device) - query_embedding = encode_reference(model, [args.text], is_query=True, device=args.device) - - document_values = document_embedding[: args.rows, : args.dims] - query_values = query_embedding[: args.rows, : args.dims] - - print("# Generated from the PyLate reference implementation, not FastEmbed.") - print("#") - print("# Reference code:") - print("# from pylate import models") - print(f"# model = models.ColBERT(model_name_or_path={args.model!r})") - print( - "# model.encode([text], is_query=False/True, " - "convert_to_numpy=True, normalize_embeddings=True)" - ) - print(f"# text = {args.text!r}") - print(f"# document_shape = {tuple(document_embedding.shape)}") - print(f"# query_shape = {tuple(query_embedding.shape)}") - print() - print("# CANONICAL_COLUMN_VALUES entry") - print(format_dict_entry(args.model, document_values, precision=args.precision)) - print() - print("# CANONICAL_QUERY_VALUES entry") - print(format_dict_entry(args.model, query_values, precision=args.precision)) - - -if __name__ == "__main__": - main() From 2b65417b48b9f003bf3001cdf1309e8292e4e38b Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Sat, 3 Oct 2026 01:11:31 +0700 Subject: [PATCH 7/8] new: add mLateOn and rework LateOn for the current colbert base --- fastembed/late_interaction/colbert.py | 6 +- .../late_interaction_text_embedding.py | 3 +- fastembed/late_interaction/lateon.py | 170 ++++++++++-------- tests/test_late_interaction_embeddings.py | 49 +++++ 4 files changed, 152 insertions(+), 76 deletions(-) diff --git a/fastembed/late_interaction/colbert.py b/fastembed/late_interaction/colbert.py index 9115115bd..d855b0bad 100644 --- a/fastembed/late_interaction/colbert.py +++ b/fastembed/late_interaction/colbert.py @@ -41,7 +41,9 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]): QUERY_MARKER_TOKEN_ID = 1 DOCUMENT_MARKER_TOKEN_ID = 2 - MIN_QUERY_LENGTH = 31 # it's 32, we add one additional special token in the beginning + # it's 32, we add one additional special token in the beginning. `None` turns the query + # expansion off: the query is then padded to the longest one in its batch, as a document is + MIN_QUERY_LENGTH: int | None = 31 MASK_TOKEN = "[MASK]" def _post_process_onnx_output( @@ -120,7 +122,7 @@ def token_count( token_num += sum(tokens.attention_mask) else: attend_count = sum(tokens.attention_mask) - if include_extension: + if include_extension and self.MIN_QUERY_LENGTH is not None: token_num += max(attend_count, self.MIN_QUERY_LENGTH) else: diff --git a/fastembed/late_interaction/late_interaction_text_embedding.py b/fastembed/late_interaction/late_interaction_text_embedding.py index c57b044ab..bd23c6b6f 100644 --- a/fastembed/late_interaction/late_interaction_text_embedding.py +++ b/fastembed/late_interaction/late_interaction_text_embedding.py @@ -6,7 +6,7 @@ from fastembed.common import OnnxProvider from fastembed.late_interaction.colbert import Colbert from fastembed.late_interaction.jina_colbert import JinaColbert -from fastembed.late_interaction.lateon import LateOn +from fastembed.late_interaction.lateon import LateOn, MLateOn from fastembed.late_interaction.late_interaction_embedding_base import ( LateInteractionTextEmbeddingBase, ) @@ -17,6 +17,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase): Colbert, JinaColbert, LateOn, + MLateOn, ] @classmethod diff --git a/fastembed/late_interaction/lateon.py b/fastembed/late_interaction/lateon.py index 7c5816bc0..23c191595 100644 --- a/fastembed/late_interaction/lateon.py +++ b/fastembed/late_interaction/lateon.py @@ -1,4 +1,4 @@ -import string +from pathlib import Path from typing import Any, Iterable, Type import numpy as np @@ -6,8 +6,6 @@ from fastembed.common.model_description import DenseModelDescription, ModelSource from fastembed.common.onnx_model import OnnxOutputContext from fastembed.common.types import NumpyArray -from fastembed.common.preprocessor_utils import load_tokenizer -from fastembed.common.utils import iter_batch from fastembed.late_interaction.colbert import Colbert, ColbertEmbeddingWorker @@ -16,24 +14,49 @@ model="lightonai/LateOn", dim=128, description=( - "PyLate/ColBERT late-interaction English model based on ModernBERT, " - "300 document tokens, 32 query tokens, 2025 year" + "Text embeddings, Unimodal (text), English, 299 input tokens truncation, 2026 year" ), license="apache-2.0", size_in_GB=0.616, sources=ModelSource(hf="lightonai/LateOn"), model_file="model.onnx", - additional_files=["onnx_config.json"], + ), +] + +supported_mlateon_models: list[DenseModelDescription] = [ + DenseModelDescription( + model="lightonai/mLateOn", + dim=128, + description=( + "Text embeddings, Unimodal (text), Multilingual and code, " + "8192 input tokens truncation, 2026 year" + ), + license="apache-2.0", + size_in_GB=1.25, + sources=ModelSource(hf="lightonai/mLateOn"), + model_file="model.onnx", ), ] class LateOn(Colbert): + """LightOn's English ColBERT model, trained and exported with PyLate. + + It differs from colbert in two ways: a query is not expanded up to a fixed length with + mask tokens, and the padding reuses the mask token id. + + TODO: PyLate writes `model_max_length` as its `document_length` minus the [D] marker, 299 + here, and colbert subtracts the marker a second time. Until the model repository reports + the real limits, a document is cut at 299 tokens instead of 300, and a query is not cut at + PyLate's `query_length` of 32 at all, which is how every other colbert model behaves here. + """ + QUERY_MARKER_TOKEN_ID = 50368 DOCUMENT_MARKER_TOKEN_ID = 50369 - QUERY_LENGTH = 32 - DOCUMENT_LENGTH = 300 MASK_TOKEN = "[MASK]" + # exported with `do_query_expansion=false`, so a query is padded to the longest one in its + # batch instead of to a fixed length, and that padding is dropped from the output + MIN_QUERY_LENGTH = None @classmethod def _get_worker_class(cls) -> Type[ColbertEmbeddingWorker]: @@ -41,81 +64,72 @@ def _get_worker_class(cls) -> Type[ColbertEmbeddingWorker]: @classmethod def _list_supported_models(cls) -> list[DenseModelDescription]: - """Lists the supported LateOn models.""" - return supported_lateon_models + """Lists the supported LateOn models. - def load_onnx_model(self) -> None: - self._load_onnx_model( - model_dir=self._model_dir, - model_file=self.model_description.model_file, - threads=self.threads, - providers=self.providers, - cuda=self.cuda, - device_id=self.device_id, - extra_session_options=self._extra_session_options, - ) - self.query_tokenizer, _ = load_tokenizer(model_dir=self._model_dir) - - assert self.tokenizer is not None - self.mask_token_id = self.special_token_to_id[self.MASK_TOKEN] - self.pad_token_id = self.mask_token_id - self.skip_list = { - self.tokenizer.encode(symbol, add_special_tokens=False).ids[0] - for symbol in string.punctuation - } - # LateOn's PyLate config uses document_length/query_length including the inserted - # [D]/[Q] prefix token. Configure the tokenizer for the pre-prefix lengths. - self.tokenizer.enable_truncation(max_length=self.DOCUMENT_LENGTH - 1) - self.query_tokenizer.enable_truncation(max_length=self.QUERY_LENGTH - 1) + Returns: + list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information. + """ + return supported_lateon_models def _post_process_onnx_output( self, output: OnnxOutputContext, is_doc: bool = True, **kwargs: Any ) -> Iterable[NumpyArray]: - if is_doc: - yield from super()._post_process_onnx_output(output, is_doc=is_doc, **kwargs) - return - if output.attention_mask is None: - raise ValueError("attention_mask must be provided for query post-processing") - - for embedding, attention_mask in zip(output.model_output, output.attention_mask): - # LateOn was exported with do_query_expansion=false, so query embeddings should - # only include non-padding query tokens instead of ColBERT mask-token expansion. - embedding = embedding[attention_mask == 1] + raise ValueError("attention_mask must be provided for post-processing") + + # with `lazy_load` and `parallel`, inference runs in the workers and the parent + # never calls `load_onnx_model`, so `skip_list` might not be set yet + self._ensure_tokenizer() + + # padding reuses the mask token id, so it can only be told apart from a mask token in + # the input by the attention mask, which is also what drops the padding from a query + keep = output.attention_mask == 1 + if is_doc and self.skip_list: + if output.input_ids is None: + raise ValueError("input_ids must be provided for document post-processing") + keep &= ~np.isin(output.input_ids, list(self.skip_list)) + + for embedding, keep_tokens in zip(output.model_output, keep): + embedding = embedding[keep_tokens] norm = np.linalg.norm(embedding, ord=2, axis=1, keepdims=True) norm_clamped = np.maximum(norm, 1e-12) yield embedding / norm_clamped - def token_count( - self, - texts: str | Iterable[str], - batch_size: int = 1024, - is_doc: bool = True, - include_extension: bool = False, - **kwargs: Any, - ) -> int: - if is_doc: - return super().token_count( - texts, - batch_size=batch_size, - is_doc=is_doc, - include_extension=include_extension, - **kwargs, - ) - - if not hasattr(self, "model") or self.model is None: - self.load_onnx_model() - - token_num = 0 - texts = [texts] if isinstance(texts, str) else texts - assert self.query_tokenizer is not None - for batch in iter_batch(texts, batch_size): - for tokens in self.query_tokenizer.encode_batch(batch): - token_num += sum(tokens.attention_mask) - if include_extension: - token_num += len(batch) # add one [Q] prefix token per query - - return token_num + +class MLateOn(LateOn): + """The multilingual and code sibling of LateOn, based on mmBERT. + + Besides its own tokenizer and context length, it was trained with an empty skiplist, so + punctuation is a part of a document rather than dropped from it. + """ + + QUERY_MARKER_TOKEN_ID = 256000 + DOCUMENT_MARKER_TOKEN_ID = 256001 + MASK_TOKEN = "" + + @classmethod + def _get_worker_class(cls) -> Type[ColbertEmbeddingWorker]: + return MLateOnEmbeddingWorker + + @classmethod + def _list_supported_models(cls) -> list[DenseModelDescription]: + """Lists the supported mLateOn models. + + Returns: + list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information. + """ + return supported_mlateon_models + + def _load_tokenizer(self, model_dir: Path) -> None: + super()._load_tokenizer(model_dir) + self.skip_list = set() + + # TODO: drop once the model repository ships its own tokenizer metadata. It carries + # LateOn's 299, and this model reads 8192 tokens, so without this every input is cut + # at 299. The marker colbert inserts takes one of the 8192 positions. + assert self.tokenizer is not None and self.query_tokenizer is not None + self.tokenizer.enable_truncation(max_length=8192 - 1) + self.query_tokenizer.enable_truncation(max_length=8192 - 1) class LateOnEmbeddingWorker(ColbertEmbeddingWorker): @@ -126,3 +140,13 @@ def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> Late threads=1, **kwargs, ) + + +class MLateOnEmbeddingWorker(ColbertEmbeddingWorker): + def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> MLateOn: + return MLateOn( + model_name=model_name, + cache_dir=cache_dir, + threads=1, + **kwargs, + ) diff --git a/tests/test_late_interaction_embeddings.py b/tests/test_late_interaction_embeddings.py index e004bc363..8db3b11ee 100644 --- a/tests/test_late_interaction_embeddings.py +++ b/tests/test_late_interaction_embeddings.py @@ -47,6 +47,17 @@ [-0.02461, -0.02876, 0.03014, -0.0035, -0.00431], ] ), + # computed with PyLate from the safetensors checkpoint. lightonai/mLateOn's published + # model.onnx was exported from different weights and does not reproduce them yet + "lightonai/mLateOn": np.array( + [ + [0.06085, -0.0992, -0.00024, -0.01035, -0.04497], + [0.0511, -0.09208, -0.00751, -0.0224, -0.03931], + [0.04596, -0.10963, -0.01414, 0.00134, -0.03307], + [0.04688, -0.09261, -0.00664, -0.01437, -0.0551], + [0.06737, -0.12068, -0.00429, -0.00843, -0.04418], + ] + ), } CANONICAL_QUERY_VALUES = { @@ -167,6 +178,15 @@ [-0.02992, -0.03874, 0.10582, 0.06303, 0.05831], ] ), + "lightonai/mLateOn": np.array( + [ + [0.08036, -0.09783, 0.00963, -0.00164, -0.04674], + [0.10542, -0.10502, 0.06527, -0.01749, -0.04493], + [0.09968, -0.1013, 0.02409, -0.04671, -0.02976], + [0.09301, -0.09293, 0.03714, -0.07039, -0.06817], + [0.10895, -0.10447, 0.01298, -0.0075, -0.05043], + ] + ), } _MODELS_TO_CACHE = ("answerdotai/answerai-colbert-small-v1",) @@ -319,6 +339,9 @@ def test_get_embedding_size(): model_name = "lightonai/LateOn" assert LateInteractionTextEmbedding.get_embedding_size(model_name) == 128 + model_name = "lightonai/mLateOn" + assert LateInteractionTextEmbedding.get_embedding_size(model_name) == 128 + def test_embedding_size(): is_ci = os.getenv("CI") @@ -365,3 +388,29 @@ def test_token_count(model_cache, model_name) -> None: assert model.token_count(documents, is_doc=False, include_extension=True) == 64 very_long_query = "It's a very long query which definitely contains more than 32 tokens and we're using it to check whether the method can handle large query properly without cutting it to 32 tokens" assert model.token_count(very_long_query, is_doc=False, include_extension=True) > 32 + + +@pytest.mark.parametrize("model_name", ["lightonai/LateOn"]) +def test_token_count_without_query_expansion(model_cache, model_name: str) -> None: + is_ci = os.getenv("CI") + is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch" + if is_ci and not is_manual: + pytest.skip("Skipping LateOn in CI non-manual mode") + + with model_cache(model_name) as model: + queries = ["short query", "a longer query, which is still shorter than 32 tokens"] + query_token_count = model.token_count(queries, is_doc=False) + assert query_token_count == sum( + model.token_count(query, is_doc=False) for query in queries + ) + # a query is not expanded up to 32 tokens, the extension is one [Q] marker per query + assert ( + model.token_count(queries, is_doc=False, include_extension=True) + == query_token_count + 2 + ) + + # TODO: PyLate cuts a query at 32 tokens and a document at 300, the marker included. + # The tokenizer metadata reports neither of them, see the note on the LateOn class. + long_text = "token " * 400 + assert model.token_count(long_text, is_doc=False, include_extension=True) == 299 + assert model.token_count(long_text, include_extension=True) == 299 From bb274c753c79ffdf7a179ad918d56aac94c5e469 Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Mon, 5 Oct 2026 22:49:39 +0700 Subject: [PATCH 8/8] fix: stop truncating LateOn documents one token early --- fastembed/late_interaction/colbert.py | 4 ++- fastembed/late_interaction/lateon.py | 38 +++++------------------ tests/test_late_interaction_embeddings.py | 10 +++--- 3 files changed, 14 insertions(+), 38 deletions(-) diff --git a/fastembed/late_interaction/colbert.py b/fastembed/late_interaction/colbert.py index d855b0bad..ecbf8972e 100644 --- a/fastembed/late_interaction/colbert.py +++ b/fastembed/late_interaction/colbert.py @@ -41,6 +41,8 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]): QUERY_MARKER_TOKEN_ID = 1 DOCUMENT_MARKER_TOKEN_ID = 2 + # tokens `_preprocess_onnx_input` inserts that the tokenizer metadata does not count + RESERVED_MARKER_TOKENS = 1 # it's 32, we add one additional special token in the beginning. `None` turns the query # expansion off: the query is then padded to the longest one in its batch, as a document is MIN_QUERY_LENGTH: int | None = 31 @@ -244,7 +246,7 @@ def _load_tokenizer(self, model_dir: Path) -> None: for symbol in string.punctuation } # ensure not to overflow after adding document-marker - max_length = tokenizer.truncation["max_length"] - 1 + 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) query_tokenizer.enable_padding( diff --git a/fastembed/late_interaction/lateon.py b/fastembed/late_interaction/lateon.py index 23c191595..833aac87e 100644 --- a/fastembed/late_interaction/lateon.py +++ b/fastembed/late_interaction/lateon.py @@ -14,7 +14,7 @@ model="lightonai/LateOn", dim=128, description=( - "Text embeddings, Unimodal (text), English, 299 input tokens truncation, 2026 year" + "Text embeddings, Unimodal (text), English, 300 input tokens truncation, 2026 year" ), license="apache-2.0", size_in_GB=0.616, @@ -40,23 +40,13 @@ class LateOn(Colbert): - """LightOn's English ColBERT model, trained and exported with PyLate. - - It differs from colbert in two ways: a query is not expanded up to a fixed length with - mask tokens, and the padding reuses the mask token id. - - TODO: PyLate writes `model_max_length` as its `document_length` minus the [D] marker, 299 - here, and colbert subtracts the marker a second time. Until the model repository reports - the real limits, a document is cut at 299 tokens instead of 300, and a query is not cut at - PyLate's `query_length` of 32 at all, which is how every other colbert model behaves here. - """ - QUERY_MARKER_TOKEN_ID = 50368 DOCUMENT_MARKER_TOKEN_ID = 50369 MASK_TOKEN = "[MASK]" - # exported with `do_query_expansion=false`, so a query is padded to the longest one in its - # batch instead of to a fixed length, and that padding is dropped from the output + # exported with `do_query_expansion=false` MIN_QUERY_LENGTH = None + # PyLate writes `model_max_length` with the [Q]/[D] marker already taken off + RESERVED_MARKER_TOKENS = 0 @classmethod def _get_worker_class(cls) -> Type[ColbertEmbeddingWorker]: @@ -77,12 +67,10 @@ def _post_process_onnx_output( if output.attention_mask is None: raise ValueError("attention_mask must be provided for post-processing") - # with `lazy_load` and `parallel`, inference runs in the workers and the parent - # never calls `load_onnx_model`, so `skip_list` might not be set yet + # inference can run in a worker, leaving `skip_list` unset in this process self._ensure_tokenizer() - # padding reuses the mask token id, so it can only be told apart from a mask token in - # the input by the attention mask, which is also what drops the padding from a query + # padding reuses the mask token id, only the attention mask tells them apart keep = output.attention_mask == 1 if is_doc and self.skip_list: if output.input_ids is None: @@ -97,12 +85,6 @@ def _post_process_onnx_output( class MLateOn(LateOn): - """The multilingual and code sibling of LateOn, based on mmBERT. - - Besides its own tokenizer and context length, it was trained with an empty skiplist, so - punctuation is a part of a document rather than dropped from it. - """ - QUERY_MARKER_TOKEN_ID = 256000 DOCUMENT_MARKER_TOKEN_ID = 256001 MASK_TOKEN = "" @@ -122,15 +104,9 @@ def _list_supported_models(cls) -> list[DenseModelDescription]: def _load_tokenizer(self, model_dir: Path) -> None: super()._load_tokenizer(model_dir) + # trained with an empty skiplist, punctuation is a part of a document self.skip_list = set() - # TODO: drop once the model repository ships its own tokenizer metadata. It carries - # LateOn's 299, and this model reads 8192 tokens, so without this every input is cut - # at 299. The marker colbert inserts takes one of the 8192 positions. - assert self.tokenizer is not None and self.query_tokenizer is not None - self.tokenizer.enable_truncation(max_length=8192 - 1) - self.query_tokenizer.enable_truncation(max_length=8192 - 1) - class LateOnEmbeddingWorker(ColbertEmbeddingWorker): def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> LateOn: diff --git a/tests/test_late_interaction_embeddings.py b/tests/test_late_interaction_embeddings.py index 8db3b11ee..224667e2f 100644 --- a/tests/test_late_interaction_embeddings.py +++ b/tests/test_late_interaction_embeddings.py @@ -47,8 +47,6 @@ [-0.02461, -0.02876, 0.03014, -0.0035, -0.00431], ] ), - # computed with PyLate from the safetensors checkpoint. lightonai/mLateOn's published - # model.onnx was exported from different weights and does not reproduce them yet "lightonai/mLateOn": np.array( [ [0.06085, -0.0992, -0.00024, -0.01035, -0.04497], @@ -409,8 +407,8 @@ def test_token_count_without_query_expansion(model_cache, model_name: str) -> No == query_token_count + 2 ) - # TODO: PyLate cuts a query at 32 tokens and a document at 300, the marker included. - # The tokenizer metadata reports neither of them, see the note on the LateOn class. + # a document is cut where PyLate cuts it, at 300 tokens with the marker included, and + # a query keeps that limit rather than PyLate's 32 long_text = "token " * 400 - assert model.token_count(long_text, is_doc=False, include_extension=True) == 299 - assert model.token_count(long_text, include_extension=True) == 299 + assert model.token_count(long_text, include_extension=True) == 300 + assert model.token_count(long_text, is_doc=False, include_extension=True) == 300