new: add qwen embedding (#678)

NOTICE

@@ -15,6 +15,8 @@ These models are developed by Jina (https://jina.ai/) and are subject to Jina AI
 This distribution includes the following Google models, each with its respective license:
 - vidore/colpali-v1.3
     - License: gemma
+- google/embeddinggemma-300m
+    - License: gemma
 
 Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms
 

fastembed/common/model_description.py

@@ -49,4 +49,5 @@ class SparseModelDescription(BaseModelDescription):
 class PoolingType(str, Enum):
     CLS = "CLS"
     MEAN = "MEAN"
+    LAST_TOKEN = "LAST_TOKEN"
     DISABLED = "DISABLED"

fastembed/common/utils.py

@@ -32,6 +32,12 @@ def mean_pooling(input_array: NumpyArray, attention_mask: NDArray[np.int64]) ->
     return pooled_embeddings
 
 
+def last_token_pooling(input_array: NumpyArray, attention_mask: NDArray[np.int64]) -> NumpyArray:
+    """Take the embedding of the last non-padding token of each sequence."""
+    last_token_indices = np.maximum(attention_mask.sum(axis=1) - 1, 0)
+    return input_array[np.arange(input_array.shape[0]), last_token_indices]
+
+
 def iter_batch(iterable: Iterable[T], size: int) -> Iterable[list[T]]:
     """
     >>> list(iter_batch([1,2,3,4,5], 3))

fastembed/text/builtin_sentence_embedding.py

@@ -16,7 +16,7 @@ supported_builtin_sentence_embedding_models: list[DenseModelDescription] = [
             "Prefixes for queries/documents: `task: search result | query: {content}` for query, "
             "`title: {title | 'none'} | text: {content}` for documents, 2025 year."
         ),
-        license="apache-2.0",
+        license="gemma",
         size_in_GB=1.24,
         sources=ModelSource(
             hf="onnx-community/embeddinggemma-300m-ONNX",

fastembed/text/custom_text_embedding.py

@@ -11,7 +11,7 @@ from fastembed.common.model_description import (
 )
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.types import NumpyArray, Device
-from fastembed.common.utils import normalize, mean_pooling
+from fastembed.common.utils import normalize, mean_pooling, last_token_pooling
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
 
 
@@ -73,12 +73,18 @@ class CustomTextEmbedding(OnnxTextEmbedding):
                 raise ValueError("attention_mask must be provided for mean pooling")
             return mean_pooling(embeddings, attention_mask)
 
+        if self._pooling == PoolingType.LAST_TOKEN:
+            if attention_mask is None:
+                raise ValueError("attention_mask must be provided for last token pooling")
+            return last_token_pooling(embeddings, attention_mask)
+
         if self._pooling == PoolingType.DISABLED:
             return embeddings
 
         raise ValueError(
             f"Unsupported pooling type {self._pooling}. "
-            f"Supported types are: {PoolingType.CLS}, {PoolingType.MEAN}, {PoolingType.DISABLED}."
+            f"Supported types are: {PoolingType.CLS}, {PoolingType.MEAN}, "
+            f"{PoolingType.LAST_TOKEN}, {PoolingType.DISABLED}."
         )
 
     def _normalize(self, embeddings: NumpyArray) -> NumpyArray:

fastembed/text/last_token_normalized_embedding.py

@@ -0,0 +1,94 @@
+from typing import Any, Iterable, Type
+
+import onnxruntime as ort
+
+from fastembed.common.types import NumpyArray
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.utils import last_token_pooling, normalize
+from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+
+supported_last_token_normalized_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="Qwen/Qwen3-Embedding-0.6B",
+        dim=1024,
+        description=(
+            "Text embeddings, Unimodal (text), multilingual, 32768 input tokens truncation, "
+            "Prefixes for queries/documents: `Instruct: {task_description}\\nQuery:{query}` "
+            "for queries, none for documents, 2025 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=2.38,
+        sources=ModelSource(hf="Qdrant/Qwen3-Embedding-0.6B-onnx"),
+        model_file="onnx/model.onnx",
+        additional_files=["onnx/model.onnx.data"],
+    ),
+    DenseModelDescription(
+        model="Qwen/Qwen3-Embedding-0.6B-Q",
+        dim=1024,
+        description=(
+            "Text embeddings, Unimodal (text), multilingual, 32768 input tokens truncation, "
+            "Prefixes for queries/documents: `Instruct: {task_description}\\nQuery:{query}` "
+            "for queries, none for documents, int8 weights, requires onnxruntime>=1.23, "
+            "2025 year."
+        ),
+        license="apache-2.0",
+        size_in_GB=1.12,
+        sources=ModelSource(hf="Qdrant/Qwen3-Embedding-0.6B-onnx"),
+        model_file="onnx/model_quantized.onnx",
+    ),
+]
+
+
+class LastTokenNormalizedEmbedding(OnnxTextEmbedding):
+    """Decoder-based embedding models, which pool the last non-padding token and normalize it"""
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
+        return LastTokenNormalizedEmbeddingWorker
+
+    def load_onnx_model(self) -> None:
+        try:
+            super().load_onnx_model()
+        except Exception as e:
+            # int8 weights are stored as 8-bit MatMulNBits, which onnxruntime only
+            # implements since 1.23; older versions fail with "nbits_ == 4 was false"
+            if "nbits" not in str(e).lower():
+                raise
+            raise RuntimeError(
+                f"Could not load {self.model_name}: its int8 weights require "
+                f"onnxruntime>=1.23, but onnxruntime {ort.__version__} is installed. "
+                f"Either upgrade onnxruntime or use a non-quantized model."
+            ) from e
+
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        """Lists the supported models.
+
+        Returns:
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+        """
+        return supported_last_token_normalized_models
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
+        if output.attention_mask is None:
+            raise ValueError("attention_mask must be provided for last token pooling")
+
+        return normalize(last_token_pooling(output.model_output, output.attention_mask))
+
+
+class LastTokenNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> OnnxTextEmbedding:
+        return LastTokenNormalizedEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/text/text_embedding.py

@@ -9,6 +9,7 @@ from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
 from fastembed.text.pooled_embedding import PooledEmbedding
 from fastembed.text.multitask_embedding import JinaEmbeddingV3
 from fastembed.text.builtin_sentence_embedding import BuiltinSentenceEmbedding
+from fastembed.text.last_token_normalized_embedding import LastTokenNormalizedEmbedding
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
 from fastembed.text.text_embedding_base import TextEmbeddingBase
 from fastembed.common.model_description import DenseModelDescription, ModelSource, PoolingType
@@ -22,6 +23,7 @@ class TextEmbedding(TextEmbeddingBase):
         PooledEmbedding,
         JinaEmbeddingV3,
         BuiltinSentenceEmbedding,
+        LastTokenNormalizedEmbedding,
         CustomTextEmbedding,
     ]
 

tests/test_common.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 from fastembed import (
     TextEmbedding,
     SparseTextEmbedding,
@@ -5,6 +7,7 @@ from fastembed import (
     LateInteractionMultimodalEmbedding,
     LateInteractionTextEmbedding,
 )
+from fastembed.common.utils import last_token_pooling
 
 
 def test_text_list_supported_models():
@@ -28,3 +31,17 @@ def test_text_list_supported_models():
         assert "model_file" in description and description["model_file"]
         assert "sources" in description and description["sources"]
         assert "hf" in description["sources"] or "url" in description["sources"]
+
+
+def test_last_token_pooling():
+    token_embeddings = np.array(
+        [
+            [[1.0, 1.0], [2.0, 2.0], [9.0, 9.0], [9.0, 9.0]],  # 2 real tokens, then padding
+            [[3.0, 3.0], [4.0, 4.0], [5.0, 5.0], [6.0, 6.0]],  # no padding
+        ]
+    )
+    attention_mask = np.array([[1, 1, 0, 0], [1, 1, 1, 1]], dtype=np.int64)
+
+    pooled = last_token_pooling(token_embeddings, attention_mask)
+
+    assert np.allclose(pooled, [[2.0, 2.0], [6.0, 6.0]])

tests/test_custom_models.py

@@ -10,7 +10,7 @@ from fastembed.common.model_description import (
     BaseModelDescription,
 )
 from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.utils import normalize, mean_pooling
+from fastembed.common.utils import normalize, mean_pooling, last_token_pooling
 from fastembed.text.custom_text_embedding import CustomTextEmbedding, PostprocessingConfig
 from fastembed.rerank.cross_encoder.custom_text_cross_encoder import CustomTextCrossEncoder
 from fastembed.rerank.cross_encoder import TextCrossEncoder
@@ -136,6 +136,8 @@ def test_mock_add_custom_models():
         f"{PoolingType.MEAN.lower()}": dummy_token_output,
         f"{PoolingType.CLS.lower()}-normalized": dummy_token_output,
         f"{PoolingType.CLS.lower()}": dummy_token_output,
+        f"{PoolingType.LAST_TOKEN.lower()}-normalized": dummy_token_output,
+        f"{PoolingType.LAST_TOKEN.lower()}": dummy_token_output,
         f"{PoolingType.DISABLED.lower()}-normalized": dummy_pooled_output,
         f"{PoolingType.DISABLED.lower()}": dummy_pooled_output,
     }
@@ -147,12 +149,19 @@ def test_mock_add_custom_models():
         f"{PoolingType.MEAN.lower()}": mean_pooling(dummy_token_embedding, dummy_attention_mask),
         f"{PoolingType.CLS.lower()}-normalized": normalize(dummy_token_embedding[:, 0]),
         f"{PoolingType.CLS.lower()}": dummy_token_embedding[:, 0],
+        f"{PoolingType.LAST_TOKEN.lower()}-normalized": normalize(
+            last_token_pooling(dummy_token_embedding, dummy_attention_mask)
+        ),
+        f"{PoolingType.LAST_TOKEN.lower()}": last_token_pooling(
+            dummy_token_embedding, dummy_attention_mask
+        ),
         f"{PoolingType.DISABLED.lower()}-normalized": normalize(dummy_pooled_embedding),
         f"{PoolingType.DISABLED.lower()}": dummy_pooled_embedding,
     }
 
     for pooling, normalization in itertools.product(
-        (PoolingType.MEAN, PoolingType.CLS, PoolingType.DISABLED), (True, False)
+        (PoolingType.MEAN, PoolingType.CLS, PoolingType.LAST_TOKEN, PoolingType.DISABLED),
+        (True, False),
     ):
         model_name = f"{pooling.name.lower()}{'-normalized' if normalization else ''}"
         TextEmbedding.add_custom_model(

tests/test_text_onnx_embeddings.py

@@ -5,6 +5,8 @@ from contextlib import contextmanager
 import numpy as np
 import pytest
 
+from fastembed.text.last_token_normalized_embedding import LastTokenNormalizedEmbedding
+from fastembed.text.onnx_embedding import OnnxTextEmbedding
 from fastembed.text.text_embedding import TextEmbedding
 from tests.utils import delete_model_cache, should_test_model
 
@@ -71,19 +73,36 @@ CANONICAL_VECTOR_VALUES = {
     "google/embeddinggemma-300m": np.array(
         [-0.08181356, 0.0214127, 0.05120273, -0.03690156, -0.0254504]
     ),
+    "Qwen/Qwen3-Embedding-0.6B": np.array(
+        [-0.01476084, 0.01723184, -0.01195498, -0.07275258, 0.00281229]
+    ),
+    "Qwen/Qwen3-Embedding-0.6B-Q": np.array(
+        [-0.01599521, 0.01676456, -0.01195119, -0.07132675, 0.00346729]
+    ),
 }
 
+QWEN3_INSTRUCT_PREFIX = (
+    "Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery:"
+)
 
 DOC_PREFIXES = {
     "google/embeddinggemma-300m": "title: none | text: ",
 }
 QUERY_PREFIXES = {
     "google/embeddinggemma-300m": "task: search result | query: ",
+    "Qwen/Qwen3-Embedding-0.6B": QWEN3_INSTRUCT_PREFIX,
+    "Qwen/Qwen3-Embedding-0.6B-Q": QWEN3_INSTRUCT_PREFIX,
 }
 CANONICAL_QUERY_VECTOR_VALUES = {
     "google/embeddinggemma-300m": np.array(
         [-0.22990295, 0.03311195, 0.04290345, -0.03558498, -0.01399477]
-    )
+    ),
+    "Qwen/Qwen3-Embedding-0.6B": np.array(
+        [-0.01908712, 0.01635596, -0.00356586, -0.03947155, -0.01387356]
+    ),
+    "Qwen/Qwen3-Embedding-0.6B-Q": np.array(
+        [-0.02221339, 0.01932909, -0.00361797, -0.03888897, -0.01362813]
+    ),
 }
 
 MULTI_TASK_MODELS = ["jinaai/jina-embeddings-v3"]
@@ -181,6 +200,23 @@ def test_query_embedding(model_cache) -> None:
             ), model_desc.model
 
 
+def test_quantized_model_reports_onnxruntime_requirement(monkeypatch) -> None:
+    """Old onnxruntime only implements 4-bit MatMulNBits, the error should say so."""
+    monkeypatch.setattr(
+        OnnxTextEmbedding,
+        "load_onnx_model",
+        lambda self: (_ for _ in ()).throw(RuntimeError("nbits_ == 4 was false")),
+    )
+    model = LastTokenNormalizedEmbedding(
+        "Qwen/Qwen3-Embedding-0.6B-Q",
+        lazy_load=True,
+        specific_model_path="./",  # disable model downloading and loading
+    )
+
+    with pytest.raises(RuntimeError, match="onnxruntime>=1.23"):
+        model.load_onnx_model()
+
+
 @pytest.mark.parametrize("n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5")])
 def test_batch_embedding(model_cache, n_dims: int, model_name: str) -> None:
     with model_cache(model_name) as model: