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

Attention sparse embeddings (#235)

Browse Source 
* WIP: sparse embeddings using attention

* support for stopwords

* apply stopwords

* proceed implementation of sparse attention embeddings (#234)

* complete inference

* query embed + comment

* use simpler weights formula instead of sorting of words

* update tests

* fix: fix bm42 usage, add query_embed to SparseTextEmbedding, update tests

---------

Co-authored-by: George <george.panchuk@qdrant.tech>
This commit is contained in:
Andrey Vasnetsov
2024-05-24 15:25:34 +02:00
committed by GitHub
parent 316c33634b
commit dfd25d41c9
15 changed files with 486 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
fastembed/common/onnx_model.py
View File

@@ -1,3 +1,4 @@
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Generic, Iterable, Optional, Tuple, Type, TypeVar, Sequence
import warnings
@@ -13,13 +14,19 @@ from fastembed.parallel_processor import Worker
T = TypeVar("T")
@dataclass
class OnnxOutputContext:
model_output: np.ndarray
attention_mask: Optional[np.ndarray] = None
input_ids: Optional[np.ndarray] = None
class OnnxModel(Generic[T]):
@classmethod
def _get_worker_class(cls) -> Type["EmbeddingWorker"]:
raise NotImplementedError("Subclasses must implement this method")
@classmethod
def _post_process_onnx_output(cls, output: Tuple[np.ndarray, np.ndarray]) -> Iterable[T]:
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
raise NotImplementedError("Subclasses must implement this method")
def __init__(self) -> None:
@@ -74,7 +81,7 @@ class OnnxModel(Generic[T]):
RuntimeWarning,
)
def onnx_embed(self, *args, **kwargs) -> Tuple[np.ndarray, np.ndarray]:
def onnx_embed(self, *args, **kwargs) -> OnnxOutputContext:
raise NotImplementedError("Subclasses must implement this method")

32
fastembed/common/preprocessor_utils.py
View File

@@ -1,12 +1,24 @@
import json
from pathlib import Path
from typing import Tuple
from tokenizers import Tokenizer, AddedToken
from fastembed.image.transform.operators import Compose
def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
def load_special_tokens(model_dir: Path) -> dict:
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}")
with open(str(tokens_map_path)) as tokens_map_file:
tokens_map = json.load(tokens_map_file)
return tokens_map
def load_tokenizer(model_dir: Path, max_length: int = 512) -> 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}")
@@ -19,18 +31,13 @@ def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
if not tokenizer_config_path.exists():
raise ValueError(f"Could not find tokenizer_config.json in {model_dir}")
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}")
with open(str(config_path)) as config_file:
config = json.load(config_file)
with open(str(tokenizer_config_path)) as tokenizer_config_file:
tokenizer_config = json.load(tokenizer_config_file)
with open(str(tokens_map_path)) as tokens_map_file:
tokens_map = json.load(tokens_map_file)
tokens_map = load_special_tokens(model_dir)
tokenizer = Tokenizer.from_file(str(tokenizer_path))
tokenizer.enable_truncation(max_length=min(tokenizer_config["model_max_length"], max_length))
@@ -44,7 +51,16 @@ def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
elif isinstance(token, dict):
tokenizer.add_special_tokens([AddedToken(**token)])
return tokenizer
special_token_to_id = {}
for token in tokens_map.values():
if isinstance(token, str):
special_token_to_id[token] = tokenizer.token_to_id(token)
elif isinstance(token, dict):
token_str = token.get("content", "")
special_token_to_id[token_str] = tokenizer.token_to_id(token_str)
return tokenizer, special_token_to_id
def load_preprocessor(model_dir: Path) -> Compose:

6
fastembed/image/onnx_embedding.py
View File

@@ -2,6 +2,7 @@ from typing import Dict, Optional, Iterable, Type, List, Any, Sequence
import numpy as np
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common.utils import normalize, define_cache_dir
from fastembed.common import ImageInput, OnnxProvider
from fastembed.image.image_embedding_base import ImageEmbeddingBase
@@ -108,9 +109,8 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
return onnx_input
@classmethod
def _post_process_onnx_output(cls, output: np.ndarray) -> Iterable[np.ndarray]:
return normalize(output).astype(np.float32)
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
return normalize(output.model_output).astype(np.float32)
class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):

