chore: Add missing returns in defs (#451)

* chore: Add missing returns in defs

* remove return type from init

* remove incorrect ndarray specifier

---------

Co-authored-by: George Panchuk <george.panchuk@qdrant.tech>

Author: hh-space-invader

docs/examples/ColBERT_with_FastEmbed.ipynb

@@ -54,7 +54,14 @@
     },
     {
      "data": {
-      "text/plain": "[{'model': 'colbert-ir/colbertv2.0',\n  'dim': 128,\n  'description': 'Late interaction model',\n  'size_in_GB': 0.44,\n  'sources': {'hf': 'colbert-ir/colbertv2.0'},\n  'model_file': 'model.onnx'}]"
+      "text/plain": [
+       "[{'model': 'colbert-ir/colbertv2.0',\n",
+       "  'dim': 128,\n",
+       "  'description': 'Late interaction model',\n",
+       "  'size_in_GB': 0.44,\n",
+       "  'sources': {'hf': 'colbert-ir/colbertv2.0'},\n",
+       "  'model_file': 'model.onnx'}]"
+      ]
      },
      "execution_count": 1,
      "metadata": {},
@@ -212,7 +219,9 @@
    "outputs": [
     {
      "data": {
-      "text/plain": "((26, 128), (32, 128))"
+      "text/plain": [
+       "((26, 128), (32, 128))"
+      ]
      },
      "execution_count": 18,
      "metadata": {},
@@ -271,7 +280,9 @@
     "import numpy as np\n",
     "\n",
     "\n",
-    "def compute_relevance_scores(query_embedding: np.array, document_embeddings: np.array, k: int):\n",
+    "def compute_relevance_scores(\n",
+    "    query_embedding: np.array, document_embeddings: np.array, k: int\n",
+    ") -> list[int]:\n",
     "    \"\"\"\n",
     "    Compute relevance scores for top-k documents given a query.\n",
     "\n",

docs/examples/FastEmbed_vs_HF_Comparison.ipynb

@@ -152,7 +152,7 @@
     "    HuggingFace Transformer implementation of FlagEmbedding\n",
     "    \"\"\"\n",
     "\n",
-    "    def __init__(self, model_id: str):\n",
+    "    def __init__(self, model_id: str) -> None:\n",
     "        self.model = AutoModel.from_pretrained(model_id)\n",
     "        self.tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
     "\n",

docs/examples/Hybrid_Search.ipynb

@@ -488,7 +488,7 @@
     }
    ],
    "source": [
-    "def make_sparse_embedding(texts: list[str]):\n",
+    "def make_sparse_embedding(texts: list[str]) -> list[SparseEmbedding]:\n",
     "    return list(sparse_model.embed(texts, batch_size=32))\n",
     "\n",
     "\n",
@@ -615,7 +615,7 @@
     }
    ],
    "source": [
-    "def get_tokens_and_weights(sparse_embedding, model_name):\n",
+    "def get_tokens_and_weights(sparse_embedding, model_name) -> dict[str, float]:\n",
     "    # Find the tokenizer for the model\n",
     "    tokenizer_source = None\n",
     "    for model_info in SparseTextEmbedding.list_supported_models():\n",
@@ -626,7 +626,7 @@
     "            raise ValueError(f\"Model {model_name} not found in the supported models.\")\n",
     "\n",
     "    tokenizer = AutoTokenizer.from_pretrained(tokenizer_source)\n",
-    "    token_weight_dict = {}\n",
+    "    token_weight_dict: dict[str, float] = {}\n",
     "    for i in range(len(sparse_embedding.indices)):\n",
     "        token = tokenizer.decode([sparse_embedding.indices[i]])\n",
     "        weight = sparse_embedding.values[i]\n",

fastembed/common/model_management.py

@@ -255,7 +255,7 @@ def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int]]) -> Non
         return result
 
     @classmethod
