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: Added type hints (#454)

Browse Source 
* chore: Added type hints

* fix: Fix generic type

* fix: Fix generic type

* new: Add type hints for parallel processor

* fix: Revert queue sub type as its not supported

* fix: Revert queue sub type as its not supported

* fix: Fixed type hints

* chore: Updated type hints

* fix: Update task id to be public

* chore: Updated type hints

* chore: Updated type hints

* fix: minor reverts in parallel processor and onnx text model
This commit is contained in:
Hossam Hagag
2025-02-04 18:06:05 +02:00
committed by GitHub
parent f1a3a6d082
commit 6dbdd6d171
15 changed files with 98 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
fastembed/common/model_management.py
View File

@@ -4,7 +4,7 @@ import json
import shutil
import tarfile
from pathlib import Path
from typing import Any, Optional
from typing import Any, Optional, Union
import requests
from huggingface_hub import snapshot_download, model_info, list_repo_tree
@@ -148,8 +148,8 @@ class ModelManagement:
def _collect_file_metadata(
model_dir: Path, repo_files: list[RepoFile]
) -> dict[str, dict[str, int]]:
meta = {}
) -> dict[str, dict[str, Union[int, str]]]:
meta: dict[str, dict[str, Union[int, str]]] = {}
file_info_map = {f.path: f for f in repo_files}
for file_path in model_dir.rglob("*"):
if file_path.is_file() and file_path.name != cls.METADATA_FILE:
@@ -161,7 +161,9 @@ class ModelManagement:
}
return meta
def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int]]) -> None:
def _save_file_metadata(
model_dir: Path, meta: dict[str, dict[str, Union[int, str]]]
) -> None:
try:
if not model_dir.exists():
model_dir.mkdir(parents=True, exist_ok=True)
@@ -338,7 +340,7 @@ class ModelManagement:
@classmethod
def download_model(
cls, model: dict[str, Any], cache_dir: Path, retries: int = 3, **kwargs: Any
cls, model: dict[str, Any], cache_dir: str, retries: int = 3, **kwargs: Any
) -> Path:
"""
Downloads a model from HuggingFace Hub or Google Cloud Storage.
@@ -384,7 +386,7 @@ class ModelManagement:
return Path(
cls.download_files_from_huggingface(
hf_source,
cache_dir=str(cache_dir),
cache_dir=cache_dir,
extra_patterns=extra_patterns,
**kwargs,
)

23
fastembed/common/onnx_model.py
View File

@@ -3,10 +3,9 @@ from dataclasses import dataclass
from pathlib import Path
from typing import Any, Generic, Iterable, Optional, Sequence, Type, TypeVar
import numpy as np
import onnxruntime as ort
from fastembed.common.types import OnnxProvider
from fastembed.common.types import OnnxProvider, NumpyArray
from fastembed.parallel_processor import Worker
# Holds type of the embedding result
@@ -15,14 +14,14 @@ T = TypeVar("T")
@dataclass
class OnnxOutputContext:
model_output: np.ndarray
attention_mask: Optional[np.ndarray] = None
input_ids: Optional[np.ndarray] = None
model_output: NumpyArray
attention_mask: Optional[NumpyArray] = None
input_ids: Optional[NumpyArray] = None
class OnnxModel(Generic[T]):
@classmethod
def _get_worker_class(cls) -> Type["EmbeddingWorker"]:
def _get_worker_class(cls) -> Type["EmbeddingWorker[T]"]:
raise NotImplementedError("Subclasses must implement this method")
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
@@ -33,8 +32,8 @@ class OnnxModel(Generic[T]):
self.tokenizer = None
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.
"""
@@ -70,7 +69,7 @@ class OnnxModel(Generic[T]):
onnx_providers = ["CPUExecutionProvider"]
available_providers = ort.get_available_providers()
requested_provider_names = []
requested_provider_names: list[str] = []
for provider in onnx_providers:
# check providers available
provider_name = provider if isinstance(provider, str) else provider[0]
@@ -107,13 +106,13 @@ class OnnxModel(Generic[T]):
raise NotImplementedError("Subclasses must implement this method")
class EmbeddingWorker(Worker):
class EmbeddingWorker(Worker, Generic[T]):
def init_embedding(
self,
model_name: str,
cache_dir: str,
**kwargs: Any,
) -> OnnxModel:
) -> OnnxModel[T]:
raise NotImplementedError()
def __init__(
@@ -125,7 +124,7 @@ class EmbeddingWorker(Worker):
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, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker[T]":
return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:

7
fastembed/common/preprocessor_utils.py
View File

@@ -1,4 +1,5 @@
import json
from typing import Any
from pathlib import Path
from tokenizers import AddedToken, Tokenizer
@@ -6,7 +7,7 @@ from tokenizers import AddedToken, Tokenizer
from fastembed.image.transform.operators import Compose
def load_special_tokens(model_dir: Path) -> dict:
def load_special_tokens(model_dir: Path) -> dict[str, Any]:
tokens_map_path = model_dir / "special_tokens_map.json"
if not tokens_map_path.exists():
raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
@@ -17,7 +18,7 @@ def load_special_tokens(model_dir: Path) -> dict:
return tokens_map
def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict]:
def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
config_path = model_dir / "config.json"
if not config_path.exists():
raise ValueError(f"Could not find config.json in {model_dir}")
@@ -59,7 +60,7 @@ def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict]:
elif isinstance(token, dict):
tokenizer.add_special_tokens([AddedToken(**token)])
special_token_to_id = {}
special_token_to_id: dict[str, int] = {}
for token in tokens_map.values():
if isinstance(token, str):

