fix: fix types

Author: joein

fastembed/common/utils.py

@@ -8,6 +8,7 @@
 from typing import Iterable, Optional, TypeVar
 
 import numpy as np
+from numpy.typing import NDArray
 
 from fastembed.common.types import NumpyArray
 
@@ -22,10 +23,9 @@ def normalize(input_array: NumpyArray, p: int = 2, dim: int = 1, eps: float = 1e
     return normalized_array
 
 
-def mean_pooling(input_array: NumpyArray, attention_mask: NumpyArray) -> NumpyArray:
-    input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
+def mean_pooling(input_array: NumpyArray, attention_mask: NDArray[np.int64]) -> NumpyArray:
+    input_mask_expanded = np.expand_dims(attention_mask, axis=-1).astype(np.int64)
     input_mask_expanded = np.tile(input_mask_expanded, (1, 1, input_array.shape[-1]))
-    input_mask_expanded = input_mask_expanded.astype(np.float32)
     sum_embeddings = np.sum(input_array * input_mask_expanded, axis=1)
     sum_mask = np.sum(input_mask_expanded, axis=1)
     pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)

fastembed/text/custom_text_embedding.py

@@ -1,6 +1,10 @@
 from typing import Optional, Sequence, Any, Iterable
 
 from dataclasses import dataclass
+
+import numpy as np
+from numpy.typing import NDArray
+
 from fastembed.common import OnnxProvider
 from fastembed.common.model_description import (
     PoolingType,
@@ -58,7 +62,7 @@ def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[Numpy
         return self._normalize(self._pool(output.model_output, output.attention_mask))
 
     def _pool(
-        self, embeddings: NumpyArray, attention_mask: Optional[NumpyArray] = None
+        self, embeddings: NumpyArray, attention_mask: Optional[NDArray[np.int64]] = None
     ) -> NumpyArray:
         if self._pooling == PoolingType.CLS:
             return embeddings[:, 0]

fastembed/text/pooled_embedding.py

@@ -1,6 +1,7 @@
 from typing import Any, Iterable, Type
 
 import numpy as np
+from numpy.typing import NDArray
 
 from fastembed.common.types import NumpyArray
 from fastembed.common.onnx_model import OnnxOutputContext
@@ -94,10 +95,12 @@ def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
         return PooledEmbeddingWorker
 
     @classmethod
-    def mean_pooling(cls, model_output: NumpyArray, attention_mask: NumpyArray) -> NumpyArray:
-        token_embeddings = model_output.astype(np.float32)
-        attention_mask = attention_mask.astype(np.float32)
-        return mean_pooling(token_embeddings, attention_mask)
+    def mean_pooling(
+        cls, model_output: NumpyArray, attention_mask: NDArray[np.int64]
+    ) -> NumpyArray:
+        # token_embeddings = model_output.astype(np.float32)
+        # attention_mask = attention_mask.astype(np.float32)
+        return mean_pooling(model_output, attention_mask)
 
     @classmethod
     def _list_supported_models(cls) -> list[DenseModelDescription]: