diff --git a/experiments/if_splade_to_onnx.py b/experiments/if_splade_to_onnx.py new file mode 100644 index 000000000..e7b652337 --- /dev/null +++ b/experiments/if_splade_to_onnx.py @@ -0,0 +1,134 @@ +"""Export an inference-free SPLADE document encoder to ONNX. + +Converts `opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte` (an MLM head +over a GTE backbone) into an onnx model producing token logits, and assembles a model dir +with everything fastembed's `IfSplade` needs: model.onnx, tokenizer files and idf.json. + +Usage: + python experiments/if_splade_to_onnx.py --output-dir models/opensearch-neural-sparse-encoding-doc-v3-gte +""" + +import argparse +import shutil +from pathlib import Path + +import torch +from huggingface_hub import hf_hub_download +from transformers import AutoModelForMaskedLM, AutoTokenizer + +MODEL_ID = "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte" +# revision of the remote modeling code (Alibaba-NLP/new-impl), pinned in the model card +CODE_REVISION = "40ced75c3017eb27626c9d4ea981bde21a2662f4" + +TOKENIZER_FILES = [ + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "vocab.txt", + "idf.json", +] + + +class LogitsOnly(torch.nn.Module): + def __init__(self, model: torch.nn.Module): + super().__init__() + self.model = model + + def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: + return self.model(input_ids=input_ids, attention_mask=attention_mask).logits + + +def export(model_id: str, output_dir: Path, opset: int = 14) -> Path: + output_dir.mkdir(parents=True, exist_ok=True) + + model = AutoModelForMaskedLM.from_pretrained( + model_id, trust_remote_code=True, code_revision=CODE_REVISION + ) + model.eval() + wrapped = LogitsOnly(model) + + tokenizer = AutoTokenizer.from_pretrained(model_id) + dummy = tokenizer( + ["fastembed is a library", "onnx export"], + padding=True, + truncation=True, + return_tensors="pt", + return_token_type_ids=False, + ) + + onnx_path = output_dir / "model.onnx" + with torch.inference_mode(): + torch.onnx.export( + wrapped, + (dummy["input_ids"], dummy["attention_mask"]), + f=onnx_path.as_posix(), + input_names=["input_ids", "attention_mask"], + output_names=["logits"], + dynamic_axes={ + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + "logits": {0: "batch_size", 1: "sequence_length"}, + }, + do_constant_folding=True, + opset_version=opset, + dynamo=False, + ) + + for file_name in TOKENIZER_FILES: + local_path = hf_hub_download(repo_id=model_id, filename=file_name) + shutil.copy(local_path, output_dir / file_name) + + return onnx_path + + +def parity_check(model_id: str, output_dir: Path) -> None: + import numpy as np + import onnxruntime as ort + + model = AutoModelForMaskedLM.from_pretrained( + model_id, trust_remote_code=True, code_revision=CODE_REVISION + ) + model.eval() + tokenizer = AutoTokenizer.from_pretrained(model_id) + + documents = [ + "Currently New York is rainy.", + "fastembed is a lightweight library for generating embeddings", + "hello world", + ] + features = tokenizer( + documents, padding=True, truncation=True, return_tensors="pt", return_token_type_ids=False + ) + + with torch.inference_mode(): + torch_logits = model(**features).logits.numpy() + + session = ort.InferenceSession(output_dir / "model.onnx") + onnx_logits = session.run( + ["logits"], + { + "input_ids": features["input_ids"].numpy(), + "attention_mask": features["attention_mask"].numpy(), + }, + )[0] + + max_diff = np.abs(torch_logits - onnx_logits).max() + print(f"max |torch - onnx| logits diff: {max_diff}") + assert max_diff < 1e-3, "onnx export does not match the torch model" + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-id", default=MODEL_ID) + parser.add_argument("--output-dir", default=f"models/{MODEL_ID.replace('/', '_')}", type=Path) + parser.add_argument("--opset", default=14, type=int) + args = parser.parse_args() + + onnx_path = export(args.model_id, args.output_dir, args.opset) + print(f"Exported to {onnx_path}") + parity_check(args.model_id, args.output_dir) + + +if __name__ == "__main__": + main() diff --git a/fastembed/sparse/if_splade.py b/fastembed/sparse/if_splade.py new file mode 100644 index 000000000..51767e85d --- /dev/null +++ b/fastembed/sparse/if_splade.py @@ -0,0 +1,246 @@ +import json +from typing import Any, Iterable, Sequence, Type + +import numpy as np + +from fastembed.common import OnnxProvider +from fastembed.common.model_description import ModelSource, SparseModelDescription +from fastembed.common.onnx_model import OnnxOutputContext +from fastembed.common.preprocessor_utils import load_tokenizer +from fastembed.common.types import Device +from fastembed.common.utils import define_cache_dir, iter_batch +from fastembed.sparse.sparse_embedding_base import ( + SparseEmbedding, + SparseTextEmbeddingBase, +) +from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker + +IDF_FILE = "idf.json" + +supported_if_splade_models: list[SparseModelDescription] = [ + SparseModelDescription( + model="opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte", + vocab_size=30522, + description="Inference-free SPLADE model. Documents are expanded with an ONNX encoder at index " + "time, queries are encoded with a tokenizer and an IDF lookup table only, " + "without any model inference.", + license="apache-2.0", + size_in_GB=0.55, + sources=ModelSource(hf="Qdrant/opensearch-neural-sparse-encoding-doc-v3-gte"), + model_file="model.onnx", + additional_files=[IDF_FILE], + requires_idf=None, + ), +] + + +class IfSplade(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]): + """Inference-free (asymmetric) SPLADE model. + + Documents are encoded with a neural encoder which expands them into a sparse vocabulary-sized + vector, while queries are encoded by tokenizing the text and looking up a precomputed IDF + weight per token — no neural inference happens at query time. + + Query and document embeddings are compared with a dot product. + Special tokens are excluded from both document and query embeddings. + """ + + def _post_process_onnx_output( + self, output: OnnxOutputContext, **kwargs: Any + ) -> Iterable[SparseEmbedding]: + if output.attention_mask is None: + raise ValueError("attention_mask must be provided for document post-processing") + + # Max-pool token logits over the sequence, masking out the padding + pooled = np.max( + output.model_output * np.expand_dims(output.attention_mask, axis=-1), axis=1 + ) + # v3 models of the opensearch-neural-sparse family use a double log activation, + # log(1 + log(1 + relu(x))), to increase sparsity of document embeddings + scores = np.log1p(np.log1p(np.maximum(pooled, 0.0))) + + if self.special_tokens_ids: + scores[:, list(self.special_tokens_ids)] = 0.0 + + for row_scores in scores: + indices = row_scores.nonzero()[0] + yield SparseEmbedding(values=row_scores[indices], indices=indices) + + def token_count( + self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any + ) -> int: + # unlike `OnnxTextModel._token_count`, does not require the onnx model to be loaded + token_num = 0 + texts = [texts] if isinstance(texts, str) else texts + for batch in iter_batch(texts, batch_size): + for tokens in self.tokenizer.encode_batch(batch): # type: ignore[union-attr] + token_num += sum(tokens.attention_mask) + return token_num + + @classmethod + def _list_supported_models(cls) -> list[SparseModelDescription]: + """Lists the supported models. + + Returns: + list[SparseModelDescription]: A list of SparseModelDescription objects containing the model information. + """ + return supported_if_splade_models + + def __init__( + self, + model_name: str, + cache_dir: str | None = None, + threads: int | None = None, + providers: Sequence[OnnxProvider] | None = None, + cuda: bool | Device = Device.AUTO, + device_ids: list[int] | None = None, + lazy_load: bool = False, + device_id: int | None = None, + specific_model_path: str | None = None, + **kwargs: Any, + ): + """ + 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 list of onnxruntime providers to use. + Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None. + cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers` + Defaults to Device. + device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in + workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, 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. + device_id (Optional[int], optional): The device id to use for loading the model in the worker process. + specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else + + Raises: + ValueError: If the model_name is not in the format / e.g. BAAI/bge-base-en. + """ + super().__init__(model_name, cache_dir, threads, **kwargs) + self.providers = providers + self.lazy_load = lazy_load + self._extra_session_options = self._select_exposed_session_options(kwargs) + + # List of device ids, that can be used for data parallel processing in workers + self.device_ids = device_ids + self.cuda = cuda + + # This device_id will be used if we need to load model in current process + self.device_id: int | None = None + if device_id is not None: + self.device_id = device_id + elif self.device_ids is not None: + self.device_id = self.device_ids[0] + + self.model_description = self._get_model_description(model_name) + self.cache_dir = str(define_cache_dir(cache_dir)) + + self._specific_model_path = specific_model_path + self._model_dir = self.download_model( + self.model_description, + self.cache_dir, + local_files_only=self._local_files_only, + specific_model_path=self._specific_model_path, + ) + + # The tokenizer and the idf table are lightweight and are required for query embedding, + # which does not involve any model inference, so they are loaded eagerly, while + # `lazy_load` only defers the initialization of the onnx model + self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=self._model_dir) + self.special_tokens_ids: set[int] = set(self.special_token_to_id.values()) + self._token_id_to_idf = self._load_idf() + + if not self.lazy_load: + self.load_onnx_model() + + def load_onnx_model(self) -> None: + self._load_onnx_model( + model_dir=self._model_dir, + model_file=self.model_description.model_file, + threads=self.threads, + providers=self.providers, + cuda=self.cuda, + device_id=self.device_id, + extra_session_options=self._extra_session_options, + ) + + def _load_idf(self) -> dict[int, float]: + with open(self._model_dir / IDF_FILE) as f: + token_to_idf: dict[str, float] = json.load(f) + + vocab: dict[str, int] = self.tokenizer.get_vocab() # type: ignore[union-attr] + return {vocab[token]: idf for token, idf in token_to_idf.items() if token in vocab} + + def embed( + self, + documents: str | Iterable[str], + batch_size: int = 256, + parallel: int | None = None, + **kwargs: Any, + ) -> Iterable[SparseEmbedding]: + """ + Encode a list of documents into list of embeddings. + + 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, + providers=self.providers, + cuda=self.cuda, + device_ids=self.device_ids, + local_files_only=self._local_files_only, + specific_model_path=self._specific_model_path, + extra_session_options=self._extra_session_options, + **kwargs, + ) + + def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]: + """ + Encode a list of queries into list of sparse embeddings without any model inference. + + A query is tokenized, and each unique token is assigned its IDF weight from + a precomputed lookup table shipped with the model. Special tokens are ignored. + """ + if isinstance(query, str): + query = [query] + + for text in query: + token_ids = set(self.tokenizer.encode(text).ids) - self.special_tokens_ids # type: ignore[union-attr] + embedding = { + token_id: self._token_id_to_idf[token_id] + for token_id in sorted(token_ids) + if token_id in self._token_id_to_idf + } + yield SparseEmbedding.from_dict(embedding) + + @classmethod + def _get_worker_class(cls) -> Type[TextEmbeddingWorker[SparseEmbedding]]: + return IfSpladeEmbeddingWorker + + +class IfSpladeEmbeddingWorker(TextEmbeddingWorker[SparseEmbedding]): + def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> IfSplade: + return IfSplade( + model_name=model_name, + cache_dir=cache_dir, + threads=1, + **kwargs, + ) diff --git a/fastembed/sparse/sparse_text_embedding.py b/fastembed/sparse/sparse_text_embedding.py index 5b5c83085..273e08f0a 100644 --- a/fastembed/sparse/sparse_text_embedding.py +++ b/fastembed/sparse/sparse_text_embedding.py @@ -5,6 +5,7 @@ from fastembed.common.types import Device from fastembed.sparse.bm25 import Bm25 from fastembed.sparse.bm42 import Bm42 +from fastembed.sparse.if_splade import IfSplade from fastembed.sparse.minicoil import MiniCOIL from fastembed.sparse.sparse_embedding_base import ( SparseEmbedding, @@ -16,7 +17,13 @@ class SparseTextEmbedding(SparseTextEmbeddingBase): - EMBEDDINGS_REGISTRY: list[Type[SparseTextEmbeddingBase]] = [SpladePP, Bm42, Bm25, MiniCOIL] + EMBEDDINGS_REGISTRY: list[Type[SparseTextEmbeddingBase]] = [ + SpladePP, + Bm42, + Bm25, + MiniCOIL, + IfSplade, + ] @classmethod def list_supported_models(cls) -> list[dict[str, Any]]: diff --git a/tests/test_sparse_embeddings.py b/tests/test_sparse_embeddings.py index c2a7e2ffa..e7c1c936c 100644 --- a/tests/test_sparse_embeddings.py +++ b/tests/test_sparse_embeddings.py @@ -8,6 +8,7 @@ from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding from tests.utils import delete_model_cache, should_test_model + CANONICAL_COLUMN_VALUES = { "prithivida/Splade_PP_en_v1": { "indices": [ @@ -58,9 +59,50 @@ -0.12508166, ], }, + # first 15 non-zero dimensions of the embedding + "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte": { + "indices": [ + 999, + 1010, + 1011, + 1024, + 1028, + 1029, + 1045, + 1074, + 1993, + 2017, + 2033, + 2054, + 2073, + 2080, + 2088, + ], + "values": [ + 0.16544909, + 0.00529129, + 0.0392109, + 0.12337475, + 0.09640586, + 0.05325737, + 0.09611791, + 0.03159865, + 0.01349991, + 0.09392473, + 0.01928805, + 0.05238346, + 0.05515401, + 0.03156782, + 0.98263124, + ], + }, } CANONICAL_QUERY_VALUES = { + "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte": { + "indices": [2088, 7592], + "values": [3.42086864, 6.93775654], + }, "Qdrant/minicoil-v1": { "indices": [80, 81, 82, 83, 6664, 6665, 6666, 6667], "values": [ @@ -82,6 +124,7 @@ "Qdrant/minicoil-v1", "Qdrant/bm25", "Qdrant/bm42-all-minilm-l6-v2-attentions", + "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte", ) MODELS_TO_CACHE = tuple([x.lower() for x in _MODELS_TO_CACHE]) @@ -142,18 +185,23 @@ def test_single_embedding(model_cache) -> None: continue if not should_test_model(model_desc, model_desc.model, is_ci, is_manual): continue - with model_cache(model_desc.model) as model: passage_result = next(iter(model.embed(docs, batch_size=6))) query_result = next(iter(model.query_embed(docs))) expected_result = CANONICAL_COLUMN_VALUES[model_desc.model] expected_query_result = CANONICAL_QUERY_VALUES.get(model_desc.model, expected_result) - assert passage_result.indices.tolist() == expected_result["indices"] - for i, value in enumerate(passage_result.values): + + # canonical values might contain only a prefix of the non-zero dimensions + num_dims = len(expected_result["indices"]) + assert passage_result.indices.tolist()[:num_dims] == expected_result["indices"] + for i, value in enumerate(passage_result.values[:num_dims]): assert pytest.approx(value, abs=0.001) == expected_result["values"][i] - assert query_result.indices.tolist() == expected_query_result["indices"] - for i, value in enumerate(query_result.values): + num_query_dims = len(expected_query_result["indices"]) + assert ( + query_result.indices.tolist()[:num_query_dims] == expected_query_result["indices"] + ) + for i, value in enumerate(query_result.values[:num_query_dims]): assert pytest.approx(value, abs=0.001) == expected_query_result["values"][i] @@ -263,6 +311,22 @@ def test_disable_stemmer_behavior(disable_stemmer: bool) -> None: assert result == expected, f"Expected {expected}, but got {result}" +def test_if_splade_query_embed_is_inference_free() -> None: + is_ci = os.getenv("CI") + model = SparseTextEmbedding( + model_name="opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte", + lazy_load=True, + ) + embeddings = list(model.query_embed(["hello world", "flag embedding"])) + # queries are embedded with a tokenizer and an idf lookup table only, + # the onnx model must stay unloaded + assert not hasattr(model.model, "model") + assert all(len(embedding.indices) > 0 for embedding in embeddings) + + if is_ci: + delete_model_cache(model.model._model_dir) + + @pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1"]) def test_lazy_load(model_name: str) -> None: is_ci = os.getenv("CI")