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

chore: Add any to kwargs (#450)

Browse Source 
* new: Add mypy and pyright deps

* chore: Add any to kwargs

* chore: Add any to args

* add missing kwargs

---------

Co-authored-by: George Panchuk <george.panchuk@qdrant.tech>
This commit is contained in:
Hossam Hagag
2025-01-28 23:49:37 +02:00
committed by GitHub
parent 314842121d
commit bb815405aa
28 changed files with 93 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
fastembed/common/model_management.py
View File

@@ -104,7 +104,7 @@ class ModelManagement:
cache_dir: str,
extra_patterns: list[str],
local_files_only: bool = False,
**kwargs,
**kwargs: Any,
) -> str:
"""
Downloads a model from HuggingFace Hub.
@@ -338,7 +338,7 @@ class ModelManagement:
@classmethod
def download_model(
cls, model: dict[str, Any], cache_dir: Path, retries: int = 3, **kwargs
cls, model: dict[str, Any], cache_dir: Path, retries: int = 3, **kwargs: Any
) -> Path:
"""
Downloads a model from HuggingFace Hub or Google Cloud Storage.

8
fastembed/common/onnx_model.py
View File

@@ -33,7 +33,7 @@ class OnnxModel(Generic[T]):
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.
@@ -103,7 +103,7 @@ 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, *args: Any, **kwargs: Any) -> OnnxOutputContext:
raise NotImplementedError("Subclasses must implement this method")
@@ -112,7 +112,7 @@ class EmbeddingWorker(Worker):
self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> OnnxModel:
raise NotImplementedError()
@@ -120,7 +120,7 @@ class EmbeddingWorker(Worker):
self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
):
self.model = self.init_embedding(model_name, cache_dir, **kwargs)

4
fastembed/embedding.py
View File

@@ -1,4 +1,4 @@
from typing import Optional
from typing import Optional, Any
from loguru import logger
@@ -19,6 +19,6 @@ class JinaEmbedding(TextEmbedding):
model_name: str = "jinaai/jina-embeddings-v2-base-en",
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
**kwargs,
**kwargs: Any,
):
super().__init__(model_name, cache_dir, threads, **kwargs)

4
fastembed/image/image_embedding.py
View File

@@ -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:
@@ -77,7 +77,7 @@ class ImageEmbedding(ImageEmbeddingBase):
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.

6
fastembed/image/image_embedding_base.py
View File

@@ -1,4 +1,4 @@
from typing import Iterable, Optional
from typing import Iterable, Optional, Any
import numpy as np
@@ -12,7 +12,7 @@ class ImageEmbeddingBase(ModelManagement):
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
@@ -24,7 +24,7 @@ class ImageEmbeddingBase(ModelManagement):
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.

8
fastembed/image/onnx_embedding.py
View File

@@ -78,7 +78,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
lazy_load: bool = False,
device_id: Optional[int] = None,
specific_model_path: Optional[str] = None,
**kwargs,
**kwargs: Any,
):
"""
Args:
@@ -158,7 +158,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
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.
@@ -193,7 +193,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
return OnnxImageEmbeddingWorker
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.
@@ -206,7 +206,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):
def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxImageEmbedding:
def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> OnnxImageEmbedding:
return OnnxImageEmbedding(
model_name=model_name,
cache_dir=cache_dir,

6
fastembed/image/onnx_image_model.py
View File

@@ -29,7 +29,7 @@ class OnnxImageModel(OnnxModel[T]):
self.processor = 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.
@@ -61,7 +61,7 @@ class OnnxImageModel(OnnxModel[T]):
def _build_onnx_input(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, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
with contextlib.ExitStack():
image_files = [
Image.open(image) if not isinstance(image, Image.Image) else image
@@ -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

8
fastembed/late_interaction/colbert.py
View File

@@ -125,7 +125,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
lazy_load: bool = False,
device_id: Optional[int] = None,
specific_model_path: Optional[str] = None,
**kwargs,
**kwargs: Any,
):
"""
Args:
@@ -205,7 +205,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
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.
@@ -234,7 +234,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
**kwargs,
)
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
if isinstance(query, str):
query = [query]
@@ -252,7 +252,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
class ColbertEmbeddingWorker(TextEmbeddingWorker):
def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> Colbert:
def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> Colbert:
return Colbert(
model_name=model_name,
cache_dir=cache_dir,

2
fastembed/late_interaction/jina_colbert.py
View File

@@ -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, model_name: str, cache_dir: str, **kwargs: Any) -> JinaColbert:
return JinaColbert(
model_name=model_name,
cache_dir=cache_dir,

12
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
import numpy as np
@@ -11,7 +11,7 @@ class LateInteractionTextEmbeddingBase(ModelManagement):
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
@@ -23,11 +23,11 @@ class LateInteractionTextEmbeddingBase(ModelManagement):
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, texts: Iterable[str], **kwargs: Any) -> Iterable[np.ndarray]:
"""
Embeds a list of text passages into a list of embeddings.
@@ -42,9 +42,7 @@ class LateInteractionTextEmbeddingBase(ModelManagement):
# 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
) -> Iterable[np.ndarray]:
def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
"""
Embeds queries

