Avoid ONNX session load for token counts

fastembed/common/onnx_model.py

@@ -26,6 +26,8 @@ class OnnxOutputContext:
 
 class OnnxModel(Generic[T]):
     EXPOSED_SESSION_OPTIONS = ("enable_cpu_mem_arena",)
+    # set by the concrete model classes in their __init__
+    _model_dir: Path
 
     @classmethod
     def _get_worker_class(cls) -> Type["EmbeddingWorker[T]"]:

fastembed/late_interaction/colbert.py

@@ -1,4 +1,5 @@
 import string
+from pathlib import Path
 from typing import Any, Iterable, Sequence, Type
 
 import numpy as np
@@ -55,6 +56,10 @@ def _post_process_onnx_output(
                     "input_ids and attention_mask must be provided for document post-processing"
                 )
 
+            # with `lazy_load` and `parallel`, inference runs in the workers and the parent
+            # never calls `load_onnx_model`, so `skip_list` and `pad_token_id` might not be set yet
+            self._ensure_tokenizer()
+
             for i, token_sequence in enumerate(output.input_ids):
                 for j, token_id in enumerate(token_sequence):  # type: ignore
                     if token_id in self.skip_list or token_id == self.pad_token_id:
@@ -104,8 +109,7 @@ def token_count(
         include_extension: bool = False,
         **kwargs: Any,
     ) -> int:
-        if not hasattr(self, "model") or self.model is None:
-            self.load_onnx_model()  # loads the tokenizer as well
+        self._ensure_tokenizer()
         token_num = 0
         texts = [texts] if isinstance(texts, str) else texts
         tokenizer = self.tokenizer if is_doc else self.query_tokenizer
@@ -218,27 +222,41 @@ def load_onnx_model(self) -> None:
             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
+    def _load_tokenizer(self, model_dir: Path) -> None:
+        """Load the document and the query tokenizers, and apply colbert's own configuration.
+
+        Both tokenizers are always (re)loaded from disk here, so the truncation length is
+        derived from the on-disk value and never decremented twice.
+        """
+        tokenizer, special_token_to_id = load_tokenizer(model_dir=model_dir)
+        query_tokenizer, _ = load_tokenizer(model_dir=model_dir)
+
         # load_tokenizer always enables both
-        assert self.tokenizer.padding is not None and self.tokenizer.truncation is not None
-        self.mask_token_id = self.special_token_to_id[self.MASK_TOKEN]
-        self.pad_token_id = self.tokenizer.padding["pad_id"]
-        self.skip_list = {
-            self.tokenizer.encode(symbol, add_special_tokens=False).ids[0]
+        assert tokenizer.padding is not None and tokenizer.truncation is not None
+        mask_token_id = special_token_to_id[self.MASK_TOKEN]
+        pad_token_id = tokenizer.padding["pad_id"]
+        skip_list = {
+            tokenizer.encode(symbol, add_special_tokens=False).ids[0]
             for symbol in string.punctuation
         }
-        current_max_length = self.tokenizer.truncation["max_length"]
         # ensure not to overflow after adding document-marker
-        self.tokenizer.enable_truncation(max_length=current_max_length - 1)
-        self.query_tokenizer.enable_truncation(max_length=current_max_length - 1)
-        self.query_tokenizer.enable_padding(
+        max_length = tokenizer.truncation["max_length"] - 1
+        tokenizer.enable_truncation(max_length=max_length)
+        query_tokenizer.enable_truncation(max_length=max_length)
+        query_tokenizer.enable_padding(
             pad_token=self.MASK_TOKEN,
-            pad_id=self.mask_token_id,
+            pad_id=mask_token_id,
             length=self.MIN_QUERY_LENGTH,
         )
 
+        self.query_tokenizer = query_tokenizer
+        self.special_token_to_id = special_token_to_id
+        self.mask_token_id = mask_token_id
+        self.pad_token_id = pad_token_id
+        self.skip_list = skip_list
+        self.tokenizer = tokenizer
+
     def embed(
         self,
         documents: str | Iterable[str],

fastembed/late_interaction_multimodal/colmodernvbert.py

@@ -203,8 +203,7 @@ def token_count(
         include_extension: bool = False,
         **kwargs: Any,
     ) -> int:
-        if not hasattr(self, "model") or self.model is None:
-            self.load_onnx_model()  # loads the tokenizer as well
+        self._ensure_tokenizer()
         token_num = 0
         texts = [texts] if isinstance(texts, str) else texts
         assert self.tokenizer is not None

fastembed/late_interaction_multimodal/colpali.py

@@ -180,8 +180,7 @@ def token_count(
         include_extension: bool = False,
         **kwargs: Any,
     ) -> int:
-        if not hasattr(self, "model") or self.model is None:
-            self.load_onnx_model()  # loads the tokenizer as well
+        self._ensure_tokenizer()
         token_num = 0
         texts = [texts] if isinstance(texts, str) else texts
         assert self.tokenizer is not None

fastembed/late_interaction_multimodal/onnx_multimodal_model.py

@@ -75,10 +75,22 @@ def _load_onnx_model(
             device_id=device_id,
             extra_session_options=extra_session_options,
         )
-        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
+        self._ensure_tokenizer()
         assert self.tokenizer is not None
         self.processor = load_preprocessor(model_dir=model_dir)
 
+    def _load_tokenizer(self, model_dir: Path) -> None:
+        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
+
+    def _ensure_tokenizer(self) -> None:
+        """Load the tokenizer if it has not been loaded yet, leaving the onnx model alone.
+
+        Tokenizer files are lightweight, so counting tokens does not have to resolve
+        `lazy_load` into a full onnx session.
+        """
+        if getattr(self, "tokenizer", None) is None:
+            self._load_tokenizer(model_dir=self._model_dir)
+
     def load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")
 

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -44,9 +44,21 @@ def _load_onnx_model(
             device_id=device_id,
             extra_session_options=extra_session_options,
         )
-        self.tokenizer, _ = load_tokenizer(model_dir=model_dir)
+        self._ensure_tokenizer()
         assert self.tokenizer is not None
 
+    def _load_tokenizer(self, model_dir: Path) -> None:
+        self.tokenizer, _ = load_tokenizer(model_dir=model_dir)
+
+    def _ensure_tokenizer(self) -> None:
+        """Load the tokenizer if it has not been loaded yet, leaving the onnx model alone.
+
+        Tokenizer files are lightweight, so counting tokens does not have to resolve
+        `lazy_load` into a full onnx session.
+        """
+        if getattr(self, "tokenizer", None) is None:
+            self._load_tokenizer(model_dir=self._model_dir)
+
     def tokenize(self, pairs: list[tuple[str, str]], **_: Any) -> list[Encoding]:
         return self.tokenizer.encode_batch(pairs)  # type: ignore[union-attr]
 
@@ -169,9 +181,7 @@ def _preprocess_onnx_input(
     def _token_count(
         self, pairs: Iterable[tuple[str, str]], batch_size: int = 1024, **_: Any
     ) -> int:
-        if not hasattr(self, "model") or self.model is None:
-            self.load_onnx_model()  # loads the tokenizer as well
-
+        self._ensure_tokenizer()
         token_num = 0
         assert self.tokenizer is not None
         for batch in iter_batch(pairs, batch_size):

fastembed/sparse/bm42.py

@@ -141,6 +141,16 @@ def __init__(
         if not self.lazy_load:
             self.load_onnx_model()
 
+    def _load_tokenizer(self, model_dir: Path) -> None:
+        super()._load_tokenizer(model_dir=model_dir)
+        # `_reconstruct_bpe` needs the special tokens, and it is also used by `query_embed`,
+        # which runs no inference, so they are derived here rather than in `load_onnx_model`.
+        # The same goes for `invert_vocab`: with `lazy_load` and `parallel`, post-processing
+        # runs in a parent process which never loads the onnx model.
+        self.special_tokens = set(self.special_token_to_id.keys())
+        self.special_tokens_ids = set(self.special_token_to_id.values())
+        self.invert_vocab = {idx: token for token, idx in self.tokenizer.get_vocab().items()}  # type: ignore[union-attr]
+
     def load_onnx_model(self) -> None:
         self._load_onnx_model(
             model_dir=self._model_dir,
@@ -152,10 +162,6 @@ def load_onnx_model(self) -> None:
             extra_session_options=self._extra_session_options,
         )
 
-        for token, idx in self.tokenizer.get_vocab().items():  # type: ignore[union-attr]
-            self.invert_vocab[idx] = token
-        self.special_tokens = set(self.special_token_to_id.keys())
-        self.special_tokens_ids = set(self.special_token_to_id.values())
         self.stopwords = set(self._load_stopwords(self._model_dir))
 
     def _filter_pair_tokens(self, tokens: list[tuple[str, Any]]) -> list[tuple[str, Any]]:
@@ -236,6 +242,10 @@ def _post_process_onnx_output(
         if output.input_ids is None:
             raise ValueError("input_ids must be provided for document post-processing")
 
+        # with `lazy_load` and `parallel`, inference runs in the workers and the parent
+        # never calls `load_onnx_model`, so the tokenizer might not be loaded here yet
+        self._ensure_tokenizer()
+
         token_ids_batch = output.input_ids.astype(int)
 
         # attention_value shape: (batch_size, num_heads, num_tokens, num_tokens)
@@ -336,8 +346,7 @@ def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[Spa
         if isinstance(query, str):
             query = [query]
 
-        if not hasattr(self, "model") or self.model is None:
-            self.load_onnx_model()
+        self._ensure_tokenizer()
 
         for text in query:
             encoded = self.tokenizer.encode(text)  # type: ignore[union-attr]
@@ -355,8 +364,6 @@ def _get_worker_class(cls) -> Type[TextEmbeddingWorker[SparseEmbedding]]:
     def token_count(
         self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
     ) -> int:
-        if not hasattr(self, "model") or self.model is None:
-            self.load_onnx_model()  # loads the tokenizer as well
         return self._token_count(texts, batch_size=batch_size, **kwargs)
 
 

fastembed/sparse/if_splade.py

@@ -6,9 +6,8 @@
 from fastembed.common import OnnxProvider
 from fastembed.common.model_description import ModelSource, SparseModelDescription
 from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.preprocessor_utils import load_tokenizer
 from fastembed.common.types import Device
-from fastembed.common.utils import define_cache_dir, iter_batch
+from fastembed.common.utils import define_cache_dir
 from fastembed.sparse.sparse_embedding_base import (
     SparseEmbedding,
     SparseTextEmbeddingBase,
@@ -69,13 +68,7 @@ def _post_process_onnx_output(
     def token_count(
         self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
     ) -> int:
-        # unlike `OnnxTextModel._token_count`, does not require the onnx model to be loaded
-        token_num = 0
-        texts = [texts] if isinstance(texts, str) else texts
-        for batch in iter_batch(texts, batch_size):
-            for tokens in self.tokenizer.encode_batch(batch):  # type: ignore[union-attr]
-                token_num += sum(tokens.attention_mask)
-        return token_num
+        return self._token_count(texts, batch_size=batch_size, **kwargs)
 
     @classmethod
     def _list_supported_models(cls) -> list[SparseModelDescription]:
@@ -151,7 +144,7 @@ def __init__(
         # The tokenizer and the idf table are lightweight and are required for query embedding,
         # which does not involve any model inference, so they are loaded eagerly, while
         # `lazy_load` only defers the initialization of the onnx model
-        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=self._model_dir)
+        self._load_tokenizer(model_dir=self._model_dir)
         self.special_tokens_ids: set[int] = set(self.special_token_to_id.values())
         self._token_id_to_idf = self._load_idf()
 

fastembed/sparse/minicoil.py

@@ -160,35 +160,48 @@ def load_onnx_model(self) -> None:
             extra_session_options=self._extra_session_options,
         )
 
+    def _load_tokenizer(self, model_dir: Path) -> None:
+        super()._load_tokenizer(model_dir=model_dir)
         assert self.tokenizer is not None
-
-        for token, idx in self.tokenizer.get_vocab().items():  # type: ignore[union-attr]
-            self.invert_vocab[idx] = token
+        self.invert_vocab = {idx: token for token, idx in self.tokenizer.get_vocab().items()}
         self.special_tokens = set(self.special_token_to_id.keys())
         self.special_tokens_ids = set(self.special_token_to_id.values())
-        self.stopwords = set(self._load_stopwords(self._model_dir))
 
+    def _load_post_processing_state(self) -> None:
+        """Load the vocab resolver, the encoder and the sparse vector converter.
+
+        None of them needs the onnx session, so the parent process can post-process the outputs
+        of parallel workers without loading it. Requires the tokenizer to be loaded.
+        """
+        assert self.tokenizer is not None
+
+        stopwords = set(self._load_stopwords(self._model_dir))
         stemmer = SnowballStemmer(get_language_by_model_name(self.model_name))
 
-        self.vocab_resolver = VocabResolver(
+        vocab_resolver = VocabResolver(
             tokenizer=VocabTokenizer(self.tokenizer),
-            stopwords=self.stopwords,
+            stopwords=stopwords,
             stemmer=stemmer,
         )
-        self.vocab_resolver.load_json_vocab(str(self._model_dir / MINICOIL_VOCAB_FILE))
+        vocab_resolver.load_json_vocab(str(self._model_dir / MINICOIL_VOCAB_FILE))
 
         weights = np.load(str(self._model_dir / MINICOIL_MODEL_FILE), mmap_mode="r")
-        self.encoder = Encoder(weights)
-        self.output_dim = self.encoder.output_dim
+        encoder = Encoder(weights)
 
-        self.sparse_vector_converter = SparseVectorConverter(
-            stopwords=self.stopwords,
+        sparse_vector_converter = SparseVectorConverter(
+            stopwords=stopwords,
             stemmer=stemmer,
             k=self.k,
             b=self.b,
             avg_len=self.avg_len,
         )
 
+        self.stopwords = stopwords
+        self.encoder = encoder
+        self.output_dim = encoder.output_dim
+        self.sparse_vector_converter = sparse_vector_converter
+        self.vocab_resolver = vocab_resolver
+
     def token_count(
         self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any
     ) -> int:
@@ -279,6 +292,12 @@ def _post_process_onnx_output(
         if output.input_ids is None:
             raise ValueError("input_ids must be provided for document post-processing")
 
+        if self.vocab_resolver is None:
+            # built on first use instead of in `load_onnx_model`: parallel workers never
+            # post-process, and with `lazy_load` the process that does may never load the model
+            self._ensure_tokenizer()
+            self._load_post_processing_state()
+
         assert self.vocab_resolver is not None
         assert self.encoder is not None
         assert self.sparse_vector_converter is not None

fastembed/text/onnx_text_model.py

@@ -38,6 +38,18 @@ def __init__(self) -> None:
         self.tokenizer: Tokenizer | None = None
         self.special_token_to_id: dict[str, int] = {}
 
+    def _load_tokenizer(self, model_dir: Path) -> None:
+        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
+
+    def _ensure_tokenizer(self) -> None:
+        """Load the tokenizer if it has not been loaded yet, leaving the onnx model alone.
+
+        Tokenizer files are lightweight, so counting tokens does not have to resolve
+        `lazy_load` into a full onnx session.
+        """
+        if getattr(self, "tokenizer", None) is None:
+            self._load_tokenizer(model_dir=self._model_dir)
+
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], **kwargs: Any
     ) -> dict[str, NumpyArray | NDArray[np.int64]]:
@@ -65,7 +77,7 @@ def _load_onnx_model(
             device_id=device_id,
             extra_session_options=extra_session_options,
         )
-        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
+        self._ensure_tokenizer()
 
     def load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")
@@ -168,9 +180,7 @@ def _embed_documents(
                 yield from self._post_process_onnx_output(batch, **kwargs)  # type: ignore
 
     def _token_count(self, texts: str | Iterable[str], batch_size: int = 1024, **_: Any) -> int:
-        if not hasattr(self, "model") or self.model is None:
-            self.load_onnx_model()  # loads the tokenizer as well
-
+        self._ensure_tokenizer()
         token_num = 0
         assert self.tokenizer is not None
         texts = [texts] if isinstance(texts, str) else texts