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

new: drop python 3.8 support, update type hints, ci (#403)

Browse Source
This commit is contained in:
George
2024-11-15 15:54:19 +01:00
committed by GitHub
parent 9841666bd5
commit 1343e55076
39 changed files with 234 additions and 240 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.github/workflows/python-tests.yml vendored
View File

@@ -14,7 +14,6 @@ jobs:
strategy:
matrix:
python-version:
- '3.8.x'
- '3.9.x'
- '3.10.x'
- '3.11.x'

4
docs/Getting Started.ipynb
View File

@@ -65,15 +65,13 @@
}
],
"source": [
"from typing import List\n",
"\n",
"import numpy as np\n",
"\n",
"from fastembed import TextEmbedding\n",
"\n",
"\n",
"# Example list of documents\n",
"documents: List[str] = [\n",
"documents: list[str] = [\n",
" \"This is built to be faster and lighter than other embedding libraries e.g. Transformers, Sentence-Transformers, etc.\",\n",
" \"fastembed is supported by and maintained by Qdrant.\",\n",
"]\n",

6
docs/examples/FastEmbed_GPU.ipynb
View File

@@ -388,8 +388,6 @@
}
],
"source": [
"from typing import List\n",
"\n",
"import numpy as np\n",
"\n",
"from fastembed import TextEmbedding\n",
@@ -407,9 +405,7 @@
"id": "iPtoHf7GeV-i"
},
"outputs": [],
"source": [
"documents: List[str] = list(np.repeat(\"Demonstrating GPU acceleration in fastembed\", 500))"
]
"source": "documents: list[str] = list(np.repeat(\"Demonstrating GPU acceleration in fastembed\", 500))"
},
{
"cell_type": "code",

11
docs/examples/FastEmbed_vs_HF_Comparison.ipynb
View File

@@ -36,7 +36,7 @@
"outputs": [],
"source": [
"import time\n",
"from typing import Callable, List, Tuple\n",
"from typing import Callable\n",
"\n",
"import matplotlib.pyplot as plt\n",
"import torch.nn.functional as F\n",
@@ -64,6 +64,7 @@
],
"source": [
"import fastembed\n",
"\n",
"fastembed.__version__"
]
},
@@ -98,7 +99,7 @@
}
],
"source": [
"documents: List[str] = [\n",
"documents: list[str] = [\n",
" \"Chandrayaan-3 is India's third lunar mission\",\n",
" \"It aimed to land a rover on the Moon's surface - joining the US, China and Russia\",\n",
" \"The mission is a follow-up to Chandrayaan-2, which had partial success\",\n",
@@ -155,7 +156,7 @@
" self.model = AutoModel.from_pretrained(model_id)\n",
" self.tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
"\n",
" def embed(self, texts: List[str]):\n",
" def embed(self, texts: list[str]):\n",
" encoded_input = self.tokenizer(\n",
" texts, max_length=512, padding=True, truncation=True, return_tensors=\"pt\"\n",
" )\n",
@@ -254,7 +255,7 @@
"\n",
"def calculate_time_stats(\n",
" embed_func: Callable, documents: list, k: int\n",
") -> Tuple[float, float, float]:\n",
") -> tuple[float, float, float]:\n",
" times = []\n",
" for _ in range(k):\n",
" # Timing the embed_func call\n",
@@ -309,7 +310,7 @@
],
"source": [
"def plot_character_per_second_comparison(\n",
" hf_stats: Tuple[float, float, float], fst_stats: Tuple[float, float, float], documents: list\n",
" hf_stats: tuple[float, float, float], fst_stats: tuple[float, float, float], documents: list\n",
"):\n",
" # Calculating total characters in documents\n",
" total_characters = sum(len(doc) for doc in documents)\n",

19
docs/examples/Hybrid_Search.ipynb
View File

@@ -50,7 +50,6 @@
"outputs": [],
"source": [
"import json\n",
"from typing import List, Tuple\n",
"\n",
"import numpy as np\n",
"import pandas as pd\n",
@@ -489,11 +488,11 @@
}
],
"source": [
"def make_sparse_embedding(texts: List[str]):\n",
"def make_sparse_embedding(texts: list[str]):\n",
" return list(sparse_model.embed(texts, batch_size=32))\n",
"\n",
"\n",
"sparse_embedding: List[SparseEmbedding] = make_sparse_embedding(\n",
"sparse_embedding: list[SparseEmbedding] = make_sparse_embedding(\n",
" [\"Fastembed is a great library for text embeddings!\"]\n",
")\n",
"sparse_embedding"
@@ -662,7 +661,7 @@
},
"outputs": [],
"source": [
"def make_dense_embedding(texts: List[str]):\n",
"def make_dense_embedding(texts: list[str]):\n",
" return list(dense_model.embed(texts))\n",
"\n",
"\n",
@@ -872,7 +871,7 @@
},
"outputs": [],
"source": [
"def make_points(df: pd.DataFrame) -> List[PointStruct]:\n",
"def make_points(df: pd.DataFrame) -> list[PointStruct]:\n",
" sparse_vectors = df[\"sparse_embedding\"].tolist()\n",
" product_texts = df[\"combined_text\"].tolist()\n",
" dense_vectors = df[\"dense_embedding\"].tolist()\n",
@@ -899,7 +898,7 @@
" return points\n",
"\n",
"\n",
"points: List[PointStruct] = make_points(df)"
"points: list[PointStruct] = make_points(df)"
]
},
{
@@ -942,8 +941,8 @@
"source": [
"def search(query_text: str):\n",
" # # Compute sparse and dense vectors\n",
" query_sparse_vectors: List[SparseEmbedding] = make_sparse_embedding([query_text])\n",
" query_dense_vector: List[np.ndarray] = make_dense_embedding([query_text])\n",
" query_sparse_vectors: list[SparseEmbedding] = make_sparse_embedding([query_text])\n",
" query_dense_vector: list[np.ndarray] = make_dense_embedding([query_text])\n",
"\n",
" search_results = client.search_batch(\n",
" collection_name=collection_name,\n",
@@ -1075,7 +1074,7 @@
"metadata": {},
"outputs": [],
"source": [
"def rank_list(search_result: List[ScoredPoint]):\n",
"def rank_list(search_result: list[ScoredPoint]):\n",
" return [(point.id, rank + 1) for rank, point in enumerate(search_result)]\n",
"\n",
"\n",
@@ -1149,7 +1148,7 @@
],
"source": [
"def find_point_by_id(\n",
" client: QdrantClient, collection_name: str, rrf_rank_list: List[Tuple[int, float]]\n",
" client: QdrantClient, collection_name: str, rrf_rank_list: list[tuple[int, float]]\n",
"):\n",
" return client.retrieve(\n",
" collection_name=collection_name, ids=[item[0] for item in rrf_rank_list]\n",

5
docs/qdrant/Retrieval_with_FastEmbed.ipynb
View File

@@ -41,7 +41,6 @@
"metadata": {},
"outputs": [],
"source": [
"from typing import List\n",
"import numpy as np\n",
"from fastembed import TextEmbedding"
]
@@ -71,7 +70,7 @@
],
"source": [
"# Example list of documents\n",
"documents: List[str] = [\n",
"documents: list[str] = [\n",
" \"Maharana Pratap was a Rajput warrior king from Mewar\",\n",
" \"He fought against the Mughal Empire led by Akbar\",\n",
" \"The Battle of Haldighati in 1576 was his most famous battle\",\n",
@@ -87,7 +86,7 @@
"embedding_model = TextEmbedding(model_name=\"BAAI/bge-small-en\")\n",
"\n",
"# We'll use the passage_embed method to get the embeddings for the documents\n",
"embeddings: List[np.ndarray] = list(\n",
"embeddings: list[np.ndarray] = list(\n",
" embedding_model.passage_embed(documents)\n",
") # notice that we are casting the generator to a list\n",
"\n",

7
docs/qdrant/Usage_With_Qdrant.ipynb
View File

@@ -46,7 +46,6 @@
"metadata": {},
"outputs": [],
"source": [
"from typing import List\n",
"from qdrant_client import QdrantClient"
]
},
@@ -67,7 +66,7 @@
"outputs": [],
"source": [
"# Example list of documents\n",
"documents: List[str] = [\n",
"documents: list[str] = [\n",
" \"Maharana Pratap was a Rajput warrior king from Mewar\",\n",
" \"He fought against the Mughal Empire led by Akbar\",\n",
" \"The Battle of Haldighati in 1576 was his most famous battle\",\n",
@@ -199,7 +198,9 @@
}
],
"source": [
"search_result = client.query(collection_name=\"demo_collection\", query_text=\"This is a query document\")\n",
"search_result = client.query(\n",
" collection_name=\"demo_collection\", query_text=\"This is a query document\"\n",
")\n",
"print(search_result)"
]
},

