feat: Add lightonai/LateOn model

fastembed/late_interaction/colbert.py

@@ -41,7 +41,11 @@
 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
+    # 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
     MASK_TOKEN = "[MASK]"
 
     def _post_process_onnx_output(
@@ -120,7 +124,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:
@@ -242,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(

fastembed/late_interaction/late_interaction_text_embedding.py

@@ -6,13 +6,19 @@
 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, MLateOn
 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,
+        MLateOn,
+    ]
 
     @classmethod
     def list_supported_models(cls) -> list[dict[str, Any]]:

fastembed/late_interaction/lateon.py

@@ -0,0 +1,128 @@
+from pathlib import Path
+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.late_interaction.colbert import Colbert, ColbertEmbeddingWorker
+
+
+supported_lateon_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="lightonai/LateOn",
+        dim=128,
+        description=(
+            "Text embeddings, Unimodal (text), English, 300 input tokens truncation, 2026 year"
+        ),
+        license="apache-2.0",
+        size_in_GB=0.616,
+        sources=ModelSource(hf="lightonai/LateOn"),
+        model_file="model.onnx",
+    ),
+]
+
+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):
+    QUERY_MARKER_TOKEN_ID = 50368
+    DOCUMENT_MARKER_TOKEN_ID = 50369
+    MASK_TOKEN = "[MASK]"
+    # 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]:
+        return LateOnEmbeddingWorker
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        """Lists the supported LateOn models.
+
+        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 output.attention_mask is None:
+            raise ValueError("attention_mask must be provided for post-processing")
+
+        # inference can run in a worker, leaving `skip_list` unset in this process
+        self._ensure_tokenizer()
+
+        # 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:
+                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
+
+
+class MLateOn(LateOn):
+    QUERY_MARKER_TOKEN_ID = 256000
+    DOCUMENT_MARKER_TOKEN_ID = 256001
+    MASK_TOKEN = "<mask>"
+
+    @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)
+        # trained with an empty skiplist, punctuation is a part of a document
+        self.skip_list = set()
+
+
+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,
+        )
+
+
+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,
+        )

tests/test_late_interaction_embeddings.py

@@ -38,6 +38,24 @@
             [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],
+        ]
+    ),
+    "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 = {
@@ -149,6 +167,24 @@
             [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],
+        ]
+    ),
+    "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",)
@@ -298,6 +334,12 @@ 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
+
+    model_name = "lightonai/mLateOn"
+    assert LateInteractionTextEmbedding.get_embedding_size(model_name) == 128
+
 
 def test_embedding_size():
     is_ci = os.getenv("CI")
@@ -344,3 +386,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
+        )
+
+        # 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, include_extension=True) == 300
+        assert model.token_count(long_text, is_doc=False, include_extension=True) == 300