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: Add sparse type hints (#460)

Browse Source 
* new: Add sparse type hints

* fix: ndarray -> numpyarray

---------

Co-authored-by: George Panchuk <george.panchuk@qdrant.tech>
This commit is contained in:
Hossam Hagag
2025-02-04 18:17:25 +02:00
committed by GitHub
parent b08febbf93
commit 37a66d9e16
5 changed files with 30 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
fastembed/sparse/bm25.py
View File

@@ -123,7 +123,7 @@ class Bm25(SparseTextEmbeddingBase):
self.avg_len = avg_len
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(
model_description,
@@ -137,7 +137,7 @@ class Bm25(SparseTextEmbeddingBase):
self.disable_stemmer = disable_stemmer
if disable_stemmer:
self.stopwords = set()
self.stopwords: set[str] = set()
self.stemmer = None
else:
self.stopwords = set(self._load_stopwords(self._model_dir, self.language))
@@ -239,7 +239,7 @@ class Bm25(SparseTextEmbeddingBase):
)
def _stem(self, tokens: list[str]) -> list[str]:
stemmed_tokens = []
stemmed_tokens: list[str] = []
for token in tokens:
lower_token = token.lower()
@@ -262,7 +262,7 @@ class Bm25(SparseTextEmbeddingBase):
self,
documents: list[str],
) -> list[SparseEmbedding]:
embeddings = []
embeddings: list[SparseEmbedding] = []
for document in documents:
document = remove_non_alphanumeric(document)
tokens = self.tokenizer.tokenize(document)
@@ -286,8 +286,8 @@ class Bm25(SparseTextEmbeddingBase):
Returns:
dict[int, float]: The token_id to term frequency mapping.
"""
tf_map = {}
counter = defaultdict(int)
tf_map: dict[int, float] = {}
counter: defaultdict[str, int] = defaultdict(int)
for stemmed_token in tokens:
counter[stemmed_token] += 1

30
fastembed/sparse/bm42.py
View File

@@ -110,7 +110,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
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,
@@ -119,10 +119,10 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
specific_model_path=specific_model_path,
)
self.invert_vocab = {}
self.invert_vocab: dict[int, str] = {}
self.special_tokens = set()
self.special_tokens_ids = set()
self.special_tokens: set[str] = set()
self.special_tokens_ids: set[int] = set()
self.punctuation = set(string.punctuation)
self.stopwords = set(self._load_stopwords(self._model_dir))
self.stemmer = SnowballStemmer(MODEL_TO_LANGUAGE[model_name])
@@ -147,7 +147,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
self.stopwords = set(self._load_stopwords(self._model_dir))
def _filter_pair_tokens(self, tokens: list[tuple[str, Any]]) -> list[tuple[str, Any]]:
result = []
result: list[tuple[str, Any]] = []
for token, value in tokens:
if token in self.stopwords or token in self.punctuation:
continue
@@ -155,7 +155,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
return result
def _stem_pair_tokens(self, tokens: list[tuple[str, Any]]) -> list[tuple[str, Any]]:
result = []
result: list[tuple[str, Any]] = []
for token, value in tokens:
processed_token = self.stemmer.stem_word(token)
result.append((processed_token, value))
@@ -165,7 +165,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
def _aggregate_weights(
cls, tokens: list[tuple[str, list[int]]], weights: list[float]
) -> list[tuple[str, float]]:
result = []
result: list[tuple[str, float]] = []
for token, idxs in tokens:
sum_weight = sum(weights[idx] for idx in idxs)
result.append((token, sum_weight))
@@ -174,9 +174,9 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
def _reconstruct_bpe(
self, bpe_tokens: Iterable[tuple[int, str]]
) -> list[tuple[str, list[int]]]:
result = []
acc = ""
acc_idx = []
result: list[tuple[str, list[int]]] = []
acc: str = ""
acc_idx: list[int] = []
continuing_subword_prefix = self.tokenizer.model.continuing_subword_prefix
continuing_subword_prefix_len = len(continuing_subword_prefix)
@@ -206,7 +206,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
So that the scoring doesn't depend on absolute values assigned by the model, but on the relative importance.
"""
new_vector = {}
new_vector: dict[int, float] = {}
for token, value in vector.items():
token_id = abs(mmh3.hash(token))
@@ -241,7 +241,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
weighted = self._aggregate_weights(stemmed, attention_value)
max_token_weight = {}
max_token_weight: dict[str, float] = {}
for token, weight in weighted:
max_token_weight[token] = max(max_token_weight.get(token, 0), weight)
@@ -304,7 +304,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
@classmethod
def _query_rehash(cls, tokens: Iterable[str]) -> dict[int, float]:
result = {}
result: dict[int, float] = {}
for token in tokens:
token_id = abs(mmh3.hash(token))
result[token_id] = 1.0
@@ -334,11 +334,11 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
yield SparseEmbedding.from_dict(self._query_rehash(token for token, _ in stemmed))
@classmethod
def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
def _get_worker_class(cls) -> Type[TextEmbeddingWorker[SparseEmbedding]]:
return Bm42TextEmbeddingWorker
class Bm42TextEmbeddingWorker(TextEmbeddingWorker):
class Bm42TextEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> Bm42:
return Bm42(
model_name=model_name,

9
fastembed/sparse/sparse_embedding_base.py
View File

@@ -3,15 +3,16 @@ from typing import Iterable, Optional, Union, Any
import numpy as np
from fastembed.common.types import NumpyArray
from fastembed.common.model_management import ModelManagement
@dataclass
class SparseEmbedding:
values: np.ndarray
indices: np.ndarray
values: NumpyArray
indices: NumpyArray
def as_object(self) -> dict[str, np.ndarray]:
def as_object(self) -> dict[str, NumpyArray]:
return {
"values": self.values,
"indices": self.indices,
@@ -81,5 +82,5 @@ class SparseTextEmbeddingBase(ModelManagement):
# 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
fastembed/sparse/sparse_text_embedding.py
View File

@@ -38,7 +38,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
]
```
"""
result = []
result: list[dict[str, Any]] = []
for embedding in cls.EMBEDDINGS_REGISTRY:
result.extend(embedding.list_supported_models())
return result

6
fastembed/sparse/splade_pp.py
View File

@@ -114,7 +114,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
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,
@@ -171,11 +171,11 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
)
@classmethod
def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
def _get_worker_class(cls) -> Type[TextEmbeddingWorker[SparseEmbedding]]:
return SpladePPEmbeddingWorker
class SpladePPEmbeddingWorker(TextEmbeddingWorker):
class SpladePPEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]):
def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> SpladePP:
return SpladePP(
model_name=model_name,