16
experiments/01_ONNX_Port.ipynb
View File

@@ -19,7 +19,7 @@
"outputs": [],
"source": [
"from pathlib import Path\n",
"from typing import List, Tuple, Any\n",
"from typing import Any\n",
"\n",
"import numpy as np\n",
"import time\n",
@@ -91,9 +91,11 @@
" return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]\n",
"\n",
"\n",
"def hf_embed(model_id: str, inputs: List[str]):\n",
"def hf_embed(model_id: str, inputs: list[str]):\n",
" # Tokenize the input texts\n",
" batch_dict = hf_tokenizer(inputs, max_length=512, padding=True, truncation=True, return_tensors=\"pt\")\n",
" batch_dict = hf_tokenizer(\n",
" inputs, max_length=512, padding=True, truncation=True, return_tensors=\"pt\"\n",
" )\n",
"\n",
" outputs = hf_model(**batch_dict)\n",
" embeddings = average_pool(outputs.last_hidden_state, batch_dict[\"attention_mask\"])\n",
@@ -133,7 +135,9 @@
"optimization_config = AutoOptimizationConfig.O4()\n",
"optimizer = ORTOptimizer.from_pretrained(model)\n",
"\n",
"optimizer.optimize(save_dir=save_dir, optimization_config=optimization_config, use_external_data_format=True)\n",
"optimizer.optimize(\n",
" save_dir=save_dir, optimization_config=optimization_config, use_external_data_format=True\n",
")\n",
"model = ORTModelForFeatureExtraction.from_pretrained(save_dir)\n",
"\n",
"tokenizer.save_pretrained(save_dir)\n",
@@ -171,7 +175,9 @@
"metadata": {},
"outputs": [],
"source": [
"def measure_pipeline_time(pipeline, input_texts: List[str], num_runs=10, **kwargs: Any) -> Tuple[float, float]:\n",
"def measure_pipeline_time(\n",
" pipeline, input_texts: list[str], num_runs=10, **kwargs: Any\n",
") -> tuple[float, float]:\n",
" \"\"\"Measures the time it takes to run the pipeline on the input texts.\"\"\"\n",
" times = []\n",
" total_chars = sum(len(text) for text in input_texts)\n",

27
experiments/Throughput_Across_Models.ipynb
View File

File diff suppressed because one or more lines are too long

18
fastembed/common/model_management.py
View File

@@ -3,7 +3,7 @@ import time
import shutil
import tarfile
from pathlib import Path
from typing import Any, Dict, List, Optional
from typing import Any, Optional
import requests
from huggingface_hub import snapshot_download
@@ -14,16 +14,16 @@ from tqdm import tqdm
class ModelManagement:
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls) -> list[dict[str, Any]]:
"""Lists the supported models.
Returns:
List[Dict[str, Any]]: A list of dictionaries containing the model information.
list[dict[str, Any]]: A list of dictionaries containing the model information.
"""
raise NotImplementedError()
@classmethod
def _get_model_description(cls, model_name: str) -> Dict[str, Any]:
def _get_model_description(cls, model_name: str) -> dict[str, Any]:
"""
Gets the model description from the model_name.
@@ -34,7 +34,7 @@ class ModelManagement:
ValueError: If the model_name is not supported.
Returns:
Dict[str, Any]: The model description.
dict[str, Any]: The model description.
"""
for model in cls.list_supported_models():
if model_name.lower() == model["model"].lower():
@@ -94,7 +94,7 @@ class ModelManagement:
cls,
hf_source_repo: str,
cache_dir: Optional[str] = None,
extra_patterns: Optional[List[str]] = None,
extra_patterns: Optional[list[str]] = None,
local_files_only: bool = False,
**kwargs,
) -> str:
@@ -103,7 +103,7 @@ class ModelManagement:
Args:
hf_source_repo (str): Name of the model on HuggingFace Hub, e.g. "qdrant/all-MiniLM-L6-v2-onnx".
cache_dir (Optional[str]): The path to the cache directory.
extra_patterns (Optional[List[str]]): extra patterns to allow in the snapshot download, typically
extra_patterns (Optional[list[str]]): extra patterns to allow in the snapshot download, typically
includes the required model files.
local_files_only (bool, optional): Whether to only use local files. Defaults to False.
Returns:
@@ -211,13 +211,13 @@ 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
) -> Path:
"""
Downloads a model from HuggingFace Hub or Google Cloud Storage.
Args:
model (Dict[str, Any]): The model description.
model (dict[str, Any]): The model description.
Example:
```
{

8
fastembed/common/onnx_model.py
View File

@@ -1,7 +1,7 @@
import warnings
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Generic, Iterable, Optional, Sequence, Tuple, Type, TypeVar
from typing import Any, Generic, Iterable, Optional, Sequence, Type, TypeVar
import numpy as np
import onnxruntime as ort
@@ -33,8 +33,8 @@ class OnnxModel(Generic[T]):
self.tokenizer = None
def _preprocess_onnx_input(
self, onnx_input: Dict[str, np.ndarray], **kwargs
) -> Dict[str, np.ndarray]:
self, onnx_input: dict[str, np.ndarray], **kwargs
) -> dict[str, np.ndarray]:
"""
Preprocess the onnx input.
"""
@@ -121,5 +121,5 @@ class EmbeddingWorker(Worker):
def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker":
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, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
raise NotImplementedError("Subclasses must implement this method")

4
fastembed/common/preprocessor_utils.py
View File

@@ -1,6 +1,6 @@
import json
from pathlib import Path
from typing import Tuple
from tokenizers import AddedToken, Tokenizer
from fastembed.image.transform.operators import Compose
@@ -17,7 +17,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]:
config_path = model_dir / "config.json"
if not config_path.exists():
raise ValueError(f"Could not find config.json in {model_dir}")

4
fastembed/common/types.py
View File

@@ -1,7 +1,7 @@
import os
import sys
from PIL import Image
from typing import Any, Dict, Iterable, Tuple, Union
from typing import Any, Iterable, Union
if sys.version_info >= (3, 10):
from typing import TypeAlias
@@ -13,4 +13,4 @@ PathInput: TypeAlias = Union[str, os.PathLike]
PilInput: TypeAlias = Union[Image.Image, Iterable[Image.Image]]
ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput], PilInput]
OnnxProvider: TypeAlias = Union[str, Tuple[str, Dict[Any, Any]]]
OnnxProvider: TypeAlias = Union[str, tuple[str, dict[Any, Any]]]

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, Iterable, Optional, Sequence, Type
import numpy as np
@@ -8,15 +8,15 @@ from fastembed.image.onnx_embedding import OnnxImageEmbedding
class ImageEmbedding(ImageEmbeddingBase):
EMBEDDINGS_REGISTRY: List[Type[ImageEmbeddingBase]] = [OnnxImageEmbedding]
EMBEDDINGS_REGISTRY: list[Type[ImageEmbeddingBase]] = [OnnxImageEmbedding]
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls) -> list[dict[str, Any]]:
"""
Lists the supported models.
Returns:
List[Dict[str, Any]]: A list of dictionaries containing the model information.
list[dict[str, Any]]: A list of dictionaries containing the model information.
Example:
```
@@ -47,7 +47,7 @@ class ImageEmbedding(ImageEmbeddingBase):
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
**kwargs,
):

10
fastembed/image/onnx_embedding.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type
from typing import Any, Iterable, Optional, Sequence, Type
import numpy as np
@@ -64,7 +64,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
@@ -129,7 +129,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
)
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls) -> list[dict[str, Any]]:
"""
Lists the supported models.
@@ -178,8 +178,8 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
return OnnxImageEmbeddingWorker
def _preprocess_onnx_input(
self, onnx_input: Dict[str, np.ndarray], **kwargs
) -> Dict[str, np.ndarray]:
self, onnx_input: dict[str, np.ndarray], **kwargs
) -> dict[str, np.ndarray]:
"""
Preprocess the onnx input.
"""

14
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, Iterable, Optional, Sequence, Type
import numpy as np
from PIL import Image
@@ -29,8 +29,8 @@ class OnnxImageModel(OnnxModel[T]):
self.processor = None
def _preprocess_onnx_input(
self, onnx_input: Dict[str, np.ndarray], **kwargs
) -> Dict[str, np.ndarray]:
self, onnx_input: dict[str, np.ndarray], **kwargs
) -> dict[str, np.ndarray]:
"""
Preprocess the onnx input.
"""
@@ -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, 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) -> OnnxOutputContext:
with contextlib.ExitStack():
image_files = [
Image.open(image) if not isinstance(image, Image.Image) else image
@@ -83,7 +83,7 @@ class OnnxImageModel(OnnxModel[T]):
parallel: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
device_ids: Optional[list[int]] = None,
**kwargs,
) -> 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, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
for idx, batch in items:
embeddings = self.model.onnx_embed(batch)
yield idx, embeddings

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

@@ -1,4 +1,4 @@
from typing import Sized, Tuple, Union
from typing import Sized, Union
import numpy as np
from PIL import Image
@@ -14,7 +14,7 @@ def convert_to_rgb(image: Image.Image) -> Image.Image:
def center_crop(
image: Union[Image.Image, np.ndarray],
size: Tuple[int, int],
size: tuple[int, int],
) -> np.ndarray:
if isinstance(image, np.ndarray):
_, orig_height, orig_width = image.shape
@@ -97,7 +97,7 @@ def normalize(
def resize(
image: Image,
size: Union[int, Tuple[int, int]],
size: Union[int, tuple[int, int]],
resample: Image.Resampling = Image.Resampling.BILINEAR,
) -> Image:
if isinstance(size, tuple):

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

@@ -1,4 +1,4 @@
from typing import Any, Dict, List, Tuple, Union
from typing import Any, Union
import numpy as np
from PIL import Image
@@ -14,78 +14,74 @@ from fastembed.image.transform.functional import (
class Transform:
def __call__(self, images: List) -> Union[List[Image.Image], List[np.ndarray]]:
def __call__(self, images: list) -> Union[list[Image.Image], list[np.ndarray]]:
raise NotImplementedError("Subclasses must implement this method")
class ConvertToRGB(Transform):
def __call__(self, images: List[Image.Image]) -> List[Image.Image]:
def __call__(self, images: list[Image.Image]) -> list[Image.Image]:
return [convert_to_rgb(image=image) for image in images]
class CenterCrop(Transform):
def __init__(self, size: Tuple[int, int]):
def __init__(self, size: tuple[int, int]):
self.size = size
def __call__(self, images: List[Image.Image]) -> List[np.ndarray]:
def __call__(self, images: list[Image.Image]) -> list[np.ndarray]:
return [center_crop(image=image, size=self.size) for image in images]
class Normalize(Transform):
def __init__(self, mean: Union[float, List[float]], std: Union[float, List[float]]):
def __init__(self, mean: Union[float, list[float]], std: Union[float, list[float]]):
self.mean = mean
self.std = std
def __call__(self, images: List[np.ndarray]) -> List[np.ndarray]:
def __call__(self, images: list[np.ndarray]) -> list[np.ndarray]:
return [normalize(image, mean=self.mean, std=self.std) for image in images]
class Resize(Transform):
def __init__(
self,
size: Union[int, Tuple[int, int]],
size: Union[int, tuple[int, int]],
resample: Image.Resampling = Image.Resampling.BICUBIC,
):
self.size = size
self.resample = resample
def __call__(self, images: List[Image.Image]) -> List[Image.Image]:
return [
resize(image, size=self.size, resample=self.resample) for image in images
]
def __call__(self, images: list[Image.Image]) -> list[Image.Image]:
return [resize(image, size=self.size, resample=self.resample) for image in images]
class Rescale(Transform):
def __init__(self, scale: float = 1 / 255):
self.scale = scale
def __call__(self, images: List[np.ndarray]) -> List[np.ndarray]:
def __call__(self, images: list[np.ndarray]) -> list[np.ndarray]:
return [rescale(image, scale=self.scale) for image in images]
class PILtoNDarray(Transform):
def __call__(
self, images: List[Union[Image.Image, np.ndarray]]
) -> List[np.ndarray]:
def __call__(self, images: list[Union[Image.Image, np.ndarray]]) -> list[np.ndarray]:
return [pil2ndarray(image) for image in images]
class Compose:
def __init__(self, transforms: List[Transform]):
def __init__(self, transforms: list[Transform]):
self.transforms = transforms
def __call__(
self, images: Union[List[Image.Image], List[np.ndarray]]
) -> Union[List[np.ndarray], List[Image.Image]]:
self, images: Union[list[Image.Image], list[np.ndarray]]
) -> Union[list[np.ndarray], list[Image.Image]]:
for transform in self.transforms:
images = transform(images)
return images
@classmethod
def from_config(cls, config: Dict[str, Any]) -> "Compose":
def from_config(cls, config: dict[str, Any]) -> "Compose":
"""Creates processor from a config dict.
Args:
config (Dict[str, Any]): Configuration dictionary.
config (dict[str, Any]): Configuration dictionary.
Valid keys:
- do_resize
@@ -114,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]):
transforms.append(ConvertToRGB())
@staticmethod
def _get_resize(transforms: List[Transform], config: Dict[str, Any]):
def _get_resize(transforms: list[Transform], config: dict[str, Any]):
mode = config.get("image_processor_type", "CLIPImageProcessor")
if mode == "CLIPImageProcessor":
if config.get("do_resize", False):
@@ -163,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]):
mode = config.get("image_processor_type", "CLIPImageProcessor")
if mode == "CLIPImageProcessor":
if config.get("do_center_crop", False):
@@ -181,18 +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]):
transforms.append(PILtoNDarray())
@staticmethod
def _get_rescale(transforms: List[Transform], config: Dict[str, Any]):
def _get_rescale(transforms: list[Transform], config: dict[str, Any]):
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]):
if config.get("do_normalize", False):
transforms.append(
Normalize(mean=config["image_mean"], std=config["image_std"])
)
transforms.append(Normalize(mean=config["image_mean"], std=config["image_std"]))

20
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, Iterable, Optional, Sequence, Type, Union
import numpy as np
from tokenizers import Encoding
@@ -68,21 +68,21 @@ 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, **kwargs: Any
) -> Dict[str, np.ndarray]:
self, onnx_input: dict[str, np.ndarray], is_doc: bool = True, **kwargs: Any
) -> dict[str, np.ndarray]:
marker_token = self.DOCUMENT_MARKER_TOKEN_ID if is_doc else self.QUERY_MARKER_TOKEN_ID
onnx_input["input_ids"] = np.insert(onnx_input["input_ids"], 1, marker_token, axis=1)
onnx_input["attention_mask"] = np.insert(onnx_input["attention_mask"], 1, 1, axis=1)
return onnx_input
def tokenize(self, documents: List[str], is_doc: bool = True, **kwargs: Any) -> List[Encoding]:
def tokenize(self, documents: list[str], is_doc: bool = True, **kwargs: Any) -> 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, query: str) -> list[Encoding]:
encoded = self.tokenizer.encode_batch([query])
# colbert authors recommend to pad queries with [MASK] tokens for query augmentation to improve performance
if len(encoded[0].ids) < self.MIN_QUERY_LENGTH:
@@ -101,16 +101,16 @@ 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, documents: list[str]) -> list[Encoding]:
encoded = self.tokenizer.encode_batch(documents)
return encoded
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls) -> list[dict[str, Any]]:
"""Lists the supported models.
Returns:
List[Dict[str, Any]]: A list of dictionaries containing the model information.
list[dict[str, Any]]: A list of dictionaries containing the model information.
"""
return supported_colbert_models
@@ -121,7 +121,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
@@ -137,7 +137,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to False.
device_ids (Optional[List[int]], optional): The list of device ids to use for data parallel processing in
device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.

10
fastembed/late_interaction/jina_colbert.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, List, Type
from typing import Any, Type
import numpy as np
@@ -33,17 +33,17 @@ class JinaColbert(Colbert):
return JinaColbertEmbeddingWorker
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls) -> list[dict[str, Any]]:
"""Lists the supported models.
Returns:
List[Dict[str, Any]]: A list of dictionaries containing the model information.
list[dict[str, Any]]: A list of dictionaries containing the model information.
"""
return supported_jina_colbert_models
def _preprocess_onnx_input(
self, onnx_input: Dict[str, np.ndarray], is_doc: bool = True, **kwargs: Any
) -> Dict[str, np.ndarray]:
self, onnx_input: dict[str, np.ndarray], is_doc: bool = True, **kwargs: Any
) -> dict[str, np.ndarray]:
onnx_input = super()._preprocess_onnx_input(onnx_input, is_doc)
# the attention mask for jina-colbert-v2 is always 1 in queries

10
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, Iterable, Optional, Sequence, Type, Union
import numpy as np
@@ -11,15 +11,15 @@ from fastembed.late_interaction.late_interaction_embedding_base import (
class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
EMBEDDINGS_REGISTRY: List[Type[LateInteractionTextEmbeddingBase]] = [Colbert, JinaColbert]
EMBEDDINGS_REGISTRY: list[Type[LateInteractionTextEmbeddingBase]] = [Colbert, JinaColbert]
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls) -> list[dict[str, Any]]:
"""
Lists the supported models.
Returns:
List[Dict[str, Any]]: A list of dictionaries containing the model information.
list[dict[str, Any]]: A list of dictionaries containing the model information.
Example:
```
@@ -50,7 +50,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
**kwargs,
):