11
fastembed/common/types.py
View File

@@ -3,6 +3,9 @@ import sys
from PIL import Image
from typing import Any, Iterable, Union
import numpy as np
from numpy.typing import NDArray
if sys.version_info >= (3, 10):
from typing import TypeAlias
else:
@@ -14,3 +17,11 @@ PilInput: TypeAlias = Union[Image.Image, Iterable[Image.Image]]
ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput], PilInput]
OnnxProvider: TypeAlias = Union[str, tuple[str, dict[Any, Any]]]
NumpyArray = Union[
NDArray[np.float32],
NDArray[np.float16],
NDArray[np.int8],
NDArray[np.int64],
NDArray[np.int32],
]

4
fastembed/common/utils.py
View File

@@ -9,10 +9,12 @@ from typing import Iterable, Optional, TypeVar
import numpy as np
from fastembed.common.types import NumpyArray
T = TypeVar("T")
def normalize(input_array: np.ndarray, p: int = 2, dim: int = 1, eps: float = 1e-12) -> np.ndarray:
def normalize(input_array: NumpyArray, p: int = 2, dim: int = 1, eps: float = 1e-12) -> NumpyArray:
# 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

2
fastembed/parallel_processor.py
View File

@@ -140,7 +140,7 @@ class ParallelWorkerPool:
self.processes.append(process)
def ordered_map(self, stream: Iterable[Any], *args: Any, **kwargs: Any) -> Iterable[Any]:
buffer = defaultdict(Any)
buffer: defaultdict[int, Any] = defaultdict(Any)
next_expected = 0
for idx, item in self.semi_ordered_map(stream, *args, **kwargs):

8
fastembed/text/clip_embedding.py
View File

