Rerank type hints (#459)

* chore: Update type hints

* remove redundant array creation, update type hints

---------

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

Author: hh-space-invader

fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py

@@ -134,15 +134,14 @@ def __init__(
             )
 
         # This device_id will be used if we need to load model in current process
+        self.device_id: Optional[int] = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
             self.device_id = self.device_ids[0]
-        else:
-            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,

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -1,10 +1,9 @@
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Optional, Sequence, Type
 
 import numpy as np
-from numpy.typing import NDArray
 from tokenizers import Encoding
 
 from fastembed.common.onnx_model import (
@@ -13,6 +12,7 @@
     OnnxOutputContext,
     OnnxProvider,
 )
+from fastembed.common.types import NumpyArray
 from fastembed.common.preprocessor_utils import load_tokenizer
 from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool
@@ -47,11 +47,9 @@ def _load_onnx_model(
     def tokenize(self, pairs: list[tuple[str, str]], **_: Any) -> list[Encoding]:
         return self.tokenizer.encode_batch(pairs)
 
-    def _build_onnx_input(
-        self, tokenized_input
-    ) -> dict[str, NDArray[Union[np.float32, np.int64]]]:
-        input_names = {node.name for node in self.model.get_inputs()}
-        inputs = {
+    def _build_onnx_input(self, tokenized_input: list[Encoding]) -> dict[str, NumpyArray]:
+        input_names: set[str] = {node.name for node in self.model.get_inputs()}
+        inputs: dict[str, NumpyArray] = {
             "input_ids": np.array([enc.ids for enc in tokenized_input], dtype=np.int64),
         }
         if "token_type_ids" in input_names:
@@ -74,7 +72,7 @@ def onnx_embed_pairs(self, pairs: list[tuple[str, str]], **kwargs: Any) -> OnnxO
         onnx_input = self._preprocess_onnx_input(inputs, **kwargs)
         outputs = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)
         relevant_output = outputs[0]
-        scores = relevant_output[:, 0]
+        scores: NumpyArray = relevant_output[:, 0]
         return OnnxOutputContext(model_output=scores)
 
     def _rerank_documents(
@@ -100,7 +98,7 @@ def _rerank_pairs(
         is_small = False
 
         if isinstance(pairs, tuple):
-            pairs = [pairs]
+            pairs = [pairs]  # type: ignore
             is_small = True
 
         if isinstance(pairs, list):
@@ -138,15 +136,32 @@ def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[float
         raise NotImplementedError("Subclasses must implement this method")
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
-    ) -> dict[str, np.ndarray]:
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
         """
         Preprocess the onnx input.
         """
         return onnx_input
 
 
-class TextRerankerWorker(EmbeddingWorker):
+class TextRerankerWorker(EmbeddingWorker[float]):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ):
+        self.model: OnnxCrossEncoderModel
+        super().__init__(model_name, cache_dir, **kwargs)
+
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> OnnxCrossEncoderModel:
+        raise NotImplementedError()
+
     def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
             onnx_output = self.model.onnx_embed_pairs(batch)

fastembed/rerank/cross_encoder/text_cross_encoder.py

@@ -33,7 +33,7 @@ def list_supported_models(cls) -> list[dict[str, Any]]:
                 ]
                 ```
         """
-        result = []
+        result: list[dict[str, Any]] = []
         for encoder in cls.CROSS_ENCODER_REGISTRY:
             result.extend(encoder.list_supported_models())
         return result