refactor: split tokenize into _tokenize and tokenize to respect MRO (#566)

joein · dancixx · commit 7dba1d0c53d5 · 2025-11-18T14:00:13.000+01:00
diff --git a/fastembed/late_interaction/colbert.py b/fastembed/late_interaction/colbert.py
@@ -80,11 +80,16 @@ def _preprocess_onnx_input(
         )
         return onnx_input
 
-    def tokenize(self, texts: list[str], is_doc: bool = True, **kwargs: Any) -> list[Encoding]:  # type: ignore[override]
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
+        return self._tokenize(documents, **kwargs)
+
+    def _tokenize(
+        self, documents: list[str], is_doc: bool = True, **kwargs: Any
+    ) -> list[Encoding]:
         return (
-            self._tokenize_documents(documents=texts)
+            self._tokenize_documents(documents=documents)
             if is_doc
-            else self._tokenize_query(query=next(iter(texts)))
+            else self._tokenize_query(query=next(iter(documents)))
         )
 
     def _tokenize_query(self, query: str) -> list[Encoding]:
diff --git a/fastembed/late_interaction/late_interaction_embedding_base.py b/fastembed/late_interaction/late_interaction_embedding_base.py
@@ -21,7 +21,7 @@ def __init__(
         self._local_files_only = kwargs.pop("local_files_only", False)
         self._embedding_size: Optional[int] = None
 
-    def tokenize(self, texts: list[str], **kwargs: Any) -> list[Encoding]:
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
         raise NotImplementedError()
 
     def embed(
diff --git a/fastembed/late_interaction/late_interaction_text_embedding.py b/fastembed/late_interaction/late_interaction_text_embedding.py
@@ -116,18 +116,18 @@ def get_embedding_size(cls, model_name: str) -> int:
             )
         return embedding_size
 
-    def tokenize(self, texts: list[str], **kwargs: Any) -> list[Encoding]:
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
         """
         Tokenize input texts using the model's tokenizer.
 
         Args:
-            texts: List of strings to tokenize
+            documents: List of strings to tokenize
             **kwargs: Additional arguments passed to the tokenizer
 
         Returns:
             List of tokenizer Encodings
         """
-        return self.model.tokenize(texts, **kwargs)
+        return self.model.tokenize(documents, **kwargs)
 
     def embed(
         self,
diff --git a/fastembed/late_interaction/token_embeddings.py b/fastembed/late_interaction/token_embeddings.py
@@ -25,7 +25,7 @@
 ]
 
 
-class TokenEmbeddingsModel(OnnxTextEmbedding, LateInteractionTextEmbeddingBase):  # type: ignore[misc]
+class TokenEmbeddingsModel(OnnxTextEmbedding, LateInteractionTextEmbeddingBase):
     @classmethod
     def _list_supported_models(cls) -> list[DenseModelDescription]:
         """Lists the supported models.
diff --git a/fastembed/late_interaction_multimodal/colpali.py b/fastembed/late_interaction_multimodal/colpali.py
@@ -160,9 +160,12 @@ def _post_process_onnx_text_output(
         """
         return output.model_output
 
-    def tokenize(self, texts: list[str], **kwargs: Any) -> list[Encoding]:  # type: ignore[override]
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
+        return self._tokenize(documents, **kwargs)
+
+    def _tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
         texts_query: list[str] = []
-        for query in texts:
+        for query in documents:
             query = self.BOS_TOKEN + self.QUERY_PREFIX + query + self.PAD_TOKEN * 10
             query += "\n"
 
diff --git a/fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py b/fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py
@@ -119,18 +119,18 @@ def get_embedding_size(cls, model_name: str) -> int:
             )
         return embedding_size
 
-    def tokenize(self, texts: list[str], **kwargs: Any) -> list[Encoding]:
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
         """
         Tokenize input texts using the model's tokenizer.
 
         Args:
-            texts: List of strings to tokenize
+            documents: List of strings to tokenize
             **kwargs: Additional arguments passed to the tokenizer
 
         Returns:
             List of tokenizer Encodings
         """
-        return self.model.tokenize(texts, **kwargs)
+        return self.model.tokenize(documents, **kwargs)
 
     def embed_text(
         self,
diff --git a/fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding_base.py b/fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding_base.py
@@ -22,7 +22,7 @@ def __init__(
         self._local_files_only = kwargs.pop("local_files_only", False)
         self._embedding_size: Optional[int] = None
 
-    def tokenize(self, texts: list[str], **kwargs: Any) -> list[Encoding]:
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
         raise NotImplementedError()
 
     def embed_text(
diff --git a/fastembed/late_interaction_multimodal/onnx_multimodal_model.py b/fastembed/late_interaction_multimodal/onnx_multimodal_model.py
@@ -80,17 +80,17 @@ def _load_onnx_model(
     def load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def tokenize(self, texts: list[str], **kwargs: Any) -> list[Encoding]:
+    def _tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
         if self.tokenizer is None:
             raise RuntimeError("Tokenizer not initialized")
-        return self.tokenizer.encode_batch(texts, **kwargs)  # type: ignore[union-attr]
+        return self.tokenizer.encode_batch(documents, **kwargs)  # type: ignore[union-attr]
 
     def onnx_embed_text(
         self,
         documents: list[str],
         **kwargs: Any,
     ) -> OnnxOutputContext:
-        encoded = self.tokenize(documents, **kwargs)
+        encoded = self._tokenize(documents, **kwargs)
         input_ids = np.array([e.ids for e in encoded])
         attention_mask = np.array([e.attention_mask for e in encoded])  # type: ignore[union-attr]
         input_names = {node.name for node in self.model.get_inputs()}  # type: ignore[union-attr]
diff --git a/fastembed/rerank/cross_encoder/onnx_text_model.py b/fastembed/rerank/cross_encoder/onnx_text_model.py
@@ -46,8 +46,6 @@ def _load_onnx_model(
         assert self.tokenizer is not None
 
     def tokenize(self, pairs: list[tuple[str, str]], **kwargs: Any) -> list[Encoding]:
-        if self.tokenizer is None:
-            raise RuntimeError("Tokenizer not initialized")
         return self.tokenizer.encode_batch(pairs, **kwargs)  # type: ignore[union-attr]
 
     def _build_onnx_input(self, tokenized_input: list[Encoding]) -> dict[str, NumpyArray]:
diff --git a/fastembed/sparse/bm25.py b/fastembed/sparse/bm25.py
@@ -137,25 +137,25 @@ def __init__(
 
         self.tokenizer = SimpleTokenizer
 
-    def tokenize(self, texts: list[str], **kwargs: Any) -> list[Encoding]:
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
         """Tokenize texts using SimpleTokenizer.
 
         Returns a list of simple Encoding-like objects with token strings.
         Note: BM25 uses a simple word tokenizer, not a learned tokenizer.
         """
         result = []
-        for text in texts:
-            tokens = self.tokenizer.tokenize(text)
 
-            # Create a simple object that mimics Encoding interface
-            class SimpleEncoding:
-                def __init__(self, tokens: list[str]):
-                    self.tokens = tokens
-                    self.ids = tokens  # For BM25, tokens are the IDs
-                    self.attention_mask = [1] * len(tokens)
+        class SimpleEncoding:
+            def __init__(self, tokens: list[str]):
+                self.tokens = tokens
+                self.ids = tokens  # For BM25, tokens are the IDs
+                self.attention_mask = [1] * len(tokens)
+
+        for document in documents:
+            tokens = self.tokenizer.tokenize(document)
+            result.append(SimpleEncoding(tokens))
 
-            result.append(SimpleEncoding(tokens))  # type: ignore[arg-type]
-        return result  # type: ignore[return-value]
+        return result
 
     @classmethod
     def _list_supported_models(cls) -> list[SparseModelDescription]:
diff --git a/fastembed/sparse/bm42.py b/fastembed/sparse/bm42.py
@@ -45,7 +45,7 @@ def get_language_by_model_name(model_name: str) -> str:
     return MODEL_TO_LANGUAGE[model_name.lower()]
 
 
-class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):  # type: ignore[misc]
+class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     """
     Bm42 is an extension of BM25, which tries to better evaluate importance of tokens in the documents,
     by extracting attention weights from the transformer model.
@@ -139,8 +139,8 @@ def __init__(
         if not self.lazy_load:
             self.load_onnx_model()
 
-    def tokenize(self, texts: list[str], **kwargs: Any) -> list[Encoding]:
-        return OnnxTextModel.tokenize(self, list(texts), **kwargs)
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
+        return self._tokenize(documents, **kwargs)
 
     def load_onnx_model(self) -> None:
         self._load_onnx_model(
diff --git a/fastembed/sparse/minicoil.py b/fastembed/sparse/minicoil.py
@@ -58,7 +58,7 @@ def get_language_by_model_name(model_name: str) -> str:
     return MODEL_TO_LANGUAGE[model_name.lower()]
 
 
-class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):  # type: ignore[misc]
+class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     """
         MiniCOIL is a sparse embedding model, that resolves semantic meaning of the words,
         while keeping exact keyword match behavior.
@@ -145,8 +145,8 @@ def __init__(
         if not self.lazy_load:
             self.load_onnx_model()
 
-    def tokenize(self, texts: list[str], **kwargs: Any) -> list[Encoding]:
-        return OnnxTextModel.tokenize(self, list(texts), **kwargs)
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
+        return self._tokenize(documents, **kwargs)
 
     def load_onnx_model(self) -> None:
         self._load_onnx_model(
diff --git a/fastembed/sparse/sparse_embedding_base.py b/fastembed/sparse/sparse_embedding_base.py
@@ -45,7 +45,7 @@ def __init__(
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
 
-    def tokenize(self, texts: list[str], **kwargs: Any) -> list[Encoding]:
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
         raise NotImplementedError()
 
     def embed(
diff --git a/fastembed/sparse/sparse_text_embedding.py b/fastembed/sparse/sparse_text_embedding.py
@@ -93,18 +93,18 @@ def __init__(
             "Please check the supported models using `SparseTextEmbedding.list_supported_models()`"
         )
 
-    def tokenize(self, texts: list[str], **kwargs: Any) -> list[Encoding]:
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
         """
         Tokenize input texts using the model's tokenizer.
 
         Args:
-            texts: List of strings to tokenize
+            documents: List of strings to tokenize
             **kwargs: Additional arguments passed to the tokenizer
 
         Returns:
             List of tokenizer Encodings
         """
-        return self.model.tokenize(texts, **kwargs)
+        return self.model.tokenize(documents, **kwargs)
 
     def embed(
         self,
diff --git a/fastembed/sparse/splade_pp.py b/fastembed/sparse/splade_pp.py
@@ -137,20 +137,18 @@ def load_onnx_model(self) -> None:
             device_id=self.device_id,
         )
 
-    def tokenize(self, texts: list[str], **kwargs: Any) -> list[Encoding]:  # type: ignore[override]
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
         """
         Tokenize input texts using the model's tokenizer.
 
         Args:
-            texts: List of strings to tokenize
+            documents: List of strings to tokenize
             **kwargs: Additional arguments passed to the tokenizer
 
         Returns:
             List of tokenizer Encodings
         """
-        if self.tokenizer is None:
-            raise RuntimeError("Tokenizer not initialized")
-        return self.tokenizer.encode_batch(texts, **kwargs)
+        return self._tokenize(documents, **kwargs)
 
     def embed(
         self,
diff --git a/fastembed/text/onnx_embedding.py b/fastembed/text/onnx_embedding.py
@@ -185,7 +185,7 @@
 ]
 
 
-class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):  # type: ignore[misc]
+class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
     """Implementation of the Flag Embedding model."""
 
     @classmethod
@@ -321,17 +321,8 @@ def _post_process_onnx_output(
             raise ValueError(f"Unsupported embedding shape: {embeddings.shape}")
         return normalize(processed_embeddings)
 
-    def tokenize(self, texts: list[str], **kwargs: Any) -> list[Encoding]:
-        """Tokenize the input texts.
-
-        Args:
-            texts: A list of strings to tokenize.
-            **kwargs: Additional keyword arguments.
-
-        Returns:
-            list[Encoding]: List of tokenized encodings.
-        """
-        return OnnxTextModel.tokenize(self, texts, **kwargs)
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
+        return self._tokenize(documents, **kwargs)
 
     def load_onnx_model(self) -> None:
         self._load_onnx_model(
diff --git a/fastembed/text/onnx_text_model.py b/fastembed/text/onnx_text_model.py
@@ -68,15 +68,15 @@ def _load_onnx_model(
     def load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def tokenize(self, texts: list[str], **kwargs: Any) -> list[Encoding]:
-        return self.tokenizer.encode_batch(texts)  # type: ignore[union-attr]
+    def _tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
+        return self.tokenizer.encode_batch(documents)  # type:ignore[union-attr]
 
     def onnx_embed(
         self,
         documents: list[str],
         **kwargs: Any,
     ) -> OnnxOutputContext:
-        encoded = self.tokenize(documents, **kwargs)
+        encoded = self._tokenize(documents, **kwargs)
         input_ids = np.array([e.ids for e in encoded])
         attention_mask = np.array([e.attention_mask for e in encoded])
         input_names = {node.name for node in self.model.get_inputs()}  # type: ignore[union-attr]
diff --git a/fastembed/text/text_embedding.py b/fastembed/text/text_embedding.py
@@ -163,18 +163,18 @@ def get_embedding_size(cls, model_name: str) -> int:
             )
         return embedding_size
 
-    def tokenize(self, texts: list[str], **kwargs: Any) -> list[Encoding]:
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
         """
         Tokenize input texts using the model's tokenizer.
 
         Args:
-            texts: List of strings to tokenize
+            documents: List of strings to tokenize
             **kwargs: Additional arguments passed to the tokenizer
 
         Returns:
             List of tokenizer Encodings
         """
-        return self.model.tokenize(texts, **kwargs)
+        return self.model.tokenize(documents, **kwargs)
 
     def embed(
         self,
diff --git a/fastembed/text/text_embedding_base.py b/fastembed/text/text_embedding_base.py
@@ -20,7 +20,7 @@ def __init__(
         self._local_files_only = kwargs.pop("local_files_only", False)
         self._embedding_size: Optional[int] = None
 
-    def tokenize(self, texts: list[str], **kwargs: Any) -> list[Encoding]:
+    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
         raise NotImplementedError()
 
     def embed(
diff --git a/tests/test_late_interaction_embeddings.py b/tests/test_late_interaction_embeddings.py
@@ -249,23 +249,27 @@ def test_single_embedding_query(model_cache, model_name: str):
 
 
 @pytest.mark.parametrize("token_dim,model_name", [(96, "answerdotai/answerai-colbert-small-v1")])
-def test_parallel_processing(model_cache, token_dim: int, model_name: str):
-    with model_cache(model_name) as model:
-        docs = ["hello world", "flag embedding"] * 100
-        embeddings = list(model.embed(docs, batch_size=10, parallel=2))
 
-        embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
+def test_parallel_processing(token_dim: int, model_name: str):
+    # this test loads a copy of a model per process, might cause oom in parallel=0 on machines with
+    # an insufficient mem-to-cpus-ratio
+    is_ci = os.getenv("CI")
+    model = LateInteractionTextEmbedding(model_name=model_name)
+    docs = ["hello world", "flag embedding"] * 100
+    embeddings = list(model.embed(docs, batch_size=10, parallel=2))
+
+    embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
 
-        # embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))  # inherits OnnxTextModel which
-        #         # is tested in TextEmbedding, disabling it here to reduce number of requests to hf
-        #         # multiprocessing is enough to test with `parallel=2`, and `parallel=None` is okay to tests since it reuses
-        #         # model from cache
+    # embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))  # inherits OnnxTextModel which
+    #         # is tested in TextEmbedding, disabling it here to reduce number of requests to hf
+    #         # multiprocessing is enough to test with `parallel=2`, and `parallel=None` is okay to tests since it reuses
+    #         # model from cache
 
-        assert len(embeddings) == len(docs) and embeddings[0].shape[-1] == token_dim
+    assert len(embeddings) == len(docs) and embeddings[0].shape[-1] == token_dim
 
-        for i in range(len(embeddings)):
-            assert np.allclose(embeddings[i], embeddings_2[i], atol=1e-3)
-            # assert np.allclose(embeddings[i], embeddings_3[i], atol=1e-3)
+    for i in range(len(embeddings)):
+        assert np.allclose(embeddings[i], embeddings_2[i], atol=1e-3)
+        # assert np.allclose(embeddings[i], embeddings_3[i], atol=1e-3)
 
 
 @pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
@@ -325,7 +329,8 @@ def test_tokenize(model_name: str) -> None:
     enc_query = model.tokenize(["hello world"], is_doc=False)
     assert len(enc_query) == 1
     assert enc_query[0].ids is not None
-    assert len(enc_query[0].ids) > 0
+    assert len(enc_query[0].ids) == 31  # colbert requires query to be at least 32 tokens,
+    # padding is done during tokenization, the last token is added preprocess onnx input
 
     doc_ids = list(enc_doc[0].ids)
     query_ids = list(enc_query[0].ids)
diff --git a/tests/test_sparse_embeddings.py b/tests/test_sparse_embeddings.py
diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py

Original file line number	Diff line number	Diff line change
`@@ -25,7 +25,7 @@`
`25`	`25`	`]`
`26`	`26`
`27`	`27`
`28`		`-class TokenEmbeddingsModel(OnnxTextEmbedding, LateInteractionTextEmbeddingBase): # type: ignore[misc]`
	`28`	`+class TokenEmbeddingsModel(OnnxTextEmbedding, LateInteractionTextEmbeddingBase):`
`29`	`29`	`@classmethod`
`30`	`30`	`def _list_supported_models(cls) -> list[DenseModelDescription]:`
`31`	`31`	`"""Lists the supported models.`