-    def decompress_to_cache(cls, targz_path: str, cache_dir: str):
+    def decompress_to_cache(cls, targz_path: str, cache_dir: str) -> str:
         """
         Decompresses a .tar.gz file to a cache directory.
 

fastembed/image/onnx_image_model.py

@@ -24,7 +24,7 @@ def _get_worker_class(cls) -> Type["ImageEmbeddingWorker"]:
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def __init__(self) -> None:
+    def __init__(self):
         super().__init__()
         self.processor = None
 

fastembed/image/transform/functional.py

@@ -118,7 +118,7 @@ def rescale(image: np.ndarray, scale: float, dtype=np.float32) -> np.ndarray:
     return (image * scale).astype(dtype)
 
 
-def pil2ndarray(image: Union[Image.Image, np.ndarray]):
+def pil2ndarray(image: Union[Image.Image, np.ndarray]) -> np.ndarray:
     if isinstance(image, Image.Image):
         return np.asarray(image).transpose((2, 0, 1))
     return image

fastembed/image/transform/operators.py

@@ -133,11 +133,11 @@ def from_config(cls, config: dict[str, Any]) -> "Compose":
         return cls(transforms=transforms)
 
     @staticmethod
-    def _get_convert_to_rgb(transforms: list[Transform], config: dict[str, Any]):
+    def _get_convert_to_rgb(transforms: list[Transform], config: dict[str, Any]) -> None:
         transforms.append(ConvertToRGB())
 
     @classmethod
-    def _get_resize(cls, transforms: list[Transform], config: dict[str, Any]):
+    def _get_resize(cls, transforms: list[Transform], config: dict[str, Any]) -> None:
         mode = config.get("image_processor_type", "CLIPImageProcessor")
         if mode == "CLIPImageProcessor":
             if config.get("do_resize", False):
@@ -200,7 +200,7 @@ def _get_resize(cls, transforms: list[Transform], config: dict[str, Any]):
             raise ValueError(f"Preprocessor {mode} is not supported")
 
     @staticmethod
-    def _get_center_crop(transforms: list[Transform], config: dict[str, Any]):
+    def _get_center_crop(transforms: list[Transform], config: dict[str, Any]) -> None:
         mode = config.get("image_processor_type", "CLIPImageProcessor")
         if mode == "CLIPImageProcessor":
             if config.get("do_center_crop", False):
@@ -220,24 +220,24 @@ def _get_center_crop(transforms: list[Transform], config: dict[str, Any]):
             raise ValueError(f"Preprocessor {mode} is not supported")
 
     @staticmethod
-    def _get_pil2ndarray(transforms: list[Transform], config: dict[str, Any]):
+    def _get_pil2ndarray(transforms: list[Transform], config: dict[str, Any]) -> None:
         transforms.append(PILtoNDarray())
 
     @staticmethod
-    def _get_rescale(transforms: list[Transform], config: dict[str, Any]):
+    def _get_rescale(transforms: list[Transform], config: dict[str, Any]) -> None:
         if config.get("do_rescale", True):
             rescale_factor = config.get("rescale_factor", 1 / 255)
             transforms.append(Rescale(scale=rescale_factor))
 
     @staticmethod
-    def _get_normalize(transforms: list[Transform], config: dict[str, Any]):
+    def _get_normalize(transforms: list[Transform], config: dict[str, Any]) -> None:
         if config.get("do_normalize", False):
             transforms.append(Normalize(mean=config["image_mean"], std=config["image_std"]))
         elif "mean" in config and "std" in config:
             transforms.append(Normalize(mean=config["mean"], std=config["std"]))
 
     @staticmethod
-    def _get_pad2square(transforms: list[Transform], config: dict[str, Any]):
+    def _get_pad2square(transforms: list[Transform], config: dict[str, Any]) -> None:
         mode = config.get("image_processor_type", "CLIPImageProcessor")
         if mode == "CLIPImageProcessor":
             pass

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -1,9 +1,10 @@
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Iterable, Optional, Sequence, Type
+from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
+from numpy.typing import NDArray
 from tokenizers import Encoding
 
 from fastembed.common.onnx_model import (
@@ -46,7 +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):
+    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 = {
             "input_ids": np.array([enc.ids for enc in tokenized_input], dtype=np.int64),

fastembed/text/onnx_text_model.py

@@ -23,7 +23,7 @@ def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def __init__(self) -> None:
+    def __init__(self):
         super().__init__()
         self.tokenizer = None
         self.special_token_to_id = {}

fastembed/text/text_embedding_base.py

@@ -43,6 +43,7 @@ def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[np.ndar
         yield from self.embed(texts, **kwargs)
 
     def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
+
         """
         Embeds queries