15
fastembed/parallel_processor.py
View File

@@ -1,14 +1,15 @@
import logging
import os
from collections import defaultdict
from copy import deepcopy
from enum import Enum
from multiprocessing import Queue, get_context
from multiprocessing.context import BaseContext
from multiprocessing.process import BaseProcess
from multiprocessing.sharedctypes import Synchronized as BaseValue
from queue import Empty
from typing import Any, Dict, Iterable, List, Optional, Tuple, Type
from copy import deepcopy
from typing import Any, Iterable, Optional, Type
# Single item should be processed in less than:
processing_timeout = 10 * 60 # seconds
@@ -27,7 +28,7 @@ class Worker:
def start(cls, *args: Any, **kwargs: Any) -> "Worker":
raise NotImplementedError()
def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
raise NotImplementedError()
@@ -37,7 +38,7 @@ def _worker(
output_queue: Queue,
num_active_workers: BaseValue,
worker_id: int,
kwargs: Optional[Dict[str, Any]] = None,
kwargs: Optional[dict[str, Any]] = None,
) -> None:
"""
A worker that pulls data pints off the input queue, and places the execution result on the output queue.
@@ -93,7 +94,7 @@ class ParallelWorkerPool:
num_workers: int,
worker: Type[Worker],
start_method: Optional[str] = None,
device_ids: Optional[List[int]] = None,
device_ids: Optional[list[int]] = None,
cuda: bool = False,
):
self.worker_class = worker
@@ -101,7 +102,7 @@ class ParallelWorkerPool:
self.input_queue: Optional[Queue] = None
self.output_queue: Optional[Queue] = None
self.ctx: BaseContext = get_context(start_method)
self.processes: List[BaseProcess] = []
self.processes: list[BaseProcess] = []
self.queue_size = self.num_workers * max_internal_batch_size
self.emergency_shutdown = False
self.device_ids = device_ids
@@ -150,7 +151,7 @@ class ParallelWorkerPool:
def semi_ordered_map(
self, stream: Iterable[Any], *args: Any, **kwargs: Any
) -> Iterable[Tuple[int, Any]]:
) -> Iterable[tuple[int, Any]]:
try:
self.start(**kwargs)

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

@@ -1,4 +1,4 @@
from typing import List, Iterable, Dict, Any, Sequence, Optional
from typing import Iterable, Any, Sequence, Optional
from loguru import logger
@@ -73,11 +73,11 @@ supported_onnx_models = [
class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls) -> list[dict[str, Any]]:
"""Lists the supported models.
Returns:
List[Dict[str, Any]]: A list of dictionaries containing the model information.
list[dict[str, Any]]: A list of dictionaries containing the model information.
"""
return supported_onnx_models
@@ -88,7 +88,7 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
@@ -104,7 +104,7 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to False.
device_ids (Optional[List[int]], optional): The list of device ids to use for data parallel processing in
device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.

12
fastembed/rerank/cross_encoder/onnx_text_model.py
View File

@@ -1,4 +1,4 @@
from typing import Sequence, Optional, List, Dict, Iterable
from typing import Sequence, Optional, Iterable
from pathlib import Path
import numpy as np
@@ -10,7 +10,7 @@ from fastembed.common.utils import iter_batch
class OnnxCrossEncoderModel(OnnxModel):
ONNX_OUTPUT_NAMES: Optional[List[str]] = None
ONNX_OUTPUT_NAMES: Optional[list[str]] = None
def _load_onnx_model(
self,
@@ -31,10 +31,10 @@ class OnnxCrossEncoderModel(OnnxModel):
)
self.tokenizer, _ = load_tokenizer(model_dir=model_dir)
def tokenize(self, query: str, documents: List[str], **kwargs) -> List[Encoding]:
def tokenize(self, query: str, documents: list[str], **kwargs) -> list[Encoding]:
return self.tokenizer.encode_batch([(query, doc) for doc in documents])
def onnx_embed(self, query: str, documents: List[str], **kwargs) -> OnnxOutputContext:
def onnx_embed(self, query: str, documents: list[str], **kwargs) -> OnnxOutputContext:
tokenized_input = self.tokenize(query, documents, **kwargs)
inputs = {
@@ -62,8 +62,8 @@ class OnnxCrossEncoderModel(OnnxModel):
yield from self.onnx_embed(query, batch, **kwargs).model_output
def _preprocess_onnx_input(
self, onnx_input: Dict[str, np.ndarray], **kwargs
) -> Dict[str, np.ndarray]:
self, onnx_input: dict[str, np.ndarray], **kwargs
) -> dict[str, np.ndarray]:
"""
Preprocess the onnx input.
"""

10
fastembed/rerank/cross_encoder/text_cross_encoder.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type
from typing import Any, Iterable, Optional, Sequence, Type
from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase
from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
@@ -6,16 +6,16 @@ from fastembed.common import OnnxProvider
class TextCrossEncoder(TextCrossEncoderBase):
CROSS_ENCODER_REGISTRY: List[Type[TextCrossEncoderBase]] = [
CROSS_ENCODER_REGISTRY: list[Type[TextCrossEncoderBase]] = [
OnnxTextCrossEncoder,
]
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls) -> list[dict[str, Any]]:
"""Lists the supported models.
Returns:
List[Dict[str, Any]]: A list of dictionaries containing the model information.
list[dict[str, Any]]: A list of dictionaries containing the model information.
Example:
```
@@ -45,7 +45,7 @@ class TextCrossEncoder(TextCrossEncoderBase):
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
**kwargs,
):

22
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, Iterable, Optional, Type, Union
import mmh3
import numpy as np
@@ -133,16 +133,16 @@ class Bm25(SparseTextEmbeddingBase):
self.tokenizer = SimpleTokenizer
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls) -> list[dict[str, Any]]:
"""Lists the supported models.
Returns:
List[Dict[str, Any]]: A list of dictionaries containing the model information.
list[dict[str, Any]]: A list of dictionaries containing the model information.
"""
return supported_bm25_models
@classmethod
def _load_stopwords(cls, model_dir: Path, language: str) -> List[str]:
def _load_stopwords(cls, model_dir: Path, language: str) -> list[str]:
stopwords_path = model_dir / f"{language}.txt"
if not stopwords_path.exists():
return []
@@ -222,7 +222,7 @@ class Bm25(SparseTextEmbeddingBase):
parallel=parallel,
)
def _stem(self, tokens: List[str]) -> List[str]:
def _stem(self, tokens: list[str]) -> list[str]:
stemmed_tokens = []
for token in tokens:
if token in self.punctuation:
@@ -242,8 +242,8 @@ class Bm25(SparseTextEmbeddingBase):
def raw_embed(
self,
documents: List[str],
) -> List[SparseEmbedding]:
documents: list[str],
) -> list[SparseEmbedding]:
embeddings = []
for document in documents:
document = remove_non_alphanumeric(document)
@@ -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, tokens: list[str]) -> dict[int, float]:
"""Calculate the term frequency part of the BM25 formula.
(
@@ -263,10 +263,10 @@ class Bm25(SparseTextEmbeddingBase):
)
Args:
tokens (List[str]): The list of tokens in the document.
tokens (list[str]): The list of tokens in the document.
Returns:
Dict[int, float]: The token_id to term frequency mapping.
dict[int, float]: The token_id to term frequency mapping.
"""
tf_map = {}
counter = defaultdict(int)
@@ -323,7 +323,7 @@ class Bm25Worker(Worker):
def start(cls, 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, 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

28
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, Iterable, Optional, Sequence, Type, Union
import mmh3
import numpy as np
@@ -63,7 +63,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
providers: Optional[Sequence[OnnxProvider]] = None,
alpha: float = 0.5,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
@@ -81,7 +81,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
It is recommended to only change this parameter based on training data for a specific dataset.
cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to False.
device_ids (Optional[List[int]], optional): The list of device ids to use for data parallel processing in
device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
@@ -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, 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, tokens: list[tuple[str, Any]]) -> list[tuple[str, Any]]:
result = []
for token, value in tokens:
processed_token = self.stemmer.stem_word(token)
@@ -158,8 +158,8 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
@classmethod
def _aggregate_weights(
cls, tokens: List[Tuple[str, List[int]]], weights: List[float]
) -> List[Tuple[str, float]]:
cls, tokens: list[tuple[str, list[int]]], weights: list[float]
) -> list[tuple[str, float]]:
result = []
for token, idxs in tokens:
sum_weight = sum(weights[idx] for idx in idxs)
@@ -167,8 +167,8 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
return result
def _reconstruct_bpe(
self, bpe_tokens: Iterable[Tuple[int, str]]
) -> List[Tuple[str, List[int]]]:
self, bpe_tokens: Iterable[tuple[int, str]]
) -> list[tuple[str, list[int]]]:
result = []
acc = ""
acc_idx = []
@@ -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, 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.
@@ -246,16 +246,16 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
yield SparseEmbedding.from_dict(rescored)
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls) -> list[dict[str, Any]]:
"""Lists the supported models.
Returns:
List[Dict[str, Any]]: A list of dictionaries containing the model information.
list[dict[str, Any]]: A list of dictionaries containing the model information.
"""
return supported_bm42_models
@classmethod
def _load_stopwords(cls, model_dir: Path) -> List[str]:
def _load_stopwords(cls, model_dir: Path) -> list[str]:
stopwords_path = model_dir / "stopwords.txt"
if not stopwords_path.exists():
return []
@@ -298,7 +298,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
)
@classmethod
def _query_rehash(cls, tokens: Iterable[str]) -> Dict[int, float]:
def _query_rehash(cls, tokens: Iterable[str]) -> dict[int, float]:
result = {}
for token in tokens:
token_id = abs(mmh3.hash(token))

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 Iterable, Optional, Union
import numpy as np
@@ -11,17 +11,17 @@ class SparseEmbedding:
values: np.ndarray
indices: np.ndarray
def as_object(self) -> Dict[str, np.ndarray]:
def as_object(self) -> dict[str, np.ndarray]:
return {
"values": self.values,
"indices": self.indices,
}
def as_dict(self) -> Dict[int, float]:
def as_dict(self) -> dict[int, float]:
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, data: dict[int, float]) -> "SparseEmbedding":
if len(data) == 0:
return cls(values=np.array([]), indices=np.array([]))
indices, values = zip(*data.items())
@@ -50,9 +50,7 @@ class SparseTextEmbeddingBase(ModelManagement):
) -> Iterable[SparseEmbedding]:
raise NotImplementedError()
def passage_embed(
self, texts: Iterable[str], **kwargs
) -> Iterable[SparseEmbedding]:
def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[SparseEmbedding]:
"""
Embeds a list of text passages into a list of embeddings.
@@ -67,9 +65,7 @@ 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) -> Iterable[SparseEmbedding]:
"""
Embeds queries

10
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, Iterable, Optional, Sequence, Type, Union
from fastembed.common import OnnxProvider
from fastembed.sparse.bm25 import Bm25
@@ -12,15 +12,15 @@ import warnings
class SparseTextEmbedding(SparseTextEmbeddingBase):
EMBEDDINGS_REGISTRY: List[Type[SparseTextEmbeddingBase]] = [SpladePP, Bm42, Bm25]
EMBEDDINGS_REGISTRY: list[Type[SparseTextEmbeddingBase]] = [SpladePP, Bm42, Bm25]
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls) -> list[dict[str, Any]]:
"""
Lists the supported models.
Returns:
List[Dict[str, Any]]: A list of dictionaries containing the model information.
list[dict[str, Any]]: A list of dictionaries containing the model information.
Example:
```
@@ -50,7 +50,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
**kwargs,
):

