DO NOT MERGE: run static embeddings with tokenizers v1

.github/workflows/python-tests.yml

@@ -65,4 +65,5 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
-          poetry run pytest --durations=10
+          # Only the static embeddings support tokenizers 1.0, the other models' tests are not run.
+          poetry run pytest tests/test_static_embeddings.py --durations=10

README.md

@@ -1,3 +1,22 @@
+# Experimental branch: static embeddings with tokenizers 1.0, not meant to be merged
+
+The minishlab potion models (`minishlab/potion-base-8M`, `minishlab/potion-retrieval-32M` and `minishlab/potion-multilingual-128M`) are static: an embedding is the mean of token vectors, computed with numpy. They run on [tokenizers](https://github.com/huggingface/tokenizers) 1.0, which makes them several times faster than with tokenizers 0.x.
+
+```python
+from fastembed import TextEmbedding
+
+model = TextEmbedding(model_name="minishlab/potion-base-8M")
+embeddings = list(model.embed(documents))
+```
+
+The tokenizers 1.0 wheels on PyPI are built without a regex backend, so tokenizers with a regex pattern, such as the one of `minishlab/potion-multilingual-128M`, fail to load with them. To use that model, build tokenizers from source with the regex feature (this needs a [Rust toolchain](https://rustup.rs)):
+
+```bash
+MATURIN_PEP517_ARGS="--features tk-encode/fancy-regex,tk-serialize/fancy-regex" pip install --no-binary tokenizers --pre "tokenizers==1.0.0rc2"
+```
+
+Models other than the static ones still use the tokenizers 0.x API and are not supported with tokenizers 1.0.
+
 # ⚡️ What is FastEmbed?
 
 FastEmbed is a lightweight, fast, Python library built for embedding generation. We [support popular text models](https://qdrant.github.io/fastembed/examples/Supported_Models/). Please [open a GitHub issue](https://github.com/qdrant/fastembed/issues/new) if you want us to add a new model.

fastembed/common/preprocessor_utils.py

@@ -3,7 +3,7 @@
 from typing import Any, Iterator
 from pathlib import Path
 
-from tokenizers import AddedToken, Tokenizer
+from tokenizers import Tokenizer
 
 from fastembed.image.transform.operators import Compose
 
@@ -109,6 +109,10 @@ def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
         if isinstance(token, str):
             tokenizer.add_special_tokens([token])
         elif isinstance(token, dict):
+            # Imported here: tokenizers 1.0 has no AddedToken, and the static models, which work
+            # with 1.0, must not fail on importing fastembed.
+            from tokenizers import AddedToken
+
             tokenizer.add_special_tokens([AddedToken(**token)])
 
     # Padding is always normalized to batch-longest. A serialized fixed length shorter than the

fastembed/text/static_embedding.py

@@ -1,4 +1,3 @@
-import itertools
 import json
 import os
 import warnings
@@ -8,7 +7,7 @@
 
 import numpy as np
 from numpy.typing import NDArray
-from tokenizers import Tokenizer
+from tokenizers import Tokenizer, Truncation
 
 from fastembed.common.model_description import DenseModelDescription, ModelSource
 from fastembed.common.safetensors_utils import load_safetensors
@@ -71,8 +70,8 @@ def load_static_tokenizer(model_dir: Path) -> tuple[Tokenizer, int | None]:
         raise ValueError(f"Could not find tokenizer.json in {model_dir}")
 
     tokenizer = Tokenizer.from_file(str(tokenizer_path))
-    tokenizer.enable_truncation(max_length=MAX_LENGTH)
-    tokenizer.no_padding()
+    tokenizer.truncation = Truncation(MAX_LENGTH)
+    tokenizer.padding = None
 
     with open(tokenizer_path, encoding="utf-8") as f:
         tokenizer_model = json.load(f)["model"]
@@ -179,16 +178,16 @@ def _ensure_model(self) -> None:
         if self.embeddings is None or self.tokenizer is None:
             self.load_model()
 
-    def _token_ids(self, documents: list[str]) -> tuple[NDArray[np.int64], NDArray[np.int64]]:
+    def _token_ids(
+        self, documents: list[str]
+    ) -> tuple[NDArray[np.integer[Any]], NDArray[np.int64]]:
         """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)
-        id_lists = [encoding.ids for encoding in encodings]
-        lengths = np.fromiter(map(len, id_lists), dtype=np.int64, count=len(documents))
-        ids = np.fromiter(
-            itertools.chain.from_iterable(id_lists), dtype=np.int64, count=int(lengths.sum())
-        )
+        # The ids come as numpy views over the encodings, so no Python lists are built.
+        encodings = self.tokenizer.encode_batch(documents, add_special_tokens=False)
+        ids_per_doc = [encoding.ids_array for encoding in encodings]
+        lengths = np.fromiter(map(len, ids_per_doc), dtype=np.int64, count=len(documents))
+        ids = np.concatenate(ids_per_doc) if ids_per_doc else np.empty(0, dtype=np.uint32)
         return ids, lengths
 
     def raw_embed(self, documents: list[str]) -> NumpyArray:
@@ -209,7 +208,7 @@ def raw_embed(self, documents: list[str]) -> NumpyArray:
         # direction. Empty documents stay zero vectors.
         return normalize(sums)
 
-    def _sum_rows(self, ids: NDArray[np.int64], lengths: NDArray[np.int64]) -> NumpyArray:
+    def _sum_rows(self, ids: NDArray[np.integer[Any]], lengths: NDArray[np.int64]) -> NumpyArray:
         """Sum the vectors of each document's tokens.
 
         The sums are a product of a (documents x distinct tokens) count matrix with the rows of

pyproject.toml

@@ -27,7 +27,7 @@ onnxruntime = [
 ]
 tqdm = "^4.66"
 requests = "^2.31"
-tokenizers = ">=0.15,<1.0"
+tokenizers = { version = ">=1.0.0rc2,<2.0", allow-prereleases = true }
 huggingface-hub = ">=0.20,<2.0"
 loguru = "^0.7.2"
 pillow = [
@@ -66,6 +66,20 @@ mypy = ">=1,<3"
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
+[[tool.mypy.overrides]]
+# Only the static embeddings support tokenizers 1.0. These modules still use the tokenizers 0.x API,
+# so they are not type checked against tokenizers 1.0.
+module = [
+    "fastembed.common.preprocessor_utils",
+    "fastembed.late_interaction.colbert",
+    "fastembed.rerank.cross_encoder.onnx_text_model",
+    "fastembed.sparse.bm42",
+    "fastembed.sparse.minicoil",
+    "fastembed.sparse.utils.vocab_resolver",
+    "fastembed.text.siglip_embedding",
+]
+ignore_errors = true
+
 [tool.pyright]
 typeCheckingMode = "strict"