11
fastembed/image/onnx_image_model.py
View File

@@ -8,7 +8,7 @@ from PIL import Image
import numpy as np
from fastembed.common.preprocessor_utils import load_preprocessor
from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker, T
from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker, T, OnnxOutputContext
from fastembed.common import PathInput, ImageInput, OnnxProvider
from fastembed.common.utils import iter_batch
from fastembed.parallel_processor import ParallelWorkerPool
@@ -21,8 +21,7 @@ class OnnxImageModel(OnnxModel[T]):
def _get_worker_class(cls) -> Type["ImageEmbeddingWorker"]:
raise NotImplementedError("Subclasses must implement this method")
@classmethod
def _post_process_onnx_output(cls, output: np.ndarray) -> Iterable[T]:
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
raise NotImplementedError("Subclasses must implement this method")
def __init__(self) -> None:
@@ -47,7 +46,7 @@ class OnnxImageModel(OnnxModel[T]):
)
self.processor = load_preprocessor(model_dir=model_dir)
def onnx_embed(self, images: List[PathInput]) -> np.ndarray:
def onnx_embed(self, images: List[PathInput]) -> OnnxOutputContext:
with contextlib.ExitStack():
image_files = [Image.open(image) for image in images]
encoded = self.processor(image_files)
@@ -56,7 +55,9 @@ class OnnxImageModel(OnnxModel[T]):
model_output = self.model.run(None, onnx_input)
embeddings = model_output[0]
return embeddings
return OnnxOutputContext(
model_output=embeddings
)
def _embed_images(
self,

281
fastembed/sparse/bm42.py Normal file
View File

@@ -0,0 +1,281 @@
import math
import string
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type, Sequence
import numpy as np
import mmh3
from snowballstemmer import stemmer as get_stemmer
from fastembed.common import OnnxProvider
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common.utils import define_cache_dir
from fastembed.sparse.sparse_embedding_base import SparseEmbedding, SparseTextEmbeddingBase
from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
supported_bm42_models = [
{
"model": "Qdrant/bm42-all-minilm-l6-v2-attentions",
"vocab_size": 30522,
"description": "Light sparse embedding model, which assigns an importance score to each token in the text",
"size_in_GB": 0.09,
"sources": {
"hf": "Qdrant/all_miniLM_L6_v2_with_attentions",
},
"model_file": "model.onnx",
"additional_files": ["stopwords.txt"],
},
]
MODEL_TO_LANGUAGE = {
"Qdrant/bm42-all-minilm-l6-v2-attentions": "english",
}
class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
"""
Bm42 is an extension of BM25, which tries to better evaluate importance of tokens in the documents,
by extracting attention weights from the transformer model.
Traditional BM25 uses a count of tokens in the document to evaluate the importance of the token,
but this approach doesn't work well with short documents or chunks of text, as almost all tokens
there are unique.
BM42 addresses this issue by replacing the token count with the attention weights from the transformer model.
This allows sparse embeddings to work well with short documents, handle rare tokens and leverage traditional NLP
techniques like stemming and stopwords.
WARNING: This model is expected to be used with `modifier="idf"` in the sparse vector index of Qdrant.
"""
ONNX_OUTPUT_NAMES = ["attention_6"]
def __init__(
self,
model_name: str,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
alpha: float = 0.5,
**kwargs,
):
"""
Args:
model_name (str): The name of the model to use.
cache_dir (str, optional): The path to the cache directory.
Can be set using the `FASTEMBED_CACHE_PATH` env variable.
Defaults to `fastembed_cache` in the system's temp directory.
threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
providers (Optional[Sequence[OnnxProvider]], optional): The providers to use for onnxruntime.
alpha (float, optional): Parameter, that defines the importance of the token weight in the document
versus the importance of the token frequency in the corpus. Defaults to 0.5, based on empirical testing.
It is recommended to only change this parameter based on training data for a specific dataset.
Raises:
ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
"""
super().__init__(model_name, cache_dir, threads, **kwargs)
model_description = self._get_model_description(model_name)
cache_dir = define_cache_dir(cache_dir)
model_dir = self.download_model(
model_description, cache_dir, local_files_only=self._local_files_only
)
self.load_onnx_model(
model_dir=model_dir,
model_file=model_description["model_file"],
threads=threads,
providers=providers,
)
self.invert_vocab = {}
for token, idx in self.tokenizer.get_vocab().items():
self.invert_vocab[idx] = token
self.special_tokens = set(self.special_token_to_id.keys())
self.special_tokens_ids = set(self.special_token_to_id.values())
self.punctuation = set(string.punctuation)
self.stopwords = set(self._load_stopwords(model_dir))
self.stemmer = get_stemmer(MODEL_TO_LANGUAGE[model_name])
self.alpha = alpha
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:
continue
result.append((token, value))
return result
def _stem_pair_tokens(self, tokens: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
result = []
for token, value in tokens:
processed_token = self.stemmer.stemWord(token)
result.append((processed_token, value))
return result
@classmethod
def _aggregate_weights(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)
result.append((token, sum_weight))
return result
def _reconstruct_bpe(
self, bpe_tokens: Iterable[Tuple[int, str]]
) -> List[Tuple[str, List[int]]]:
result = []
acc = ""
acc_idx = []
continuing_subword_prefix = self.tokenizer.model.continuing_subword_prefix
continuing_subword_prefix_len = len(continuing_subword_prefix)
for idx, token in bpe_tokens:
if token in self.special_tokens:
continue
if token.startswith(continuing_subword_prefix):
acc += token[continuing_subword_prefix_len:]
acc_idx.append(idx)
else:
if acc:
result.append((acc, acc_idx))
acc_idx = []
acc = token
acc_idx.append(idx)
if acc:
result.append((acc, acc_idx))
return result
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.
"""
new_vector = {}
for token, value in vector.items():
token_id = abs(mmh3.hash(token))
# Examples:
# Num 0: Log(1/1 + 1) = 0.6931471805599453
# Num 1: Log(1/2 + 1) = 0.4054651081081644
# Num 2: Log(1/3 + 1) = 0.28768207245178085
new_vector[token_id] = math.log(1. + value) ** self.alpha # value
return new_vector
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[SparseEmbedding]:
token_ids_batch = output.input_ids
# attention_value shape: (batch_size, num_heads, num_tokens, num_tokens)
pooled_attention = np.mean(output.model_output[:, :, 0], axis=1) * output.attention_mask
for document_token_ids, attention_value in zip(token_ids_batch, pooled_attention):
document_tokens_with_ids = (
(idx, self.invert_vocab[token_id])
for idx, token_id in enumerate(document_token_ids)
)
reconstructed = self._reconstruct_bpe(document_tokens_with_ids)
filtered = self._filter_pair_tokens(reconstructed)
stemmed = self._stem_pair_tokens(filtered)
weighted = self._aggregate_weights(stemmed, attention_value)
max_token_weight = {}
for token, weight in weighted:
max_token_weight[token] = max(max_token_weight.get(token, 0), weight)
rescored = self._rescore_vector(max_token_weight)
yield SparseEmbedding.from_dict(rescored)
@classmethod
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.
"""
return supported_bm42_models
@classmethod
def _load_stopwords(cls, model_dir: Path) -> List[str]:
stopwords_path = model_dir / "stopwords.txt"
if not stopwords_path.exists():
return []
with open(stopwords_path, "r") as f:
return f.read().splitlines()
def embed(
self,
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
) -> Iterable[SparseEmbedding]:
"""
Encode a list of documents into list of embeddings.
We use mean pooling with attention so that the model can handle variable-length inputs.
Args:
documents: Iterator of documents or single document to embed
batch_size: Batch size for encoding -- higher values will use more memory, but be faster
parallel:
If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
If 0, use all available cores.
If None, don't use data-parallel processing, use default onnxruntime threading instead.
Returns:
List of embeddings, one per document
"""
yield from self._embed_documents(
model_name=self.model_name,
cache_dir=str(self.cache_dir),
documents=documents,
batch_size=batch_size,
parallel=parallel,
)
@classmethod
def _query_rehash(cls, tokens: Iterable[str]) -> Dict[int, float]:
result = {}
for token in tokens:
token_id = abs(mmh3.hash(token))
result[token_id] = 1.0
return result
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> 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.
It is also faster, as we don't need to run the model for the query.
"""
if isinstance(query, str):
query = [query]
for text in query:
encoded = self.tokenizer.encode(text)
document_tokens_with_ids = enumerate(encoded.tokens)
reconstructed = self._reconstruct_bpe(document_tokens_with_ids)
filtered = self._filter_pair_tokens(reconstructed)
stemmed = self._stem_pair_tokens(filtered)
yield SparseEmbedding.from_dict(self._query_rehash(token for token, _ in stemmed))
@classmethod
def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
return TextEmbeddingWorker

57
fastembed/sparse/sparse_embedding_base.py
View File

@@ -20,14 +20,19 @@ class SparseEmbedding:
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":
indices, values = zip(*data.items())
return cls(values=np.array(values), indices=np.array(indices))
class SparseTextEmbeddingBase(ModelManagement):
def __init__(
self,
model_name: str,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
**kwargs,
self,
model_name: str,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
**kwargs,
):
self.model_name = model_name
self.cache_dir = cache_dir
@@ -35,10 +40,42 @@ class SparseTextEmbeddingBase(ModelManagement):
self._local_files_only = kwargs.pop("local_files_only", False)
def embed(
self,
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
self,
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
) -> Iterable[SparseEmbedding]:
raise NotImplementedError()
def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[SparseEmbedding]:
"""
Embeds a list of text passages into a list of embeddings.
Args:
texts (Iterable[str]): The list of texts to embed.
**kwargs: Additional keyword argument to pass to the embed method.
Yields:
Iterable[SparseEmbedding]: The sparse 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) -> Iterable[SparseEmbedding]:
"""
Embeds queries
Args:
query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
Returns:
Iterable[SparseEmbedding]: The sparse 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):
yield from self.embed(query, **kwargs)

14
fastembed/sparse/sparse_text_embedding.py
View File

@@ -1,6 +1,7 @@
from typing import List, Type, Dict, Any, Union, Iterable, Optional, Sequence
from fastembed.common import OnnxProvider
from fastembed.sparse.bm42 import Bm42
from fastembed.sparse.sparse_embedding_base import SparseTextEmbeddingBase, SparseEmbedding
from fastembed.sparse.splade_pp import SpladePP
@@ -8,6 +9,7 @@ from fastembed.sparse.splade_pp import SpladePP
class SparseTextEmbedding(SparseTextEmbeddingBase):
EMBEDDINGS_REGISTRY: List[Type[SparseTextEmbeddingBase]] = [
SpladePP,
Bm42,
]
@classmethod
@@ -84,3 +86,15 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
List of embeddings, one per document
"""
yield from self.model.embed(documents, batch_size, parallel, **kwargs)
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]:
"""
Embeds queries
Args:
query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
Returns:
Iterable[SparseEmbedding]: The sparse embeddings.
"""
yield from self.model.query_embed(query, **kwargs)

