upstream-archive/fastembed
1
0
Fork 0
mirror of https://github.com/qdrant/fastembed.git synced 2026-07-23 11:20:51 -05:00
Code Issues Packages Projects Releases Wiki Activity

Mypy fixes

Browse Source
This commit is contained in:
d.rudenko
2024-12-09 22:34:38 +01:00
parent 05834259b3
commit be25bdf6eb
31 changed files with 231 additions and 208 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
fastembed/common/model_management.py
View File

@@ -3,11 +3,11 @@ import time
import shutil
import tarfile
from pathlib import Path
from typing import Any, Dict, List, Optional
from typing import Any, Dict, List, Optional, Self
import requests
from huggingface_hub import snapshot_download
from huggingface_hub.utils import RepositoryNotFoundError
from huggingface_hub import snapshot_download # type: ignore
from huggingface_hub.utils import RepositoryNotFoundError # type: ignore
from loguru import logger
from tqdm import tqdm
@@ -91,12 +91,12 @@ class ModelManagement:
@classmethod
def download_files_from_huggingface(
cls,
cls: type[Self],
hf_source_repo: str,
cache_dir: Optional[str] = None,
extra_patterns: Optional[List[str]] = None,
local_files_only: bool = False,
**kwargs,
**kwargs: Any,
) -> str:
"""
Downloads a model from HuggingFace Hub.
@@ -119,16 +119,16 @@ class ModelManagement:
if extra_patterns is not None:
allow_patterns.extend(extra_patterns)
return snapshot_download(
return str(snapshot_download(
repo_id=hf_source_repo,
allow_patterns=allow_patterns,
cache_dir=cache_dir,
local_files_only=local_files_only,
**kwargs,
)
))
@classmethod
def decompress_to_cache(cls, targz_path: str, cache_dir: str):
def decompress_to_cache(cls: type[Self], targz_path: str, cache_dir: str) -> str:
"""
Decompresses a .tar.gz file to a cache directory.
@@ -211,7 +211,7 @@ class ModelManagement:
@classmethod
def download_model(
cls, model: Dict[str, Any], cache_dir: Path, retries: int = 3, **kwargs
cls: type[Self], model: Dict[str, Any], cache_dir: Path, retries: int = 3, **kwargs: Any
) -> Path:
"""
Downloads a model from HuggingFace Hub or Google Cloud Storage.

29
fastembed/common/onnx_model.py
View File

@@ -11,12 +11,13 @@ from typing import (
Tuple,
Type,
TypeVar,
Self
)
import numpy as np
import onnxruntime as ort
from fastembed.common.types import OnnxProvider
from fastembed.common.types import OnnxProvider # type: ignore
from fastembed.parallel_processor import Worker
# Holds type of the embedding result
@@ -39,11 +40,11 @@ class OnnxModel(Generic[T]):
raise NotImplementedError("Subclasses must implement this method")
def __init__(self) -> None:
self.model = None
self.model: ort.InferenceSession = None
self.tokenizer = None
def _preprocess_onnx_input(
self, onnx_input: Dict[str, np.ndarray], **kwargs
self, onnx_input: Dict[str, np.ndarray], **kwargs: Any
) -> Dict[str, np.ndarray]:
"""
Preprocess the onnx input.
@@ -51,7 +52,7 @@ class OnnxModel(Generic[T]):
return onnx_input
def _load_onnx_model(
self,
self: Self,
model_dir: Path,
model_file: str,
threads: Optional[int],
@@ -106,30 +107,36 @@ class OnnxModel(Generic[T]):
def load_onnx_model(self) -> None:
raise NotImplementedError("Subclasses must implement this method")
def onnx_embed(self, *args, **kwargs) -> OnnxOutputContext:
def onnx_embed(self: Self, *args: Any, **kwargs: Any) -> OnnxOutputContext:
raise NotImplementedError("Subclasses must implement this method")
class EmbeddingWorker(Worker):
def init_embedding(
self,
self: Self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> OnnxModel:
raise NotImplementedError()
def __init__(
self,
self: Self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
):
self.model = self.init_embedding(model_name, cache_dir, **kwargs)
@classmethod
def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker":
def start(cls: type[Self], **kwargs: Any) -> "EmbeddingWorker":
model_name = kwargs.get("model_name", None)
if model_name is None:
raise ValueError("model_name must be provided")
cache_dir = kwargs.get("cache_dir", None)
if cache_dir is None:
raise ValueError("cache_dir must be provided")
return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
def process(self: Self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
raise NotImplementedError("Subclasses must implement this method")

2
fastembed/common/preprocessor_utils.py
View File

@@ -1,7 +1,7 @@
import json
from pathlib import Path
from typing import Tuple
from tokenizers import AddedToken, Tokenizer
from tokenizers import AddedToken, Tokenizer # type: ignore
from fastembed.image.transform.operators import Compose

4
fastembed/common/utils.py
View File

@@ -10,12 +10,12 @@ import re
from typing import Set
def normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
def normalize(input_array: np.ndarray, p: float=2, dim:int=1, eps:float=1e-12) -> np.ndarray:
# Calculate the Lp norm along the specified dimension
norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
norm = np.maximum(norm, eps) # Avoid division by zero
normalized_array = input_array / norm
return normalized_array
return normalized_array # type: ignore[no-any-return]
def iter_batch(iterable: Union[Iterable, Generator], size: int) -> Iterable:

10
fastembed/image/image_embedding.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Self
import numpy as np
@@ -41,7 +41,7 @@ class ImageEmbedding(ImageEmbeddingBase):
return result
def __init__(
self,
self: Self,
model_name: str,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
@@ -49,7 +49,7 @@ class ImageEmbedding(ImageEmbeddingBase):
cuda: bool = False,
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
**kwargs,
**kwargs: Any,
):
super().__init__(model_name, cache_dir, threads, **kwargs)
for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
@@ -73,11 +73,11 @@ class ImageEmbedding(ImageEmbeddingBase):
)
def embed(
self,
self: Self,
images: ImageInput,
batch_size: int = 16,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[np.ndarray]:
"""
Encode a list of documents into list of embeddings.

14
fastembed/image/image_embedding_base.py
View File

@@ -1,5 +1,5 @@
from typing import Iterable, Optional
from typing import Iterable, Optional, Self, Any
from pathlib import Path
import numpy as np
from fastembed.common.model_management import ModelManagement
@@ -8,11 +8,11 @@ from fastembed.common.types import ImageInput
class ImageEmbeddingBase(ModelManagement):
def __init__(
self,
self: Self,
model_name: str,
cache_dir: Optional[str] = None,
cache_dir: Optional[str|Path] = None,
threads: Optional[int] = None,
**kwargs,
**kwargs: Any,
):
self.model_name = model_name
self.cache_dir = cache_dir
@@ -20,11 +20,11 @@ class ImageEmbeddingBase(ModelManagement):
self._local_files_only = kwargs.pop("local_files_only", False)
def embed(
self,
self: Self,
images: ImageInput,
batch_size: int = 16,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[np.ndarray]:
"""
Embeds a list of images into a list of embeddings.

24
fastembed/image/onnx_embedding.py
View File

@@ -1,7 +1,7 @@
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Self
import numpy as np
from pathlib import Path
from fastembed.common import ImageInput, OnnxProvider
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common.utils import define_cache_dir, normalize
@@ -70,7 +70,7 @@ supported_onnx_models = [
class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
def __init__(
self,
self: Self,
model_name: str,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
@@ -79,12 +79,12 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
**kwargs: Any,
):
"""
Args:
model_name (str): The name of the model to use.
cache_dir (str, optional): The path to the cache directory.
cache_dir (Path, str, optional): The path to the cache directory.
Can be set using the `FASTEMBED_CACHE_PATH` env variable.
Defaults to `fastembed_cache` in the system's temp directory.
threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
@@ -115,11 +115,9 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
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: Path = define_cache_dir(cache_dir)
self._model_dir = self.download_model(
self.model_description, self.cache_dir, local_files_only=self._local_files_only
)
@@ -151,11 +149,11 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
return supported_onnx_models
def embed(
self,
self: Self,
images: ImageInput,
batch_size: int = 16,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[np.ndarray]:
"""
Encode a list of images into list of embeddings.
@@ -190,7 +188,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
return OnnxImageEmbeddingWorker
def _preprocess_onnx_input(
self, onnx_input: Dict[str, np.ndarray], **kwargs
self: Self, onnx_input: Dict[str, np.ndarray], **kwargs: Any
) -> Dict[str, np.ndarray]:
"""
Preprocess the onnx input.
@@ -198,12 +196,12 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
return onnx_input
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
def _post_process_onnx_output(self: Self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
return normalize(output.model_output).astype(np.float32)
class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):
def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxImageEmbedding:
def init_embedding(self: Self, model_name: str, cache_dir: str, **kwargs: Any) -> OnnxImageEmbedding:
return OnnxImageEmbedding(
model_name=model_name,
cache_dir=cache_dir,

16
fastembed/image/onnx_image_model.py
View File

@@ -2,7 +2,7 @@ import contextlib
import os
from multiprocessing import get_all_start_methods
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type, Self
import numpy as np
from PIL import Image
@@ -29,7 +29,7 @@ class OnnxImageModel(OnnxModel[T]):
self.processor = None
def _preprocess_onnx_input(
self, onnx_input: Dict[str, np.ndarray], **kwargs
self: Self, onnx_input: Dict[str, np.ndarray], **kwargs: Any
) -> Dict[str, np.ndarray]:
"""
Preprocess the onnx input.
@@ -37,7 +37,7 @@ class OnnxImageModel(OnnxModel[T]):
return onnx_input
def _load_onnx_model(
self,
self: Self,
model_dir: Path,
model_file: str,
threads: Optional[int],
@@ -58,10 +58,10 @@ class OnnxImageModel(OnnxModel[T]):
def load_onnx_model(self) -> None:
raise NotImplementedError("Subclasses must implement this method")
def _build_onnx_input(self, encoded: np.ndarray) -> Dict[str, np.ndarray]:
def _build_onnx_input(self: Self, encoded: np.ndarray) -> Dict[str, np.ndarray]:
return {node.name: encoded for node in self.model.get_inputs()}
def onnx_embed(self, images: List[ImageInput], **kwargs) -> OnnxOutputContext:
def onnx_embed(self: Self, images: List[ImageInput], **kwargs: Any) -> OnnxOutputContext:
with contextlib.ExitStack():
image_files = [
Image.open(image) if not isinstance(image, Image.Image) else image
@@ -75,7 +75,7 @@ class OnnxImageModel(OnnxModel[T]):
return OnnxOutputContext(model_output=embeddings)
def _embed_images(
self,
self: Self,
model_name: str,
cache_dir: str,
images: ImageInput,
@@ -84,7 +84,7 @@ class OnnxImageModel(OnnxModel[T]):
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[T]:
is_small = False
@@ -125,7 +125,7 @@ class OnnxImageModel(OnnxModel[T]):
class ImageEmbeddingWorker(EmbeddingWorker):
def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
def process(self: Self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
for idx, batch in items:
embeddings = self.model.onnx_embed(batch)
yield idx, embeddings

21
fastembed/image/transform/functional.py
View File

@@ -62,8 +62,8 @@ def center_crop(
def normalize(
image: np.ndarray,
mean=Union[float, np.ndarray],
std=Union[float, np.ndarray],
mean:Union[float, np.ndarray],
std:Union[float, np.ndarray],
) -> np.ndarray:
if not isinstance(image, np.ndarray):
raise ValueError("image must be a numpy array")
@@ -79,7 +79,7 @@ def normalize(
f"mean must have {num_channels} elements if it is an iterable, got {len(mean)}"
)
else:
mean = [mean] * num_channels
mean = np.array([mean] * num_channels)
mean = np.array(mean, dtype=image.dtype)
if isinstance(std, Sized):
@@ -88,7 +88,7 @@ def normalize(
f"std must have {num_channels} elements if it is an iterable, got {len(std)}"
)
else:
std = [std] * num_channels
std = np.array([std] * num_channels)
std = np.array(std, dtype=image.dtype)
image = ((image.T - mean) / std).T
@@ -96,10 +96,10 @@ def normalize(
def resize(
image: Image,
image: Image.Image,
size: Union[int, Tuple[int, int]],
resample: Image.Resampling = Image.Resampling.BILINEAR,
) -> Image:
) -> Image.Image:
if isinstance(size, tuple):
return image.resize(size, resample)
@@ -114,11 +114,14 @@ def resize(
return image.resize(new_size, resample)
def rescale(image: np.ndarray, scale: float, dtype=np.float32) -> np.ndarray:
def rescale(
image: np.ndarray[Union[np.float32, np.int32]],
scale: float,
dtype: Union[np.dtype[np.float32], np.dtype[np.int32]] = 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

18
fastembed/image/transform/operators.py
View File

@@ -33,8 +33,8 @@ class CenterCrop(Transform):
class Normalize(Transform):
def __init__(self, mean: Union[float, List[float]], std: Union[float, List[float]]):
self.mean = mean
self.std = std
self.mean = np.array(mean)
self.std = np.array(std)
def __call__(self, images: List[np.ndarray]) -> List[np.ndarray]:
return [normalize(image, mean=self.mean, std=self.std) for image in images]
@@ -100,7 +100,7 @@ class Compose:
Returns:
Compose: Image processor.
"""
transforms = []
transforms: List[Transform] = []
cls._get_convert_to_rgb(transforms, config)
cls._get_resize(transforms, config)
cls._get_center_crop(transforms, config)
@@ -110,11 +110,11 @@ class 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())
@staticmethod
def _get_resize(transforms: List[Transform], config: Dict[str, Any]):
def _get_resize(transforms: List[Transform], config: Dict[str, Any]) -> None:
mode = config.get("image_processor_type", "CLIPImageProcessor")
if mode == "CLIPImageProcessor" or mode == "SiglipImageProcessor":
if config.get("do_resize", False):
@@ -159,7 +159,7 @@ class Compose:
)
@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" or mode == "SiglipImageProcessor":
if config.get("do_center_crop", False):
@@ -177,16 +177,16 @@ class Compose:
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"]))

30
fastembed/late_interaction/colbert.py
View File

@@ -1,5 +1,5 @@
import string
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union, Self
import numpy as np
from tokenizers import Encoding
@@ -45,7 +45,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
MASK_TOKEN = "[MASK]"
def _post_process_onnx_output(
self, output: OnnxOutputContext, is_doc: bool = True
self: Self, output: OnnxOutputContext, is_doc: bool = True
) -> Iterable[np.ndarray]:
if not is_doc:
return output.model_output.astype(np.float32)
@@ -67,7 +67,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
return output.model_output.astype(np.float32)
def _preprocess_onnx_input(
self, onnx_input: Dict[str, np.ndarray], is_doc: bool = True
self: Self, onnx_input: Dict[str, np.ndarray], is_doc: bool = True
) -> Dict[str, np.ndarray]:
if is_doc:
onnx_input["input_ids"][:, 1] = self.DOCUMENT_MARKER_TOKEN_ID
@@ -75,14 +75,14 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
onnx_input["input_ids"][:, 1] = self.QUERY_MARKER_TOKEN_ID
return onnx_input
def tokenize(self, documents: List[str], is_doc: bool = True) -> List[Encoding]:
def tokenize(self: Self, documents: List[str], is_doc: bool = True) -> List[Encoding]:
return (
self._tokenize_documents(documents=documents)
if is_doc
else self._tokenize_query(query=next(iter(documents)))
)
def _tokenize_query(self, query: str) -> List[Encoding]:
def _tokenize_query(self: Self, query: str) -> List[Encoding]:
# "@ " is added to a query to be replaced with a special query token
# make sure that "@ " is considered as a single token
query = f"@ {query}"
@@ -104,7 +104,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
self.tokenizer.enable_padding(**prev_padding)
return encoded
def _tokenize_documents(self, documents: List[str]) -> List[Encoding]:
def _tokenize_documents(self: Self, documents: List[str]) -> List[Encoding]:
# "@ " is added to a document to be replaced with a special document token
# make sure that "@ " is considered as a single token
documents = ["@ " + doc for doc in documents]
@@ -112,7 +112,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
return encoded
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls: type[Self]) -> List[Dict[str, Any]]:
"""Lists the supported models.
Returns:
@@ -121,7 +121,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
return supported_colbert_models
def __init__(
self,
self: Self,
model_name: str,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
@@ -130,7 +130,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
**kwargs: Any,
):
"""
Args:
@@ -182,7 +182,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
if not self.lazy_load:
self.load_onnx_model()
def load_onnx_model(self) -> None:
def load_onnx_model(self: Self) -> None:
self._load_onnx_model(
model_dir=self._model_dir,
model_file=self.model_description["model_file"],
@@ -199,11 +199,11 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
}
def embed(
self,
self: Self,
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[np.ndarray]:
"""
Encode a list of documents into list of embeddings.
@@ -232,7 +232,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
**kwargs,
)
def query_embed(self, query: Union[str, List[str]], **kwargs) -> Iterable[np.ndarray]:
def query_embed(self: Self, query: Union[str, List[str]], **kwargs: Any) -> Iterable[np.ndarray]:
if isinstance(query, str):
query = [query]
@@ -245,12 +245,12 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
)
@classmethod
def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
def _get_worker_class(cls: type[Self]) -> Type[TextEmbeddingWorker]:
return ColbertEmbeddingWorker
class ColbertEmbeddingWorker(TextEmbeddingWorker):
def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> Colbert:
def init_embedding(self: Self, model_name: str, cache_dir: str, **kwargs: Any) -> Colbert:
return Colbert(
model_name=model_name,
cache_dir=cache_dir,

4
fastembed/late_interaction/jina_colbert.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, List, Type
from typing import Any, Dict, List, Type, Self
import numpy as np
@@ -53,7 +53,7 @@ class JinaColbert(Colbert):
class JinaColbertEmbeddingWorker(TextEmbeddingWorker):
def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> JinaColbert:
def init_embedding(self: Self, model_name: str, cache_dir: str, **kwargs: Any) -> JinaColbert:
return JinaColbert(
model_name=model_name,
cache_dir=cache_dir,

14
fastembed/late_interaction/late_interaction_embedding_base.py
View File

@@ -1,4 +1,4 @@
from typing import Iterable, Optional, Union
from typing import Iterable, Optional, Union, Any, Self
import numpy as np
@@ -7,11 +7,11 @@ from fastembed.common.model_management import ModelManagement
class LateInteractionTextEmbeddingBase(ModelManagement):
def __init__(
self,
self: Self,
model_name: str,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
**kwargs,
**kwargs: Any,
):
self.model_name = model_name
self.cache_dir = cache_dir
@@ -19,15 +19,15 @@ class LateInteractionTextEmbeddingBase(ModelManagement):
self._local_files_only = kwargs.pop("local_files_only", False)
def embed(
self,
self: Self,
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[np.ndarray]:
raise NotImplementedError()
def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[np.ndarray]:
def passage_embed(self: Self, texts: Iterable[str], **kwargs: Any) -> Iterable[np.ndarray]:
"""
Embeds a list of text passages into a list of embeddings.
@@ -43,7 +43,7 @@ class LateInteractionTextEmbeddingBase(ModelManagement):
yield from self.embed(texts, **kwargs)
def query_embed(
self, query: Union[str, Iterable[str]], **kwargs
self: Self, query: Union[str, Iterable[str]], **kwargs: Any
) -> Iterable[np.ndarray]:
"""
Embeds queries

14
fastembed/late_interaction/late_interaction_text_embedding.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union, Self
import numpy as np
@@ -14,7 +14,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
EMBEDDINGS_REGISTRY: List[Type[LateInteractionTextEmbeddingBase]] = [Colbert, JinaColbert]
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls: type[Self]) -> List[Dict[str, Any]]:
"""
Lists the supported models.
@@ -44,7 +44,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
return result
def __init__(
self,
self: Self,
model_name: str,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
@@ -52,7 +52,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
cuda: bool = False,
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
**kwargs,
**kwargs: Any,
):
super().__init__(model_name, cache_dir, threads, **kwargs)
for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
@@ -76,11 +76,11 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
)
def embed(
self,
self: Self,
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[np.ndarray]:
"""
Encode a list of documents into list of embeddings.
@@ -99,7 +99,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
"""
yield from self.model.embed(documents, batch_size, parallel, **kwargs)
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
def query_embed(self: Self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
"""
Embeds queries

36
fastembed/sparse/bm25.py
View File

@@ -2,7 +2,7 @@ import os
from collections import defaultdict
from multiprocessing import get_all_start_methods
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Union
from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Union, Self
import mmh3
import numpy as np
@@ -97,7 +97,7 @@ class Bm25(SparseTextEmbeddingBase):
"""
def __init__(
self,
self: Self,
model_name: str,
cache_dir: Optional[str] = None,
k: float = 1.2,
@@ -105,7 +105,7 @@ class Bm25(SparseTextEmbeddingBase):
avg_len: float = 256.0,
language: str = "english",
token_max_length: int = 40,
**kwargs,
**kwargs: Any,
):
super().__init__(model_name, cache_dir, **kwargs)
@@ -133,7 +133,7 @@ class Bm25(SparseTextEmbeddingBase):
self.tokenizer = SimpleTokenizer
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls: type[Self]) -> List[Dict[str, Any]]:
"""Lists the supported models.
Returns:
@@ -142,7 +142,7 @@ class Bm25(SparseTextEmbeddingBase):
return supported_bm25_models
@classmethod
def _load_stopwords(cls, model_dir: Path, language: str) -> List[str]:
def _load_stopwords(cls: type[Self], model_dir: Path, language: str) -> List[str]:
stopwords_path = model_dir / f"{language}.txt"
if not stopwords_path.exists():
return []
@@ -151,7 +151,7 @@ class Bm25(SparseTextEmbeddingBase):
return f.read().splitlines()
def _embed_documents(
self,
self: Self,
model_name: str,
cache_dir: str,
documents: Union[str, Iterable[str]],
@@ -193,11 +193,11 @@ class Bm25(SparseTextEmbeddingBase):
yield record
def embed(
self,
self: Self,
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[SparseEmbedding]:
"""
Encode a list of documents into list of embeddings.
@@ -222,7 +222,7 @@ class Bm25(SparseTextEmbeddingBase):
parallel=parallel,
)
def _stem(self, tokens: List[str]) -> List[str]:
def _stem(self: Self, tokens: List[str]) -> List[str]:
stemmed_tokens = []
for token in tokens:
if token in self.punctuation:
@@ -241,7 +241,7 @@ class Bm25(SparseTextEmbeddingBase):
return stemmed_tokens
def raw_embed(
self,
self: Self,
documents: List[str],
) -> List[SparseEmbedding]:
embeddings = []
@@ -253,7 +253,7 @@ class Bm25(SparseTextEmbeddingBase):
embeddings.append(SparseEmbedding.from_dict(token_id2value))
return embeddings
def _term_frequency(self, tokens: List[str]) -> Dict[int, float]:
def _term_frequency(self: Self, tokens: List[str]) -> Dict[int, float]:
"""Calculate the term frequency part of the BM25 formula.
(
@@ -284,10 +284,10 @@ class Bm25(SparseTextEmbeddingBase):
return tf_map
@classmethod
def compute_token_id(cls, token: str) -> int:
def compute_token_id(cls: type[Self], token: str) -> int:
return abs(mmh3.hash(token))
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]:
def query_embed(self: Self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[SparseEmbedding]:
"""To emulate BM25 behaviour, we don't need to use weights in the query, and
it's enough to just hash the tokens and assign a weight of 1.0 to them.
"""
@@ -312,22 +312,22 @@ class Bm25(SparseTextEmbeddingBase):
class Bm25Worker(Worker):
def __init__(
self,
self: Self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
):
self.model = self.init_embedding(model_name, cache_dir, **kwargs)
@classmethod
def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "Bm25Worker":
def start(cls: type[Self], model_name: str, cache_dir: str, **kwargs: Any) -> "Bm25Worker":
return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
def process(self: Self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
for idx, batch in items:
onnx_output = self.model.raw_embed(batch)
yield idx, onnx_output
@staticmethod
def init_embedding(model_name: str, cache_dir: str, **kwargs) -> Bm25:
def init_embedding(model_name: str, cache_dir: str, **kwargs: Any) -> Bm25:
return Bm25(model_name=model_name, cache_dir=cache_dir, **kwargs)

24
fastembed/sparse/bm42.py
View File

@@ -1,7 +1,7 @@
import math
import string
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type, Union
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type, Union, Self
import mmh3
import numpy as np
@@ -56,7 +56,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
ONNX_OUTPUT_NAMES = ["attention_6"]
def __init__(
self,
self: Self,
model_name: str,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
@@ -66,7 +66,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
**kwargs: Any,
):
"""
Args:
@@ -141,7 +141,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
self.special_tokens_ids = set(self.special_token_to_id.values())
self.stopwords = set(self._load_stopwords(self._model_dir))
def _filter_pair_tokens(self, tokens: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
def _filter_pair_tokens(self: Self, tokens: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
result = []
for token, value in tokens:
if token in self.stopwords or token in self.punctuation:
@@ -149,7 +149,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
result.append((token, value))
return result
def _stem_pair_tokens(self, tokens: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
def _stem_pair_tokens(self: Self, tokens: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
result = []
for token, value in tokens:
processed_token = self.stemmer.stem_word(token)
@@ -167,7 +167,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
return result
def _reconstruct_bpe(
self, bpe_tokens: Iterable[Tuple[int, str]]
self: Self, bpe_tokens: Iterable[Tuple[int, str]]
) -> List[Tuple[str, List[int]]]:
result = []
acc = ""
@@ -195,7 +195,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
return result
def _rescore_vector(self, vector: Dict[str, float]) -> Dict[int, float]:
def _rescore_vector(self: Self, vector: Dict[str, float]) -> Dict[int, float]:
"""
Orders all tokens in the vector by their importance and generates a new score based on the importance order.
So that the scoring doesn't depend on absolute values assigned by the model, but on the relative importance.
@@ -213,7 +213,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
return new_vector
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[SparseEmbedding]:
def _post_process_onnx_output(self: Self, output: OnnxOutputContext) -> Iterable[SparseEmbedding]:
if output.input_ids is None:
raise ValueError("input_ids must be provided for document post-processing")
@@ -264,11 +264,11 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
return f.read().splitlines()
def embed(
self,
self: Self,
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[SparseEmbedding]:
"""
Encode a list of documents into list of embeddings.
@@ -305,7 +305,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
result[token_id] = 1.0
return result
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]:
def query_embed(self: Self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[SparseEmbedding]:
"""
To emulate BM25 behaviour, we don't need to use smart weights in the query, and
it's enough to just hash the tokens and assign a weight of 1.0 to them.
@@ -332,7 +332,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
class Bm42TextEmbeddingWorker(TextEmbeddingWorker):
def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> Bm42:
def init_embedding(self: Self, model_name: str, cache_dir: str, **kwargs: Any) -> Bm42:
return Bm42(
model_name=model_name,
cache_dir=cache_dir,

16
fastembed/sparse/sparse_embedding_base.py
View File

@@ -1,5 +1,5 @@
from dataclasses import dataclass
from typing import Dict, Iterable, Optional, Union
from typing import Dict, Iterable, Optional, Union, Self, Any
import numpy as np
@@ -21,7 +21,7 @@ class SparseEmbedding:
return {i: v for i, v in zip(self.indices, self.values)}
@classmethod
def from_dict(cls, data: Dict[int, float]) -> "SparseEmbedding":
def from_dict(cls: type[Self], data: Dict[int, float]) -> "SparseEmbedding":
if len(data) == 0:
return cls(values=np.array([]), indices=np.array([]))
indices, values = zip(*data.items())
@@ -30,11 +30,11 @@ class SparseEmbedding:
class SparseTextEmbeddingBase(ModelManagement):
def __init__(
self,
self: Self,
model_name: str,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
**kwargs,
**kwargs: Any,
):
self.model_name = model_name
self.cache_dir = cache_dir
@@ -42,16 +42,16 @@ class SparseTextEmbeddingBase(ModelManagement):
self._local_files_only = kwargs.pop("local_files_only", False)
def embed(
self,
self: Self,
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[SparseEmbedding]:
raise NotImplementedError()
def passage_embed(
self, texts: Iterable[str], **kwargs
self: Self, texts: Iterable[str], **kwargs: Any
) -> Iterable[SparseEmbedding]:
"""
Embeds a list of text passages into a list of embeddings.
@@ -68,7 +68,7 @@ class SparseTextEmbeddingBase(ModelManagement):
yield from self.embed(texts, **kwargs)
def query_embed(
self, query: Union[str, Iterable[str]], **kwargs
self: Self, query: Union[str, Iterable[str]], **kwargs: Any
) -> Iterable[SparseEmbedding]:
"""
Embeds queries

14
fastembed/sparse/sparse_text_embedding.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union, Self
from fastembed.common import OnnxProvider
from fastembed.sparse.bm25 import Bm25
@@ -15,7 +15,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
EMBEDDINGS_REGISTRY: List[Type[SparseTextEmbeddingBase]] = [SpladePP, Bm42, Bm25]
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls: type[Self]) -> List[Dict[str, Any]]:
"""
Lists the supported models.
@@ -44,7 +44,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
return result
def __init__(
self,
self: Self,
model_name: str,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
@@ -52,7 +52,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
cuda: bool = False,
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
**kwargs,
**kwargs: Any,
):
super().__init__(model_name, cache_dir, threads, **kwargs)
if model_name == "prithvida/Splade_PP_en_v1":
@@ -85,11 +85,11 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
)
def embed(
self,
self: Self,
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[SparseEmbedding]:
"""
Encode a list of documents into list of embeddings.
@@ -108,7 +108,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
"""
yield from self.model.embed(documents, batch_size, parallel, **kwargs)
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]:
def query_embed(self: Self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[SparseEmbedding]:
"""
Embeds queries

12
fastembed/sparse/splade_pp.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union, Self
import numpy as np
from fastembed.common import OnnxProvider
@@ -64,7 +64,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
return supported_splade_models
def __init__(
self,
self: Self,
model_name: str,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
@@ -73,7 +73,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
**kwargs: Any,
):
"""
Args:
@@ -132,11 +132,11 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
)
def embed(
self,
self: Self,
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[SparseEmbedding]:
"""
Encode a list of documents into list of embeddings.
@@ -171,7 +171,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
class SpladePPEmbeddingWorker(TextEmbeddingWorker):
def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> SpladePP:
def init_embedding(self: Self, model_name: str, cache_dir: str, **kwargs: Any) -> SpladePP:
return SpladePP(
model_name=model_name,
cache_dir=cache_dir,

1
fastembed/sparse/utils/tokenizer.py
View File

@@ -5,6 +5,7 @@ from typing import List
class SimpleTokenizer:
@staticmethod
def tokenize(text: str) -> List[str]:
text = re.sub(r"[^\w]", " ", text.lower())
text = re.sub(r"\s+", " ", text)

6
fastembed/text/clip_embedding.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, Iterable, List, Type
from typing import Any, Dict, Iterable, List, Type, Self
import numpy as np
@@ -41,10 +41,10 @@ class CLIPOnnxEmbedding(OnnxTextEmbedding):
class CLIPEmbeddingWorker(OnnxTextEmbeddingWorker):
def init_embedding(
self,
self: Self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> OnnxTextEmbedding:
return CLIPOnnxEmbedding(
model_name=model_name,

8
fastembed/text/e5_onnx_embedding.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, List, Type
from typing import Any, Dict, List, Type, Self
import numpy as np
@@ -48,7 +48,7 @@ class E5OnnxEmbedding(OnnxTextEmbedding):
return supported_multilingual_e5_models
def _preprocess_onnx_input(
self, onnx_input: Dict[str, np.ndarray], **kwargs
self: Self, onnx_input: Dict[str, np.ndarray], **kwargs: Any
) -> Dict[str, np.ndarray]:
"""
Preprocess the onnx input.
@@ -59,10 +59,10 @@ class E5OnnxEmbedding(OnnxTextEmbedding):
class E5OnnxEmbeddingWorker(OnnxTextEmbeddingWorker):
def init_embedding(
self,
self: Self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> E5OnnxEmbedding:
return E5OnnxEmbedding(
model_name=model_name,

26
fastembed/text/onnx_embedding.py
View File

@@ -1,7 +1,7 @@
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union, Self
import numpy as np
from pathlib import Path
from fastembed.common import OnnxProvider
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common.utils import define_cache_dir, normalize
@@ -193,16 +193,16 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
return supported_onnx_models
def __init__(
self,
self: Self,
model_name: str = "BAAI/bge-small-en-v1.5",
cache_dir: Optional[str] = None,
cache_dir: Optional[str|Path] = None,
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
**kwargs: Any,
):
"""
Args:
@@ -237,11 +237,9 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
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: Path = define_cache_dir(cache_dir)
self._model_dir = self.download_model(
self.model_description, self.cache_dir, local_files_only=self._local_files_only
)
@@ -250,11 +248,11 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
self.load_onnx_model()
def embed(
self,
self: Self,
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[np.ndarray]:
"""
Encode a list of documents into list of embeddings.
@@ -288,14 +286,14 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
return OnnxTextEmbeddingWorker
def _preprocess_onnx_input(
self, onnx_input: Dict[str, np.ndarray], **kwargs
self: Self, onnx_input: Dict[str, np.ndarray], **kwargs: Any
) -> Dict[str, np.ndarray]:
"""
Preprocess the onnx input.
"""
return onnx_input
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
def _post_process_onnx_output(self: Self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
embeddings = output.model_output
return normalize(embeddings[:, 0]).astype(np.float32)
@@ -312,10 +310,10 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
def init_embedding(
self,
self: Self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> OnnxTextEmbedding:
return OnnxTextEmbedding(
model_name=model_name,

24
fastembed/text/onnx_text_model.py
View File

@@ -1,10 +1,10 @@
import os
from multiprocessing import get_all_start_methods
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type, Union
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type, Union, Self
import numpy as np
from tokenizers import Encoding
from tokenizers import Encoding # type: ignore
from fastembed.common import OnnxProvider
from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
@@ -26,10 +26,10 @@ class OnnxTextModel(OnnxModel[T]):
def __init__(self) -> None:
super().__init__()
self.tokenizer = None
self.special_token_to_id = {}
self.special_token_to_id: dict[str, int] = {}
def _preprocess_onnx_input(
self, onnx_input: Dict[str, np.ndarray], **kwargs
self: Self, onnx_input: Dict[str, np.ndarray], **kwargs: Any
) -> Dict[str, np.ndarray]:
"""
Preprocess the onnx input.
@@ -53,18 +53,18 @@ class OnnxTextModel(OnnxModel[T]):
cuda=cuda,
device_id=device_id,
)
self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir) # type: ignore[no-any-return]
def load_onnx_model(self) -> None:
raise NotImplementedError("Subclasses must implement this method")
def tokenize(self, documents: List[str], **kwargs) -> List[Encoding]:
return self.tokenizer.encode_batch(documents)
def tokenize(self: Self, documents: List[str], **kwargs: Any) -> List[Encoding]:
return self.tokenizer.encode_batch(documents) # type: ignore[no-any-return]
def onnx_embed(
self,
self: Self,
documents: List[str],
**kwargs,
**kwargs: Any,
) -> OnnxOutputContext:
encoded = self.tokenize(documents, **kwargs)
input_ids = np.array([e.ids for e in encoded])
@@ -94,7 +94,7 @@ class OnnxTextModel(OnnxModel[T]):
)
def _embed_documents(
self,
self: Self,
model_name: str,
cache_dir: str,
documents: Union[str, Iterable[str]],
@@ -103,7 +103,7 @@ class OnnxTextModel(OnnxModel[T]):
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[T]:
is_small = False
@@ -144,7 +144,7 @@ class OnnxTextModel(OnnxModel[T]):
class TextEmbeddingWorker(EmbeddingWorker):
def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
def process(self: Self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
for idx, batch in items:
onnx_output = self.model.onnx_embed(batch)
yield idx, onnx_output

8
fastembed/text/pooled_embedding.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, Iterable, List, Type
from typing import Any, Dict, Iterable, List, Type, Self
import numpy as np
@@ -57,7 +57,7 @@ class PooledEmbedding(OnnxTextEmbedding):
sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
sum_mask = np.sum(input_mask_expanded, axis=1)
pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
return pooled_embeddings
return pooled_embeddings # type: ignore[no-any-return]
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
@@ -79,10 +79,10 @@ class PooledEmbedding(OnnxTextEmbedding):
class PooledEmbeddingWorker(OnnxTextEmbeddingWorker):
def init_embedding(
self,
self: Self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> OnnxTextEmbedding:
return PooledEmbedding(
model_name=model_name,

6
fastembed/text/pooled_normalized_embedding.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, Iterable, List, Type
from typing import Any, Dict, Iterable, List, Type, Self
import numpy as np
@@ -85,10 +85,10 @@ class PooledNormalizedEmbedding(PooledEmbedding):
class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker):
def init_embedding(
self,
self: Self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> OnnxTextEmbedding:
return PooledNormalizedEmbedding(
model_name=model_name,

0
fastembed/text/pooling_embedding.py Normal file
View File

10
fastembed/text/text_embedding.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union, Self
import numpy as np
@@ -51,7 +51,7 @@ class TextEmbedding(TextEmbeddingBase):
return result
def __init__(
self,
self: Self,
model_name: str = "BAAI/bge-small-en-v1.5",
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
@@ -59,7 +59,7 @@ class TextEmbedding(TextEmbeddingBase):
cuda: bool = False,
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
**kwargs,
**kwargs: Any,
):
super().__init__(model_name, cache_dir, threads, **kwargs)
for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
@@ -83,11 +83,11 @@ class TextEmbedding(TextEmbeddingBase):
)
def embed(
self,
self: Self,
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[np.ndarray]:
"""
Encode a list of documents into list of embeddings.

14
fastembed/text/text_embedding_base.py
View File

@@ -1,4 +1,4 @@
from typing import Iterable, Optional, Union
from typing import Iterable, Optional, Union, Self, Any
import numpy as np
@@ -7,11 +7,11 @@ from fastembed.common.model_management import ModelManagement
class TextEmbeddingBase(ModelManagement):
def __init__(
self,
self: Self,
model_name: str,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
**kwargs,
**kwargs: Any,
):
self.model_name = model_name
self.cache_dir = cache_dir
@@ -19,15 +19,15 @@ class TextEmbeddingBase(ModelManagement):
self._local_files_only = kwargs.pop("local_files_only", False)
def embed(
self,
self: Self,
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[np.ndarray]:
raise NotImplementedError()
def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[np.ndarray]:
def passage_embed(self: Self, texts: Iterable[str], **kwargs: Any) -> Iterable[np.ndarray]:
"""
Embeds a list of text passages into a list of embeddings.
@@ -43,7 +43,7 @@ class TextEmbeddingBase(ModelManagement):
yield from self.embed(texts, **kwargs)
def query_embed(
self, query: Union[str, Iterable[str]], **kwargs
self: Self, query: Union[str, Iterable[str]], **kwargs: Any
) -> Iterable[np.ndarray]:
"""
Embeds queries

10
pyproject.toml
View File

@@ -32,6 +32,8 @@ pytest = "^7.4.2"
ruff = ">=0.3.1,<1.0"
notebook = ">=7.0.2"
pre-commit = {version = "^3.6.2", python = ">=3.9,<3.12" }
mypy = "^1.13.0"
pyright = "^1.1.390"
[tool.poetry.group.docs.dependencies]
mkdocs-material = "^9.5.10"
@@ -47,3 +49,11 @@ build-backend = "poetry.core.masonry.api"
[tool.ruff]
line-length = 99
[tool.mypy]
plugins = []
strict = true
disable_error_code = ["type-arg"]

6
pyrightconfig.json Normal file
View File

@@ -0,0 +1,6 @@
{
"include": ["fastembed/"],
"exclude": ["tests/", "docs/", "experiments/"],
"strict": true,
"venvPath": ".venv"
}