Avoid ONNX session load for token counts

fastembed/late_interaction/colbert.py

@@ -104,8 +104,10 @@ 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
+        if not hasattr(self, "tokenizer") or self.tokenizer is None:
+            self._load_tokenizer(model_dir=self._model_dir)
+        if self.query_tokenizer is None:
+            self._load_query_tokenizer()
         token_num = 0
         texts = [texts] if isinstance(texts, str) else texts
         tokenizer = self.tokenizer if is_doc else self.query_tokenizer
@@ -218,6 +220,9 @@ def load_onnx_model(self) -> None:
             device_id=self.device_id,
             extra_session_options=self._extra_session_options,
         )
+        self._load_query_tokenizer()
+
+    def _load_query_tokenizer(self) -> None:
         self.query_tokenizer, _ = load_tokenizer(model_dir=self._model_dir)
 
         assert self.tokenizer is not None

fastembed/late_interaction_multimodal/colmodernvbert.py

@@ -203,8 +203,8 @@ 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
+        if not hasattr(self, "tokenizer") or self.tokenizer is None:
+            self._load_tokenizer(model_dir=self._model_dir)
         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,8 @@ 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
+        if not hasattr(self, "tokenizer") or self.tokenizer is None:
+            self._load_tokenizer(model_dir=self._model_dir)
         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,13 @@ 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._load_tokenizer(model_dir=model_dir)
         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 load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")
 

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -44,9 +44,12 @@ def _load_onnx_model(
             device_id=device_id,
             extra_session_options=extra_session_options,
         )
-        self.tokenizer, _ = load_tokenizer(model_dir=model_dir)
+        self._load_tokenizer(model_dir=model_dir)
         assert self.tokenizer is not None
 
+    def _load_tokenizer(self, model_dir: Path) -> None:
+        self.tokenizer, _ = load_tokenizer(model_dir=model_dir)
+
     def tokenize(self, pairs: list[tuple[str, str]], **_: Any) -> list[Encoding]:
         return self.tokenizer.encode_batch(pairs)  # type: ignore[union-attr]
 
@@ -168,8 +171,11 @@ 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
+        if not hasattr(self, "tokenizer") or self.tokenizer is None:
+            model_dir = getattr(self, "_model_dir", None)
+            if model_dir is None:
+                raise ValueError("Tokenizer cannot be loaded before model files are resolved.")
+            self._load_tokenizer(model_dir=Path(model_dir))
 
         token_num = 0
         assert self.tokenizer is not None

fastembed/sparse/bm42.py

@@ -355,8 +355,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/text/onnx_text_model.py

@@ -38,6 +38,9 @@ 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 _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], **kwargs: Any
     ) -> dict[str, NumpyArray | NDArray[np.int64]]:
@@ -65,7 +68,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._load_tokenizer(model_dir=model_dir)
 
     def load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")
@@ -167,8 +170,11 @@ 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
+        if not hasattr(self, "tokenizer") or self.tokenizer is None:
+            model_dir = getattr(self, "_model_dir", None)
+            if model_dir is None:
+                raise ValueError("Tokenizer cannot be loaded before model files are resolved.")
+            self._load_tokenizer(model_dir=Path(model_dir))
 
         token_num = 0
         assert self.tokenizer is not None

tests/test_late_interaction_embeddings.py

@@ -276,6 +276,8 @@ def test_lazy_load(model_name: str):
     assert not hasattr(model.model, "model")
 
     docs = ["hello world", "flag embedding"]
+    assert model.token_count(docs, is_doc=False, include_extension=True) > 0
+    assert not hasattr(model.model, "model")
     list(model.embed(docs))
     assert hasattr(model.model, "model")
 

tests/test_text_cross_encoder.py

@@ -100,6 +100,9 @@ def test_lazy_load(model_name: str) -> None:
     assert not hasattr(model.model, "model")
     query = "What is the capital of France?"
     documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
+    pairs = [(query, doc) for doc in documents]
+    assert model.token_count(pairs) > 0
+    assert not hasattr(model.model, "model")
     list(model.rerank(query, documents))
     assert hasattr(model.model, "model")
 

tests/test_text_onnx_embeddings.py

@@ -215,6 +215,8 @@ def test_lazy_load(model_name: str) -> None:
     model = TextEmbedding(model_name=model_name, lazy_load=True)
     assert not hasattr(model.model, "model")
     docs = ["hello world", "flag embedding"]
+    assert model.token_count(docs) > 0
+    assert not hasattr(model.model, "model")
     list(model.embed(docs))
     assert hasattr(model.model, "model")