6
fastembed/late_interaction/late_interaction_text_embedding.py
View File

@@ -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:
@@ -80,7 +80,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
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, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
"""
Embeds queries

2
fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py
View File

@@ -217,7 +217,7 @@ class TextCrossEncoderWorker(TextRerankerWorker):
self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> OnnxTextCrossEncoder:
return OnnxTextCrossEncoder(
model_name=model_name,

4
fastembed/rerank/cross_encoder/text_cross_encoder_base.py
View File

@@ -9,7 +9,7 @@ class TextCrossEncoderBase(ModelManagement):
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
@@ -21,7 +21,7 @@ class TextCrossEncoderBase(ModelManagement):
query: str,
documents: Iterable[str],
batch_size: int = 64,
**kwargs,
**kwargs: Any,
) -> Iterable[float]:
"""Rerank a list of documents given a query.

12
fastembed/sparse/bm25.py
View File

@@ -109,7 +109,7 @@ class Bm25(SparseTextEmbeddingBase):
token_max_length: int = 40,
disable_stemmer: bool = False,
specific_model_path: Optional[str] = None,
**kwargs,
**kwargs: Any,
):
super().__init__(model_name, cache_dir, **kwargs)
@@ -213,7 +213,7 @@ class Bm25(SparseTextEmbeddingBase):
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,9 @@ class Bm25(SparseTextEmbeddingBase):
def compute_token_id(cls, token: str) -> int:
return abs(mmh3.hash(token))
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]:
def query_embed(
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.
"""
@@ -333,7 +335,7 @@ class Bm25Worker(Worker):
self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
):
self.model = self.init_embedding(model_name, cache_dir, **kwargs)
@@ -347,5 +349,5 @@ class Bm25Worker(Worker):
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)

10
fastembed/sparse/bm42.py
View File

@@ -67,7 +67,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
lazy_load: bool = False,
device_id: Optional[int] = None,
specific_model_path: Optional[str] = None,
**kwargs,
**kwargs: Any,
):
"""
Args:
@@ -273,7 +273,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
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.
@@ -310,7 +310,9 @@ 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, 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.
@@ -337,7 +339,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, model_name: str, cache_dir: str, **kwargs: Any) -> Bm42:
return Bm42(
model_name=model_name,
cache_dir=cache_dir,

12
fastembed/sparse/sparse_embedding_base.py
View File

@@ -1,5 +1,5 @@
from dataclasses import dataclass
from typing import Iterable, Optional, Union
from typing import Iterable, Optional, Union, Any
import numpy as np
@@ -34,7 +34,7 @@ class SparseTextEmbeddingBase(ModelManagement):
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
@@ -46,11 +46,11 @@ class SparseTextEmbeddingBase(ModelManagement):
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) -> Iterable[SparseEmbedding]:
def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
"""
Embeds a list of text passages into a list of embeddings.
@@ -65,7 +65,9 @@ class SparseTextEmbeddingBase(ModelManagement):
# 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) -> Iterable[SparseEmbedding]:
def query_embed(
self, query: Union[str, Iterable[str]], **kwargs: Any
) -> Iterable[SparseEmbedding]:
"""
Embeds queries

8
fastembed/sparse/sparse_text_embedding.py
View File