9
fastembed/sparse/splade_pp.py
View File

@@ -3,6 +3,7 @@ from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type, Sequ
import numpy as np
from fastembed.common import OnnxProvider
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common.utils import define_cache_dir
from fastembed.sparse.sparse_embedding_base import SparseEmbedding, SparseTextEmbeddingBase
from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
@@ -33,14 +34,12 @@ supported_splade_models = [
class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
@classmethod
def _post_process_onnx_output(
cls, output: Tuple[np.ndarray, np.ndarray]
self, output: OnnxOutputContext
) -> Iterable[SparseEmbedding]:
logits, attention_mask = output
relu_log = np.log(1 + np.maximum(logits, 0))
relu_log = np.log(1 + np.maximum(output.model_output, 0))
weighted_log = relu_log * np.expand_dims(attention_mask, axis=-1)
weighted_log = relu_log * np.expand_dims(output.attention_mask, axis=-1)
scores = np.max(weighted_log, axis=1)

7
fastembed/text/clip_embedding.py
View File

@@ -2,6 +2,7 @@ from typing import Type, List, Dict, Any, Tuple, Iterable
import numpy as np
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
from fastembed.text.onnx_text_model import TextEmbeddingWorker
@@ -33,12 +34,10 @@ class CLIPOnnxEmbedding(OnnxTextEmbedding):
"""
return supported_clip_models
@classmethod
def _post_process_onnx_output(
cls, output: Tuple[np.ndarray, np.ndarray]
self, output: OnnxOutputContext
) -> Iterable[np.ndarray]:
embeddings, attention_mask = output
return embeddings
return output.model_output
class CLIPEmbeddingWorker(OnnxTextEmbeddingWorker):

9
fastembed/text/jina_onnx_embedding.py
View File

@@ -2,6 +2,7 @@ from typing import Type, List, Dict, Any, Tuple, Iterable
import numpy as np
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
@@ -50,12 +51,12 @@ class JinaOnnxEmbedding(OnnxTextEmbedding):
"""
return supported_jina_models
@classmethod
def _post_process_onnx_output(
cls, output: Tuple[np.ndarray, np.ndarray]
self, output: OnnxOutputContext
) -> Iterable[np.ndarray]:
embeddings, attn_mask = output
return normalize(cls.mean_pooling(embeddings, attn_mask)).astype(np.float32)
embeddings = output.model_output
attn_mask = output.attention_mask
return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
class JinaEmbeddingWorker(OnnxTextEmbeddingWorker):

6
fastembed/text/onnx_embedding.py
View File

@@ -3,6 +3,7 @@ from typing import Dict, Optional, Tuple, Union, Iterable, Type, List, Any, Sequ
import numpy as np
from fastembed.common import OnnxProvider
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common.utils import normalize, define_cache_dir
from fastembed.text.onnx_text_model import TextEmbeddingWorker, OnnxTextModel
from fastembed.text.text_embedding_base import TextEmbeddingBase
@@ -282,11 +283,10 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
"""
return onnx_input
@classmethod
def _post_process_onnx_output(
cls, output: Tuple[np.ndarray, np.ndarray]
self, output: OnnxOutputContext
) -> Iterable[np.ndarray]:
embeddings, _ = output
embeddings = output.model_output
return normalize(embeddings[:, 0]).astype(np.float32)

28
fastembed/text/onnx_text_model.py
View File

@@ -7,23 +7,25 @@ import numpy as np
from fastembed.common import OnnxProvider
from fastembed.common.preprocessor_utils import load_tokenizer
from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker, T
from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker, T, OnnxOutputContext
from fastembed.common.utils import iter_batch
from fastembed.parallel_processor import ParallelWorkerPool
class OnnxTextModel(OnnxModel[T]):
ONNX_OUTPUT_NAMES: Optional[List[str]] = None
@classmethod
def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
raise NotImplementedError("Subclasses must implement this method")
@classmethod
def _post_process_onnx_output(cls, output: Tuple[np.ndarray, np.ndarray]) -> Iterable[T]:
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
raise NotImplementedError("Subclasses must implement this method")
def __init__(self) -> None:
super().__init__()
self.tokenizer = None
self.special_token_to_id = {}
def _preprocess_onnx_input(self, onnx_input: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
"""
@@ -41,9 +43,12 @@ class OnnxTextModel(OnnxModel[T]):
super().load_onnx_model(
model_dir=model_dir, model_file=model_file, threads=threads, providers=providers
)
self.tokenizer = load_tokenizer(model_dir=model_dir)
self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
def onnx_embed(self, documents: List[str]) -> Tuple[np.ndarray, np.ndarray]:
def onnx_embed(
self,
documents: List[str],
) -> OnnxOutputContext:
encoded = self.tokenizer.encode_batch(documents)
input_ids = np.array([e.ids for e in encoded])
attention_mask = np.array([e.attention_mask for e in encoded])
@@ -60,9 +65,12 @@ class OnnxTextModel(OnnxModel[T]):
onnx_input = self._preprocess_onnx_input(onnx_input)
model_output = self.model.run(None, onnx_input)
embeddings = model_output[0]
return embeddings, attention_mask
model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)
return OnnxOutputContext(
model_output=model_output[0],
attention_mask=attention_mask,
input_ids=input_ids,
)
def _embed_documents(
self,
@@ -104,5 +112,5 @@ class OnnxTextModel(OnnxModel[T]):
class TextEmbeddingWorker(EmbeddingWorker):
def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
for idx, batch in items:
embeddings, attn_mask = self.model.onnx_embed(batch)
yield idx, (embeddings, attn_mask)
onnx_output = self.model.onnx_embed(batch)
yield idx, onnx_output

3
pyproject.toml
View File

@@ -24,6 +24,9 @@ numpy = [
{ version = ">=1.26", python = ">=3.12" }
]
pillow = "^10.3.0"
snowballstemmer = "^2.2.0"
PyStemmer = "^2.2.0"
mmh3 = "^4.0"
[tool.poetry.group.dev.dependencies]
pytest = "^7.4.2"

58
tests/test_attention_embeddings.py Normal file
View File

@@ -0,0 +1,58 @@
import numpy as np
from fastembed import SparseTextEmbedding
def test_attention_embeddings():
model = SparseTextEmbedding(model_name="Qdrant/bm42-all-minilm-l6-v2-attentions")
output = list(
model.query_embed(
[
"I must not fear. Fear is the mind-killer.",
]
)
)
assert len(output) == 1
for result in output:
assert len(result.indices) == len(result.values)
assert np.allclose(result.values, np.ones(len(result.values)))
quotes = [
"I must not fear. Fear is the mind-killer.",
"All animals are equal, but some animals are more equal than others.",
"It was a pleasure to burn.",
"The sky above the port was the color of television, tuned to a dead channel.",
"In the beginning, the universe was created."
" This has made a lot of people very angry and been widely regarded as a bad move.",
"It's a truth universally acknowledged that a zombie in possession of brains must be in want of more brains.",
"War is peace. Freedom is slavery. Ignorance is strength.",
"We're not in Infinity; we're in the suburbs.",
"I was a thousand times more evil than thou!",
"History is merely a list of surprises... It can only prepare us to be surprised yet again.",
]
output = list(model.embed(quotes))
assert len(output) == len(quotes)
for result in output:
assert len(result.indices) == len(result.values)
assert len(result.indices) > 0
# Test support for unknown languages
output = list(
model.query_embed(
[
"привет мир!",
]
)
)
assert len(output) == 1
for result in output:
assert len(result.indices) == len(result.values)
assert len(result.indices) == 2

18
tests/test_sparse_embeddings.py
View File

@@ -1,4 +1,5 @@
import pytest
from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
CANONICAL_COLUMN_VALUES = {
@@ -47,11 +48,8 @@ def test_batch_embedding():
docs_to_embed = docs * 10
for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
print("evaluating", model_name)
model = SparseTextEmbedding(model_name=model_name)
result = next(iter(model.embed(docs_to_embed, batch_size=6)))
print(result.indices)
assert result.indices.tolist() == expected_result["indices"]
for i, value in enumerate(result.values):
@@ -59,18 +57,16 @@ def test_batch_embedding():
def test_single_embedding():
docs_to_embed = docs
for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
print("evaluating", model_name)
model = SparseTextEmbedding(model_name=model_name)
result = next(iter(model.embed(docs_to_embed, batch_size=6)))
print(result.indices)
assert result.indices.tolist() == expected_result["indices"]
passage_result = next(iter(model.embed(docs, batch_size=6)))
query_result = next(iter(model.query_embed(docs)))
for result in [passage_result, query_result]:
assert result.indices.tolist() == expected_result["indices"]
for i, value in enumerate(result.values):
assert pytest.approx(value, abs=0.001) == expected_result["values"][i]
for i, value in enumerate(result.values):
assert pytest.approx(value, abs=0.001) == expected_result["values"][i]
def test_parallel_processing():