10
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, Iterable, Optional, Sequence, Type, Union
import numpy as np
from fastembed.common import OnnxProvider
@@ -55,11 +55,11 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
yield SparseEmbedding(values=scores, indices=indices)
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls) -> list[dict[str, Any]]:
"""Lists the supported models.
Returns:
List[Dict[str, Any]]: A list of dictionaries containing the model information.
list[dict[str, Any]]: A list of dictionaries containing the model information.
"""
return supported_splade_models
@@ -70,7 +70,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
@@ -86,7 +86,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to False.
device_ids (Optional[List[int]], optional): The list of device ids to use for data parallel processing in
device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.

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

@@ -1,12 +1,11 @@
# This code is a modified copy of the `NLTKWordTokenizer` class from `NLTK` library.
import re
from typing import List
class SimpleTokenizer:
@staticmethod
def tokenize(text: str) -> List[str]:
def tokenize(text: str) -> list[str]:
text = re.sub(r"[^\w]", " ", text.lower())
text = re.sub(r"\s+", " ", text)
@@ -81,7 +80,7 @@ class WordTokenizer:
]
@classmethod
def tokenize(cls, text: str) -> List[str]:
def tokenize(cls, text: str) -> list[str]:
"""Return a tokenized copy of `text`.
>>> s = '''Good muffins cost $3.88 (roughly 3,36 euros)\nin New York.'''

6
fastembed/text/clip_embedding.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, Iterable, List, Type
from typing import Any, Iterable, Type
import numpy as np
@@ -27,11 +27,11 @@ class CLIPOnnxEmbedding(OnnxTextEmbedding):
return CLIPEmbeddingWorker
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls) -> list[dict[str, Any]]:
"""Lists the supported models.
Returns:
List[Dict[str, Any]]: A list of dictionaries containing the model information.
list[dict[str, Any]]: A list of dictionaries containing the model information.
"""
return supported_clip_models

10
fastembed/text/e5_onnx_embedding.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, List, Type
from typing import Any, Type
import numpy as np
@@ -39,17 +39,17 @@ class E5OnnxEmbedding(OnnxTextEmbedding):
return E5OnnxEmbeddingWorker
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls) -> list[dict[str, Any]]:
"""Lists the supported models.
Returns:
List[Dict[str, Any]]: A list of dictionaries containing the model information.
list[dict[str, Any]]: A list of dictionaries containing the model information.
"""
return supported_multilingual_e5_models
def _preprocess_onnx_input(
self, onnx_input: Dict[str, np.ndarray], **kwargs
) -> Dict[str, np.ndarray]:
self, onnx_input: dict[str, np.ndarray], **kwargs
) -> dict[str, np.ndarray]:
"""
Preprocess the onnx input.
"""

14
fastembed/text/onnx_embedding.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union
from typing import Any, Iterable, Optional, Sequence, Type, Union
import numpy as np
@@ -171,12 +171,12 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
"""Implementation of the Flag Embedding model."""
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls) -> list[dict[str, Any]]:
"""
Lists the supported models.
Returns:
List[Dict[str, Any]]: A list of dictionaries containing the model information.
list[dict[str, Any]]: A list of dictionaries containing the model information.
"""
return supported_onnx_models
@@ -187,7 +187,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
@@ -203,7 +203,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to False.
device_ids (Optional[List[int]], optional): The list of device ids to use for data parallel processing in
device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
@@ -276,8 +276,8 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
return OnnxTextEmbeddingWorker
def _preprocess_onnx_input(
self, onnx_input: Dict[str, np.ndarray], **kwargs
) -> Dict[str, np.ndarray]:
self, onnx_input: dict[str, np.ndarray], **kwargs
) -> dict[str, np.ndarray]:
"""
Preprocess the onnx input.
"""

16
fastembed/text/onnx_text_model.py
View File

@@ -1,7 +1,7 @@
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, Iterable, Optional, Sequence, Type, Union
import numpy as np
from tokenizers import Encoding
@@ -14,7 +14,7 @@ from fastembed.parallel_processor import ParallelWorkerPool
class OnnxTextModel(OnnxModel[T]):
ONNX_OUTPUT_NAMES: Optional[List[str]] = None
ONNX_OUTPUT_NAMES: Optional[list[str]] = None
@classmethod
def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
@@ -29,8 +29,8 @@ class OnnxTextModel(OnnxModel[T]):
self.special_token_to_id = {}
def _preprocess_onnx_input(
self, onnx_input: Dict[str, np.ndarray], **kwargs
) -> Dict[str, np.ndarray]:
self, onnx_input: dict[str, np.ndarray], **kwargs
) -> dict[str, np.ndarray]:
"""
Preprocess the onnx input.
"""
@@ -58,12 +58,12 @@ 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) -> list[Encoding]:
return self.tokenizer.encode_batch(documents)
def onnx_embed(
self,
documents: List[str],
documents: list[str],
**kwargs,
) -> OnnxOutputContext:
encoded = self.tokenize(documents, **kwargs)
@@ -98,7 +98,7 @@ class OnnxTextModel(OnnxModel[T]):
parallel: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
device_ids: Optional[list[int]] = None,
**kwargs,
) -> Iterable[T]:
is_small = False
@@ -140,7 +140,7 @@ class OnnxTextModel(OnnxModel[T]):
class TextEmbeddingWorker(EmbeddingWorker):
def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
def process(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

6
fastembed/text/pooled_embedding.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, Iterable, List, Type
from typing import Any, Iterable, Type
import numpy as np
@@ -60,11 +60,11 @@ class PooledEmbedding(OnnxTextEmbedding):
return pooled_embeddings
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls) -> list[dict[str, Any]]:
"""Lists the supported models.
Returns:
List[Dict[str, Any]]: A list of dictionaries containing the model information.
list[dict[str, Any]]: A list of dictionaries containing the model information.
"""
return supported_pooled_models

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, Iterable, Type
import numpy as np
@@ -66,11 +66,11 @@ class PooledNormalizedEmbedding(PooledEmbedding):
return PooledNormalizedEmbeddingWorker
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls) -> list[dict[str, Any]]:
"""Lists the supported models.
Returns:
List[Dict[str, Any]]: A list of dictionaries containing the model information.
list[dict[str, Any]]: A list of dictionaries containing the model information.
"""
return supported_pooled_normalized_models

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, Iterable, Optional, Sequence, Type, Union
import numpy as np
@@ -12,7 +12,7 @@ from fastembed.text.text_embedding_base import TextEmbeddingBase
class TextEmbedding(TextEmbeddingBase):
EMBEDDINGS_REGISTRY: List[Type[TextEmbeddingBase]] = [
EMBEDDINGS_REGISTRY: list[Type[TextEmbeddingBase]] = [
OnnxTextEmbedding,
E5OnnxEmbedding,
CLIPOnnxEmbedding,
@@ -21,12 +21,12 @@ class TextEmbedding(TextEmbeddingBase):
]
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
def list_supported_models(cls) -> list[dict[str, Any]]:
"""
Lists the supported models.
Returns:
List[Dict[str, Any]]: A list of dictionaries containing the model information.
list[dict[str, Any]]: A list of dictionaries containing the model information.
Example:
```
@@ -57,7 +57,7 @@ class TextEmbedding(TextEmbeddingBase):
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
**kwargs,
):

4
pyproject.toml
View File

@@ -11,7 +11,7 @@ repository = "https://github.com/qdrant/fastembed"
keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-transformers"]
[tool.poetry.dependencies]
python = ">=3.8.0,<3.13"
python = ">=3.9.0,<3.13"
onnx = "^1.15.0"
onnxruntime = ">=1.17.0,<1.20.0"
tqdm = "^4.66"
@@ -31,7 +31,7 @@ py-rust-stemmers = "^0.1.0"
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" }
pre-commit = "^3.6.2"
[tool.poetry.group.docs.dependencies]
mkdocs-material = "^9.5.10"