@@ -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":
@@ -89,7 +89,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
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,9 @@ 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, query: Union[str, Iterable[str]], **kwargs: Any
) -> Iterable[SparseEmbedding]:
"""
Embeds queries

6
fastembed/sparse/splade_pp.py
View File

@@ -74,7 +74,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
lazy_load: bool = False,
device_id: Optional[int] = None,
specific_model_path: Optional[str] = None,
**kwargs,
**kwargs: Any,
):
"""
Args:
@@ -141,7 +141,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
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.
@@ -176,7 +176,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, model_name: str, cache_dir: str, **kwargs: Any) -> SpladePP:
return SpladePP(
model_name=model_name,
cache_dir=cache_dir,

2
fastembed/text/clip_embedding.py
View File

@@ -44,7 +44,7 @@ class CLIPEmbeddingWorker(OnnxTextEmbeddingWorker):
self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> OnnxTextEmbedding:
return CLIPOnnxEmbedding(
model_name=model_name,

4
fastembed/text/e5_onnx_embedding.py
View File

@@ -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, onnx_input: dict[str, np.ndarray], **kwargs: Any
) -> dict[str, np.ndarray]:
"""
Preprocess the onnx input.
@@ -62,7 +62,7 @@ class E5OnnxEmbeddingWorker(OnnxTextEmbeddingWorker):
self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> E5OnnxEmbedding:
return E5OnnxEmbedding(
model_name=model_name,

2
fastembed/text/multitask_embedding.py
View File

@@ -42,7 +42,7 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
PASSAGE_TASK = Task.RETRIEVAL_PASSAGE
QUERY_TASK = Task.RETRIEVAL_QUERY
def __init__(self, *args, **kwargs):
def __init__(self, *args: Any, **kwargs: Any):
super().__init__(*args, **kwargs)
self._current_task_id = self.PASSAGE_TASK

8
fastembed/text/onnx_embedding.py
View File

@@ -194,7 +194,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
lazy_load: bool = False,
device_id: Optional[int] = None,
specific_model_path: Optional[str] = None,
**kwargs,
**kwargs: Any,
):
"""
Args:
@@ -250,7 +250,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
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.
@@ -284,7 +284,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
return OnnxTextEmbeddingWorker
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.
@@ -317,7 +317,7 @@ class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> OnnxTextEmbedding:
return OnnxTextEmbedding(
model_name=model_name,

8
fastembed/text/onnx_text_model.py
View File

@@ -29,7 +29,7 @@ class OnnxTextModel(OnnxModel[T]):
self.special_token_to_id = {}
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.
@@ -59,13 +59,13 @@ class OnnxTextModel(OnnxModel[T]):
def load_onnx_model(self) -> None:
raise NotImplementedError("Subclasses must implement this method")
def tokenize(self, documents: list[str], **kwargs) -> list[Encoding]:
def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
return self.tokenizer.encode_batch(documents)
def onnx_embed(
self,
documents: list[str],
**kwargs,
**kwargs: Any,
) -> OnnxOutputContext:
encoded = self.tokenize(documents, **kwargs)
input_ids = np.array([e.ids for e in encoded])
@@ -99,7 +99,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

2
fastembed/text/pooled_embedding.py
View File

@@ -93,7 +93,7 @@ class PooledEmbeddingWorker(OnnxTextEmbeddingWorker):
self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> OnnxTextEmbedding:
return PooledEmbedding(
model_name=model_name,

2
fastembed/text/pooled_normalized_embedding.py
View File

@@ -115,7 +115,7 @@ class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker):
self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> OnnxTextEmbedding:
return PooledNormalizedEmbedding(
model_name=model_name,

8
fastembed/text/text_embedding.py
View File

@@ -62,7 +62,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)
if model_name == "nomic-ai/nomic-embed-text-v1.5-Q":
@@ -105,7 +105,7 @@ class TextEmbedding(TextEmbeddingBase):
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.
@@ -124,7 +124,7 @@ class TextEmbedding(TextEmbeddingBase):
"""
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, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
"""
Embeds queries
@@ -137,7 +137,7 @@ class TextEmbedding(TextEmbeddingBase):
# This is model-specific, so that different models can have specialized implementations
yield from self.model.query_embed(query, **kwargs)
def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[np.ndarray]:
def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[np.ndarray]:
"""
Embeds a list of text passages into a list of embeddings.

12
fastembed/text/text_embedding_base.py
View File

@@ -1,4 +1,4 @@
from typing import Iterable, Optional, Union
from typing import Iterable, Optional, Union, Any
import numpy as np
@@ -11,7 +11,7 @@ class TextEmbeddingBase(ModelManagement):
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
@@ -23,11 +23,11 @@ class TextEmbeddingBase(ModelManagement):
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, texts: Iterable[str], **kwargs: Any) -> Iterable[np.ndarray]:
"""
Embeds a list of text passages into a list of embeddings.
@@ -42,9 +42,7 @@ class TextEmbeddingBase(ModelManagement):
# 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
) -> Iterable[np.ndarray]:
def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
"""
Embeds queries

6
pyproject.toml
View File

@@ -48,10 +48,16 @@ pillow = "^10.2.0"
cairosvg = "^2.7.1"
mknotebooks = "^0.8.0"
[tool.poetry.group.types.dependencies]
pyright = ">=1.1.293"
mypy = "^1.0.0"
[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
[tool.pyright]
typeCheckingMode = "strict"
[tool.ruff]
line-length = 99