new: run constella-zero with numpy instead of onnx

fastembed/text/onnx_embedding.py

@@ -176,19 +176,6 @@
         sources=ModelSource(hf="jinaai/jina-clip-v1"),
         model_file="onnx/text_model.onnx",
     ),
-    DenseModelDescription(
-        model="Qdrant/constella-zero",
-        dim=1024,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Query-side lookup table with no transformer, embeds into the document space of "
-            "stella_en_400M_v5, Prefixes for queries/documents: not necessary, 2026 year."
-        ),
-        license="mit",
-        size_in_GB=0.03,
-        sources=ModelSource(hf="Qdrant/constella-zero"),
-        model_file="model.onnx",
-    ),
     DenseModelDescription(
         model="Qdrant/stella-en-400M-v5-doc-onnx",
         dim=1024,

fastembed/text/static_embedding.py

@@ -2,6 +2,7 @@
 import json
 import os
 import warnings
+from dataclasses import dataclass
 from multiprocessing import get_all_start_methods
 from pathlib import Path
 from typing import Any, Iterable, Type
@@ -20,8 +21,20 @@
 # Same cap as the onnx exports of these models and model2vec's default `max_length`.
 MAX_LENGTH = 512
 
+
+@dataclass(frozen=True)
+class StaticModelDescription(DenseModelDescription):
+    """How a static model pools its tokens. The defaults are model2vec's."""
+
+    add_special_tokens: bool = False
+    skip_unknown_tokens: bool = True
+    # Weigh a token seen c times in a text by sqrt(c) instead of c. A `weights` tensor still
+    # scales each token's row, which lets an int8 table carry its per-row scales.
+    sqrt_token_counts: bool = False
+
+
 supported_static_models: list[DenseModelDescription] = [
-    DenseModelDescription(
+    StaticModelDescription(
         model="minishlab/potion-base-8M",
         dim=256,
         description=(
@@ -33,7 +46,7 @@
         sources=ModelSource(hf="minishlab/potion-base-8M"),
         model_file="model.safetensors",
     ),
-    DenseModelDescription(
+    StaticModelDescription(
         model="minishlab/potion-retrieval-32M",
         dim=512,
         description=(
@@ -45,7 +58,7 @@
         sources=ModelSource(hf="minishlab/potion-retrieval-32M"),
         model_file="model.safetensors",
     ),
-    DenseModelDescription(
+    StaticModelDescription(
         model="minishlab/potion-multilingual-128M",
         dim=256,
         description=(
@@ -57,6 +70,21 @@
         sources=ModelSource(hf="minishlab/potion-multilingual-128M"),
         model_file="model.safetensors",
     ),
+    StaticModelDescription(
+        model="Qdrant/constella-zero",
+        dim=1024,
+        description=(
+            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Query encoder"
+            " for stella_en_400M_v5, Prefixes for queries/documents: not necessary, 2026 year."
+        ),
+        license="mit",
+        size_in_GB=0.031,
+        sources=ModelSource(hf="Qdrant/constella-zero"),
+        model_file="model.safetensors",
+        add_special_tokens=True,
+        skip_unknown_tokens=False,
+        sqrt_token_counts=True,
+    ),
 ]
 
 
@@ -133,6 +161,8 @@ def __init__(
         # on the CPU with numpy and ignore them, so only options set to something else than their
         # default are reported.
         ignored = {"threads": threads} if threads is not None else {}
+        if kwargs.get("cuda") in (False, Device.CPU):  # qdrant-client passes cuda=False
+            kwargs.pop("cuda")
         for option, default in (
             ("providers", None),
             ("cuda", Device.AUTO),
@@ -152,7 +182,9 @@ def __init__(
             )
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.lazy_load = lazy_load
-        self.model_description = self._get_model_description(model_name)
+        model_description = self._get_model_description(model_name)
+        assert isinstance(model_description, StaticModelDescription)
+        self.model_description = model_description
         self.cache_dir = str(define_cache_dir(cache_dir))
         self._specific_model_path = specific_model_path
         self.mmap = mmap
@@ -191,7 +223,9 @@ def _token_ids(self, documents: list[str]) -> tuple[NDArray[np.int64], NDArray[n
         """Tokenize a batch into one flat id array plus the number of ids of each document."""
         assert self.tokenizer is not None
         # encode_batch_fast skips the character offsets, which are not needed here.
-        encodings = self.tokenizer.encode_batch_fast(documents, add_special_tokens=False)
+        encodings = self.tokenizer.encode_batch_fast(
+            documents, add_special_tokens=self.model_description.add_special_tokens
+        )
         id_lists = [encoding.ids for encoding in encodings]
         lengths = np.fromiter(map(len, id_lists), dtype=np.int64, count=len(documents))
         ids = np.fromiter(
@@ -205,7 +239,7 @@ def raw_embed(self, documents: list[str]) -> NumpyArray:
         assert self.embeddings is not None
         ids, lengths = self._token_ids(documents)
 
-        if self.unk_token_id is not None:
+        if self.unk_token_id is not None and self.model_description.skip_unknown_tokens:
             known = ids != self.unk_token_id
             if not known.all():
                 doc_index = np.repeat(np.arange(len(documents)), lengths)
@@ -230,7 +264,8 @@ def _sum_rows(self, ids: NDArray[np.int64], lengths: NDArray[np.int64]) -> Numpy
         if not ids.size:
             return np.zeros((len(lengths), self.embeddings.shape[1]), dtype=np.float32)
         rows = ids if self.token_mapping is None else self.token_mapping[ids]
-        if len(lengths) == 1:
+        sqrt_counts = self.model_description.sqrt_token_counts
+        if len(lengths) == 1 and not sqrt_counts:
             # A single text, typically a query: summing its rows skips the fixed cost of building
             # a count matrix, which dominates for short texts.
             vectors = self.embeddings[rows]
@@ -248,12 +283,20 @@ def _sum_rows(self, ids: NDArray[np.int64], lengths: NDArray[np.int64]) -> Numpy
                 ]
             )
         doc_index = np.repeat(np.arange(len(lengths)), lengths)
-        counts = np.bincount(
-            doc_index * len(distinct_rows) + row_index,
-            weights=None if self.weights is None else self.weights[ids],
-            minlength=len(lengths) * len(distinct_rows),
-        ).astype(np.float32)
-        return counts.reshape(len(lengths), len(distinct_rows)) @ self.embeddings[distinct_rows]
+        counts = (
+            np.bincount(
+                doc_index * len(distinct_rows) + row_index,
+                weights=None if self.weights is None or sqrt_counts else self.weights[ids],
+                minlength=len(lengths) * len(distinct_rows),
+            )
+            .astype(np.float32)
+            .reshape(len(lengths), len(distinct_rows))
+        )
+        if sqrt_counts:
+            np.sqrt(counts, out=counts)
+            if self.weights is not None:  # rows are token ids: no sqrt model has a `mapping`
+                counts *= self.weights[distinct_rows]
+        return counts @ self.embeddings[distinct_rows]
 
     def embed(
         self,

tests/test_static_embeddings.py

@@ -226,3 +226,17 @@ def test_mmap(model: TextEmbedding) -> None:
     assert np.array_equal(np.stack(list(mapped.embed(docs, batch_size=16))), expected)
     parallel = np.stack(list(mapped.embed(docs, batch_size=16, parallel=2)))
     assert np.allclose(parallel, expected, atol=1e-6)
+
+
+@pytest.mark.skipif(bool(os.getenv("CI")) and not is_manual_run(), reason="large model")
+def test_constella_zero_pooling() -> None:
+    """Verify model-card pooling: keep special and unknown tokens and weigh counts by sqrt(c)."""
+    model = TextEmbedding("Qdrant/constella-zero")
+    text = "data data data 🤗 pipeline"
+    ids, counts = np.unique(model.model.tokenizer.encode(text).ids, return_counts=True)
+    assert 100 in ids and 101 in ids and 102 in ids  # [UNK], [CLS], [SEP]
+    # int8 rows, dequantized by their per-row scales in `weights`
+    scales = model.model.weights[ids]
+    expected = (np.sqrt(counts) * scales).astype(np.float32) @ model.model.embeddings[ids]
+    expected /= np.linalg.norm(expected)
+    assert np.allclose(next(iter(model.embed(text))), expected, atol=1e-6)