Late interaction type hints (#461)

* chore: Add type hints

* new: Add late_interaction type hints

* fix: ndarray -> numpy array

---------

Co-authored-by: George Panchuk <[email protected]>

Author: hh-space-invader

fastembed/late_interaction/colbert.py

@@ -4,6 +4,7 @@
 import numpy as np
 from tokenizers import Encoding
 
+from fastembed.common.types import NumpyArray
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import define_cache_dir
@@ -39,15 +40,15 @@
 ]
 
 
-class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
+class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
     QUERY_MARKER_TOKEN_ID = 1
     DOCUMENT_MARKER_TOKEN_ID = 2
     MIN_QUERY_LENGTH = 31  # it's 32, we add one additional special token in the beginning
     MASK_TOKEN = "[MASK]"
 
     def _post_process_onnx_output(
         self, output: OnnxOutputContext, is_doc: bool = True
-    ) -> Iterable[np.ndarray]:
+    ) -> Iterable[NumpyArray]:
         if not is_doc:
             return output.model_output.astype(np.float32)
 
@@ -68,11 +69,15 @@ def _post_process_onnx_output(
         return output.model_output.astype(np.float32)
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, np.ndarray], is_doc: bool = True, **kwargs: Any
-    ) -> dict[str, np.ndarray]:
+        self, onnx_input: dict[str, NumpyArray], is_doc: bool = True, **kwargs: Any
+    ) -> dict[str, NumpyArray]:
         marker_token = self.DOCUMENT_MARKER_TOKEN_ID if is_doc else self.QUERY_MARKER_TOKEN_ID
-        onnx_input["input_ids"] = np.insert(onnx_input["input_ids"], 1, marker_token, axis=1)
-        onnx_input["attention_mask"] = np.insert(onnx_input["attention_mask"], 1, 1, axis=1)
+        onnx_input["input_ids"] = np.insert(
+            onnx_input["input_ids"].astype(np.int64), 1, marker_token, axis=1
+        )
+        onnx_input["attention_mask"] = np.insert(
+            onnx_input["attention_mask"].astype(np.int64), 1, 1, axis=1
+        )
         return onnx_input
 
     def tokenize(self, documents: list[str], is_doc: bool = True, **kwargs: Any) -> list[Encoding]:
@@ -166,17 +171,17 @@ def __init__(
             self.device_id = None
 
         self.model_description = self._get_model_description(model_name)
-        self.cache_dir = define_cache_dir(cache_dir)
+        self.cache_dir = str(define_cache_dir(cache_dir))
 
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
             specific_model_path=specific_model_path,
         )
-        self.mask_token_id = None
-        self.pad_token_id = None
-        self.skip_list = set()
+        self.mask_token_id: Optional[int] = None
+        self.pad_token_id: Optional[int] = None
+        self.skip_list: set[str] = set()
 
         if not self.lazy_load:
             self.load_onnx_model()
@@ -206,7 +211,7 @@ def embed(
         batch_size: int = 256,
         parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[np.ndarray]:
+    ) -> Iterable[NumpyArray]:
         """
         Encode a list of documents into list of embeddings.
         We use mean pooling with attention so that the model can handle variable-length inputs.
@@ -234,7 +239,7 @@ def embed(
             **kwargs,
         )
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
         if isinstance(query, str):
             query = [query]
 
@@ -247,11 +252,11 @@ def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterab
             )
 
     @classmethod
-    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker[NumpyArray]]:
         return ColbertEmbeddingWorker
 
 
-class ColbertEmbeddingWorker(TextEmbeddingWorker):
+class ColbertEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
     def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> Colbert:
         return Colbert(
             model_name=model_name,

fastembed/late_interaction/jina_colbert.py

@@ -1,9 +1,7 @@
 from typing import Any, Type
 
-import numpy as np
-
-from fastembed.late_interaction.colbert import Colbert
-from fastembed.text.onnx_text_model import TextEmbeddingWorker
+from fastembed.common.types import NumpyArray
+from fastembed.late_interaction.colbert import Colbert, ColbertEmbeddingWorker
 
 
 supported_jina_colbert_models = [
@@ -29,7 +27,7 @@ class JinaColbert(Colbert):
     MASK_TOKEN = "<mask>"
 
     @classmethod
-    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
+    def _get_worker_class(cls) -> Type[ColbertEmbeddingWorker]:
         return JinaColbertEmbeddingWorker
 
     @classmethod
@@ -42,8 +40,8 @@ def list_supported_models(cls) -> list[dict[str, Any]]:
         return supported_jina_colbert_models
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, np.ndarray], is_doc: bool = True, **kwargs: Any
-    ) -> dict[str, np.ndarray]:
+        self, onnx_input: dict[str, NumpyArray], is_doc: bool = True, **kwargs: Any
+    ) -> dict[str, NumpyArray]:
         onnx_input = super()._preprocess_onnx_input(onnx_input, is_doc)
 
         # the attention mask for jina-colbert-v2 is always 1 in queries
@@ -52,7 +50,7 @@ def _preprocess_onnx_input(
         return onnx_input
 
 
-class JinaColbertEmbeddingWorker(TextEmbeddingWorker):
+class JinaColbertEmbeddingWorker(ColbertEmbeddingWorker):
     def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> JinaColbert:
         return JinaColbert(
             model_name=model_name,

fastembed/late_interaction/late_interaction_embedding_base.py

@@ -1,7 +1,6 @@
 from typing import Iterable, Optional, Union, Any
 
-import numpy as np
-
+from fastembed.common.types import NumpyArray
 from fastembed.common.model_management import ModelManagement
 
 
@@ -24,10 +23,10 @@ def embed(
         batch_size: int = 256,
         parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[np.ndarray]:
+    ) -> Iterable[NumpyArray]:
         raise NotImplementedError()
 
-    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[np.ndarray]:
+    def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds a list of text passages into a list of embeddings.
 
@@ -36,25 +35,25 @@ def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[np.ndar
             **kwargs: Additional keyword argument to pass to the embed method.
 
         Yields:
-            Iterable[np.ndarray]: The embeddings.
+            Iterable[NdArray]: The embeddings.
         """
 
         # This is model-specific, so that different models can have specialized implementations
         yield from self.embed(texts, **kwargs)
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds queries
 
         Args:
             query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
 
         Returns:
-            Iterable[np.ndarray]: The embeddings.
+            Iterable[NdArray]: The embeddings.
         """
 
         # This is model-specific, so that different models can have specialized implementations
         if isinstance(query, str):
             yield from self.embed([query], **kwargs)
-        if isinstance(query, Iterable):
+        else:
             yield from self.embed(query, **kwargs)

fastembed/late_interaction/late_interaction_text_embedding.py

@@ -1,7 +1,6 @@
 from typing import Any, Iterable, Optional, Sequence, Type, Union
 
-import numpy as np
-
+from fastembed.common.types import NumpyArray
 from fastembed.common import OnnxProvider
 from fastembed.late_interaction.colbert import Colbert
 from fastembed.late_interaction.jina_colbert import JinaColbert
@@ -38,7 +37,7 @@ def list_supported_models(cls) -> list[dict[str, Any]]:
                 ]
                 ```
         """
-        result = []
+        result: list[dict[str, Any]] = []
         for embedding in cls.EMBEDDINGS_REGISTRY:
             result.extend(embedding.list_supported_models())
         return result
@@ -81,7 +80,7 @@ def embed(
         batch_size: int = 256,
         parallel: Optional[int] = None,
         **kwargs: Any,
-    ) -> Iterable[np.ndarray]:
+    ) -> Iterable[NumpyArray]:
         """
         Encode a list of documents into list of embeddings.
         We use mean pooling with attention so that the model can handle variable-length inputs.
@@ -99,15 +98,15 @@ def embed(
         """
         yield from self.model.embed(documents, batch_size, parallel, **kwargs)
 
-    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
         """
         Embeds queries
 
         Args:
             query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
 
         Returns:
-            Iterable[np.ndarray]: The embeddings.
+            Iterable[NdArray]: The embeddings.
         """
 
         # This is model-specific, so that different models can have specialized implementations