@@ -1,10 +1,8 @@
from typing import Any, Iterable, Type
import numpy as np
from fastembed.common.types import NumpyArray
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
from fastembed.text.onnx_text_model import TextEmbeddingWorker
supported_clip_models = [
{
@@ -23,7 +21,7 @@ supported_clip_models = [
class CLIPOnnxEmbedding(OnnxTextEmbedding):
@classmethod
def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
return CLIPEmbeddingWorker
@classmethod
@@ -35,7 +33,7 @@ class CLIPOnnxEmbedding(OnnxTextEmbedding):
"""
return supported_clip_models
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
return output.model_output

30
fastembed/text/multitask_embedding.py
View File

@@ -3,9 +3,9 @@ from typing import Any, Type, Iterable, Union, Optional
import numpy as np
from fastembed.common.types import NumpyArray
from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
from fastembed.text.onnx_embedding import OnnxTextEmbeddingWorker
from fastembed.text.onnx_text_model import TextEmbeddingWorker
supported_multitask_models = [
{
@@ -44,10 +44,10 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
def __init__(self, *args: Any, **kwargs: Any):
super().__init__(*args, **kwargs)
self._current_task_id = self.PASSAGE_TASK
self.current_task_id: Union[Task, int] = self.PASSAGE_TASK
@classmethod
def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
return JinaEmbeddingV3Worker
@classmethod
@@ -55,9 +55,9 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
return supported_multitask_models
def _preprocess_onnx_input(
self, onnx_input: dict[str, np.ndarray], **kwargs
) -> dict[str, np.ndarray]:
onnx_input["task_id"] = np.array(self._current_task_id, dtype=np.int64)
self, onnx_input: dict[str, NumpyArray], **kwargs: Any
) -> dict[str, NumpyArray]:
onnx_input["task_id"] = np.array(self.current_task_id, dtype=np.int64)
return onnx_input
def embed(
@@ -66,18 +66,18 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
batch_size: int = 256,
parallel: Optional[int] = None,
task_id: int = PASSAGE_TASK,
**kwargs,
) -> Iterable[np.ndarray]:
self._current_task_id = task_id
**kwargs: Any,
) -> Iterable[NumpyArray]:
self.current_task_id = task_id
kwargs["task_id"] = task_id
yield from super().embed(documents, batch_size, parallel, **kwargs)
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
self._current_task_id = self.QUERY_TASK
def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
self.current_task_id = self.QUERY_TASK
yield from super().embed(query, **kwargs)
def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[np.ndarray]:
self._current_task_id = self.PASSAGE_TASK
def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
self.current_task_id = self.PASSAGE_TASK
yield from super().embed(texts, **kwargs)
@@ -86,7 +86,7 @@ class JinaEmbeddingV3Worker(OnnxTextEmbeddingWorker):
self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> JinaEmbeddingV3:
model = JinaEmbeddingV3(
model_name=model_name,
@@ -94,5 +94,5 @@ class JinaEmbeddingV3Worker(OnnxTextEmbeddingWorker):
threads=1,
**kwargs,
)
model._current_task_id = kwargs["task_id"]
model.current_task_id = kwargs["task_id"]
return model

21
fastembed/text/onnx_embedding.py
View File

@@ -2,7 +2,7 @@ from typing import Any, Iterable, Optional, Sequence, Type, Union
import numpy as np
from fastembed.common import OnnxProvider
from fastembed.common.types import NumpyArray, OnnxProvider
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common.utils import define_cache_dir, normalize
from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
@@ -170,7 +170,7 @@ supported_onnx_models = [
]
class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
"""Implementation of the Flag Embedding model."""
@classmethod
@@ -226,15 +226,14 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
self.cuda = cuda
# 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,
@@ -251,7 +250,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs: Any,
) -> Iterable[np.ndarray]:
) -> Iterable[NumpyArray]:
"""
Encode a list of documents into list of embeddings.
We use mean pooling with attention so that the model can handle variable-length inputs.
@@ -280,18 +279,18 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
)
@classmethod
def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
def _get_worker_class(cls) -> Type["TextEmbeddingWorker[NumpyArray]"]:
return OnnxTextEmbeddingWorker
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
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
embeddings = output.model_output
if embeddings.ndim == 3: # (batch_size, seq_len, embedding_dim)
processed_embeddings = embeddings[:, 0]
@@ -312,7 +311,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
)
class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
class OnnxTextEmbeddingWorker(TextEmbeddingWorker[NumpyArray]):
def init_embedding(
self,
model_name: str,

16
fastembed/text/onnx_text_model.py
View File

@@ -6,7 +6,7 @@ from typing import Any, Iterable, Optional, Sequence, Type, Union
import numpy as np
from tokenizers import Encoding
from fastembed.common import OnnxProvider
from fastembed.common.types import NumpyArray, OnnxProvider
from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
from fastembed.common.preprocessor_utils import load_tokenizer
from fastembed.common.utils import iter_batch
@@ -17,7 +17,7 @@ class OnnxTextModel(OnnxModel[T]):
ONNX_OUTPUT_NAMES: Optional[list[str]] = None
@classmethod
def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
def _get_worker_class(cls) -> Type["TextEmbeddingWorker[T]"]:
raise NotImplementedError("Subclasses must implement this method")
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
@@ -26,11 +26,11 @@ class OnnxTextModel(OnnxModel[T]):
def __init__(self):
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: Any
) -> dict[str, np.ndarray]:
self, onnx_input: dict[str, NumpyArray], **kwargs: Any
) -> dict[str, NumpyArray]:
"""
Preprocess the onnx input.
"""
@@ -71,7 +71,7 @@ class OnnxTextModel(OnnxModel[T]):
input_ids = np.array([e.ids for e in encoded])
attention_mask = np.array([e.attention_mask for e in encoded])
input_names = {node.name for node in self.model.get_inputs()}
onnx_input = {
onnx_input: dict[str, NumpyArray] = {
"input_ids": np.array(input_ids, dtype=np.int64),
}
if "attention_mask" in input_names:
@@ -139,8 +139,8 @@ class OnnxTextModel(OnnxModel[T]):
yield from self._post_process_onnx_output(batch)
class TextEmbeddingWorker(EmbeddingWorker):
def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
class TextEmbeddingWorker(EmbeddingWorker[T]):
def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, OnnxOutputContext]]:
for idx, batch in items:
onnx_output = self.model.onnx_embed(batch)
yield idx, onnx_output

13
fastembed/text/pooled_embedding.py
View File

