Source code for langchain_qdrant.fastembed_sparse
from typing import Any, List, Optional, Sequence
from langchain_qdrant.sparse_embeddings import SparseEmbeddings, SparseVector
[docs]class FastEmbedSparse(SparseEmbeddings):
"""An interface for sparse embedding models to use with Qdrant."""
[docs] def __init__(
self,
model_name: str = "Qdrant/bm25",
batch_size: int = 256,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
providers: Optional[Sequence[Any]] = None,
parallel: Optional[int] = None,
**kwargs: Any,
) -> None:
"""
Sparse encoder implementation using FastEmbed - https://qdrant.github.io/fastembed/
For a list of available models, see https://qdrant.github.io/fastembed/examples/Supported_Models/
Args:
model_name (str): The name of the model to use. Defaults to `"Qdrant/bm25"`.
batch_size (int): Batch size for encoding. Defaults to 256.
cache_dir (str, optional): The path to the model cache directory.\
Can also be set using the\
`FASTEMBED_CACHE_PATH` env variable.
threads (int, optional): The number of threads onnxruntime session can use.
providers (Sequence[Any], optional): List of ONNX execution providers.\
parallel (int, optional): If `>1`, data-parallel encoding will be used, r\
Recommended for encoding of large datasets.\
If `0`, use all available cores.\
If `None`, don't use data-parallel processing,\
use default onnxruntime threading instead.\
Defaults to None.
kwargs: Additional options to pass to fastembed.SparseTextEmbedding
Raises:
ValueError: If the model_name is not supported in SparseTextEmbedding.
"""
try:
from fastembed import SparseTextEmbedding # type: ignore
except ImportError:
raise ValueError(
"The 'fastembed' package is not installed. "
"Please install it with "
"`pip install fastembed` or `pip install fastembed-gpu`."
)
self._batch_size = batch_size
self._parallel = parallel
self._model = SparseTextEmbedding(
model_name=model_name,
cache_dir=cache_dir,
threads=threads,
providers=providers,
**kwargs,
)
[docs] def embed_documents(self, texts: List[str]) -> List[SparseVector]:
results = self._model.embed(
texts, batch_size=self._batch_size, parallel=self._parallel
)
return [
SparseVector(indices=result.indices.tolist(), values=result.values.tolist())
for result in results
]
[docs] def embed_query(self, text: str) -> SparseVector:
result = next(self._model.query_embed(text))
return SparseVector(
indices=result.indices.tolist(), values=result.values.tolist()
)