@@ -2,9 +2,9 @@ from typing import Any, Iterable, Type
import numpy as np
from fastembed.common.types import NumpyArray
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
from fastembed.text.onnx_text_model import TextEmbeddingWorker
supported_pooled_models = [
{
@@ -80,15 +80,16 @@ supported_pooled_models = [
class PooledEmbedding(OnnxTextEmbedding):
@classmethod
def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
return PooledEmbeddingWorker
@classmethod
def mean_pooling(cls, model_output: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
token_embeddings = model_output
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)
input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
input_mask_expanded = np.tile(input_mask_expanded, (1, 1, token_embeddings.shape[-1]))
input_mask_expanded = input_mask_expanded.astype(float)
input_mask_expanded = input_mask_expanded.astype(np.float32)
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)
@@ -103,7 +104,7 @@ class PooledEmbedding(OnnxTextEmbedding):
"""
return supported_pooled_models
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
if output.attention_mask is None:
raise ValueError("attention_mask must be provided for document post-processing")

6
fastembed/text/pooled_normalized_embedding.py
View File

@@ -2,10 +2,10 @@ from typing import Any, Iterable, Type
import numpy as np
from fastembed.common.types import NumpyArray
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common.utils import normalize
from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
from fastembed.text.onnx_text_model import TextEmbeddingWorker
from fastembed.text.pooled_embedding import PooledEmbedding
supported_pooled_normalized_models = [
@@ -89,7 +89,7 @@ supported_pooled_normalized_models = [
class PooledNormalizedEmbedding(PooledEmbedding):
@classmethod
def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
return PooledNormalizedEmbeddingWorker
@classmethod
@@ -101,7 +101,7 @@ class PooledNormalizedEmbedding(PooledEmbedding):
"""
return supported_pooled_normalized_models
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
if output.attention_mask is None:
raise ValueError("attention_mask must be provided for document post-processing")

13
fastembed/text/text_embedding.py
View File

@@ -1,8 +1,7 @@
import warnings
from typing import Any, Iterable, Optional, Sequence, Type, Union
import numpy as np
from fastembed.common import OnnxProvider
from fastembed.common.types import NumpyArray, OnnxProvider
from fastembed.text.clip_embedding import CLIPOnnxEmbedding
from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
from fastembed.text.pooled_embedding import PooledEmbedding
@@ -45,7 +44,7 @@ class TextEmbedding(TextEmbeddingBase):
]
```
"""
result = []
result: list[dict[str, Any]] = []
for embedding in cls.EMBEDDINGS_REGISTRY:
result.extend(embedding.list_supported_models())
return result
@@ -113,7 +112,7 @@ class TextEmbedding(TextEmbeddingBase):
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs: Any,
) -> Iterable[np.ndarray]:
) -> Iterable[NumpyArray]:
"""
Encode a list of documents into list of embeddings.
We use mean pooling with attention so that the model can handle variable-length inputs.
@@ -131,7 +130,7 @@ class TextEmbedding(TextEmbeddingBase):
"""
yield from self.model.embed(documents, batch_size, parallel, **kwargs)
def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
"""
Embeds queries
@@ -139,12 +138,12 @@ class TextEmbedding(TextEmbeddingBase):
query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
Returns:
Iterable[np.ndarray]: The embeddings.
Iterable[NumpyArray]: The embeddings.
"""
# 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: Any) -> Iterable[np.ndarray]:
def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
"""
Embeds a list of text passages into a list of embeddings.

16
fastembed/text/text_embedding_base.py
View File

@@ -1,7 +1,6 @@
from typing import Iterable, Optional, Union, Any
import numpy as np
from fastembed.common.types import NumpyArray
from fastembed.common.model_management import ModelManagement
@@ -24,10 +23,10 @@ class TextEmbeddingBase(ModelManagement):
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs: Any,
) -> Iterable[np.ndarray]:
) -> Iterable[NumpyArray]:
raise NotImplementedError()
def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[np.ndarray]:
def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
"""
Embeds a list of text passages into a list of embeddings.
@@ -36,14 +35,13 @@ class TextEmbeddingBase(ModelManagement):
**kwargs: Additional keyword argument to pass to the embed method.
Yields:
Iterable[np.ndarray]: The embeddings.
Iterable[NumpyArray]: The embeddings.
"""
# 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: Any) -> Iterable[np.ndarray]:
def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
"""
Embeds queries
@@ -51,11 +49,11 @@ class TextEmbeddingBase(ModelManagement):
query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
Returns:
Iterable[np.ndarray]: The embeddings.
Iterable[NumpyArray]: The embeddings.
"""
# This is model-specific, so that different models can have specialized implementations
if isinstance(query, str):
yield from self.embed([query], **kwargs)
if isinstance(query, Iterable):
else:
yield from self.embed(query, **kwargs)

2
tests/test_text_multitask_embeddings.py
View File

@@ -231,7 +231,7 @@ def test_task_assignment():
for i, task_id in enumerate(Task):
_ = list(model.embed(documents=docs, batch_size=1, task_id=i))
assert model.model._current_task_id == task_id
assert model.model.current_task_id == task_id
if is_ci:
delete_model_cache(model.model._model_dir)