chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,58 @@
|
||||
# Copyright 2025-present the zvec project
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from .bm25_embedding_function import BM25EmbeddingFunction
|
||||
from .embedding_function import DenseEmbeddingFunction, SparseEmbeddingFunction
|
||||
from .http_embedding_function import HTTPDenseEmbedding
|
||||
from .jina_embedding_function import JinaDenseEmbedding
|
||||
from .jina_function import JinaFunctionBase
|
||||
from .multi_vector_reranker import CallbackReRanker, RrfReRanker, WeightedReRanker
|
||||
from .openai_embedding_function import OpenAIDenseEmbedding
|
||||
from .openai_function import OpenAIFunctionBase
|
||||
from .qwen_embedding_function import QwenDenseEmbedding, QwenSparseEmbedding
|
||||
from .qwen_function import QwenFunctionBase
|
||||
from .qwen_rerank_function import QwenReRanker
|
||||
from .rerank_function import RerankFunction
|
||||
from .rerank_function import RerankFunction as ReRanker
|
||||
from .sentence_transformer_embedding_function import (
|
||||
DefaultLocalDenseEmbedding,
|
||||
DefaultLocalSparseEmbedding,
|
||||
)
|
||||
from .sentence_transformer_function import SentenceTransformerFunctionBase
|
||||
from .sentence_transformer_rerank_function import DefaultLocalReRanker
|
||||
|
||||
__all__ = [
|
||||
"BM25EmbeddingFunction",
|
||||
"CallbackReRanker",
|
||||
"DefaultLocalDenseEmbedding",
|
||||
"DefaultLocalReRanker",
|
||||
"DefaultLocalSparseEmbedding",
|
||||
"DenseEmbeddingFunction",
|
||||
"HTTPDenseEmbedding",
|
||||
"JinaDenseEmbedding",
|
||||
"JinaFunctionBase",
|
||||
"OpenAIDenseEmbedding",
|
||||
"OpenAIFunctionBase",
|
||||
"QwenDenseEmbedding",
|
||||
"QwenFunctionBase",
|
||||
"QwenReRanker",
|
||||
"QwenSparseEmbedding",
|
||||
"ReRanker",
|
||||
"RerankFunction",
|
||||
"RrfReRanker",
|
||||
"SentenceTransformerFunctionBase",
|
||||
"SparseEmbeddingFunction",
|
||||
"WeightedReRanker",
|
||||
]
|
||||
@@ -0,0 +1,375 @@
|
||||
# Copyright 2025-present the zvec project
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from functools import lru_cache
|
||||
from typing import Literal, Optional
|
||||
|
||||
from ..common.constants import TEXT, SparseVectorType
|
||||
from ..tool import require_module
|
||||
from .embedding_function import SparseEmbeddingFunction
|
||||
|
||||
|
||||
class BM25EmbeddingFunction(SparseEmbeddingFunction[TEXT]):
|
||||
"""BM25-based sparse embedding function using DashText SDK.
|
||||
|
||||
This class provides text-to-sparse-vector embedding capabilities using
|
||||
the DashText library with BM25 algorithm. BM25 (Best Matching 25) is a
|
||||
probabilistic retrieval function used for lexical search and document
|
||||
ranking based on term frequency and inverse document frequency.
|
||||
|
||||
BM25 generates sparse vectors where each dimension corresponds to a term in
|
||||
the vocabulary, and the value represents the BM25 score for that term. It's
|
||||
particularly effective for:
|
||||
|
||||
- Lexical search and keyword matching
|
||||
- Document ranking and information retrieval
|
||||
- Combining with dense embeddings for hybrid search
|
||||
- Traditional IR tasks where exact term matching is important
|
||||
|
||||
This implementation uses DashText's SparseVectorEncoder, which provides
|
||||
efficient BM25 computation for Chinese and English text using either a
|
||||
built-in encoder or custom corpus training.
|
||||
|
||||
Args:
|
||||
corpus (Optional[list[str]], optional): List of documents to train the
|
||||
BM25 encoder. If provided, creates a custom encoder trained on this
|
||||
corpus for better domain-specific accuracy. If ``None``, uses the
|
||||
built-in encoder. Defaults to ``None``.
|
||||
encoding_type (Literal["query", "document"], optional): Encoding mode
|
||||
for text processing. Use ``"query"`` for search queries (default) and
|
||||
``"document"`` for document indexing. This distinction optimizes the
|
||||
BM25 scoring for asymmetric retrieval tasks. Defaults to ``"query"``.
|
||||
language (Literal["zh", "en"], optional): Language for built-in encoder.
|
||||
Only used when corpus is None. ``"zh"`` for Chinese (trained on Chinese
|
||||
Wikipedia), ``"en"`` for English. Defaults to ``"zh"``.
|
||||
b (float, optional): Document length normalization parameter for BM25.
|
||||
Range [0, 1]. 0 means no normalization, 1 means full normalization.
|
||||
Only used with custom corpus. Defaults to ``0.75``.
|
||||
k1 (float, optional): Term frequency saturation parameter for BM25.
|
||||
Higher values give more weight to term frequency. Only used with
|
||||
custom corpus. Defaults to ``1.2``.
|
||||
**kwargs: Additional parameters for DashText encoder customization.
|
||||
|
||||
Attributes:
|
||||
corpus_size (int): Number of documents in the training corpus (0 if using built-in encoder).
|
||||
encoding_type (str): The encoding type being used ("query" or "document").
|
||||
language (str): The language of the built-in encoder ("zh" or "en").
|
||||
|
||||
Raises:
|
||||
ValueError: If corpus is provided but empty or contains non-string elements.
|
||||
TypeError: If input to ``embed()`` is not a string.
|
||||
RuntimeError: If DashText encoder initialization or training fails.
|
||||
|
||||
Note:
|
||||
- Requires Python 3.10, 3.11, or 3.12
|
||||
- Requires the ``dashtext`` package: ``pip install dashtext``
|
||||
- Two encoder options available:
|
||||
|
||||
1. **Built-in encoder** (no corpus needed): Pre-trained models for
|
||||
Chinese (zh) and English (en), good generalization, works out-of-the-box
|
||||
2. **Custom encoder** (corpus required): Better accuracy for domain-specific
|
||||
terminology, requires training on your full corpus with BM25 parameters
|
||||
|
||||
- Encoding types:
|
||||
|
||||
* ``encoding_type="query"``: Optimized for search queries (shorter text)
|
||||
* ``encoding_type="document"``: Optimized for document indexing (longer text)
|
||||
|
||||
- BM25 parameters (b, k1) only apply to custom encoder training
|
||||
- Output is sorted by indices (vocabulary term IDs) for consistency
|
||||
- Results are cached (LRU cache, maxsize=10) to reduce computation
|
||||
- No API key or network connectivity required (local computation)
|
||||
|
||||
Examples:
|
||||
>>> # Option 1: Using built-in encoder for Chinese (no corpus needed)
|
||||
>>> from zvec.extension import BM25EmbeddingFunction
|
||||
>>>
|
||||
>>> # For query encoding (Chinese)
|
||||
>>> bm25_query_zh = BM25EmbeddingFunction(language="zh", encoding_type="query")
|
||||
>>> query_vec = bm25_query_zh.embed("什么是机器学习")
|
||||
>>> isinstance(query_vec, dict)
|
||||
True
|
||||
>>> # query_vec: {1169440797: 0.29, 2045788977: 0.70, ...}
|
||||
|
||||
>>> # For document encoding (Chinese)
|
||||
>>> bm25_doc_zh = BM25EmbeddingFunction(language="zh", encoding_type="document")
|
||||
>>> doc_vec = bm25_doc_zh.embed("机器学习是人工智能的一个重要分支...")
|
||||
>>> isinstance(doc_vec, dict)
|
||||
True
|
||||
|
||||
>>> # Using built-in encoder for English
|
||||
>>> bm25_query_en = BM25EmbeddingFunction(language="en", encoding_type="query")
|
||||
>>> query_vec_en = bm25_query_en.embed("what is vector search service")
|
||||
>>> isinstance(query_vec_en, dict)
|
||||
True
|
||||
|
||||
>>> # Option 2: Using custom corpus for domain-specific accuracy
|
||||
>>> corpus = [
|
||||
... "机器学习是人工智能的一个重要分支",
|
||||
... "深度学习使用多层神经网络进行特征提取",
|
||||
... "自然语言处理技术用于理解和生成人类语言"
|
||||
... ]
|
||||
>>> bm25_custom = BM25EmbeddingFunction(
|
||||
... corpus=corpus,
|
||||
... encoding_type="query",
|
||||
... b=0.75,
|
||||
... k1=1.2
|
||||
... )
|
||||
>>> custom_vec = bm25_custom.embed("机器学习算法")
|
||||
>>> isinstance(custom_vec, dict)
|
||||
True
|
||||
|
||||
>>> # Hybrid search: combining with dense embeddings
|
||||
>>> from zvec.extension import DefaultLocalDenseEmbedding
|
||||
>>> dense_emb = DefaultLocalDenseEmbedding()
|
||||
>>> bm25_emb = BM25EmbeddingFunction(language="zh", encoding_type="query")
|
||||
>>>
|
||||
>>> query = "machine learning algorithms"
|
||||
>>> dense_vec = dense_emb.embed(query) # Semantic similarity
|
||||
>>> sparse_vec = bm25_emb.embed(query) # Lexical matching
|
||||
>>> # Combine scores for hybrid retrieval
|
||||
|
||||
>>> # Callable interface
|
||||
>>> sparse_vec = bm25_query_zh("information retrieval")
|
||||
>>> isinstance(sparse_vec, dict)
|
||||
True
|
||||
|
||||
>>> # Error handling
|
||||
>>> try:
|
||||
... bm25_query_zh.embed("") # Empty query
|
||||
... except ValueError as e:
|
||||
... print(f"Error: {e}")
|
||||
Error: Input text cannot be empty or whitespace only
|
||||
|
||||
See Also:
|
||||
- ``SparseEmbeddingFunction``: Base class for sparse embeddings
|
||||
- ``DefaultLocalSparseEmbedding``: SPLADE-based sparse embedding
|
||||
- ``QwenSparseEmbedding``: API-based sparse embedding using Qwen
|
||||
- ``DefaultLocalDenseEmbedding``: Dense embedding for semantic search
|
||||
|
||||
References:
|
||||
- DashText Documentation: https://help.aliyun.com/zh/document_detail/2546039.html
|
||||
- DashText PyPI: https://pypi.org/project/dashtext/
|
||||
- BM25 Algorithm: Robertson & Zaragoza (2009)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
corpus: Optional[list[str]] = None,
|
||||
encoding_type: Literal["query", "document"] = "query",
|
||||
language: Literal["zh", "en"] = "zh",
|
||||
b: float = 0.75,
|
||||
k1: float = 1.2,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize the BM25 embedding function.
|
||||
|
||||
Args:
|
||||
corpus (Optional[list[str]]): Optional corpus for training custom encoder.
|
||||
If None, uses built-in encoder. Defaults to None.
|
||||
encoding_type (Literal["query", "document"]): Text encoding mode.
|
||||
Use "query" for search queries, "document" for indexing.
|
||||
Defaults to "query".
|
||||
language (Literal["zh", "en"]): Language for built-in encoder.
|
||||
"zh" for Chinese, "en" for English. Defaults to "zh".
|
||||
b (float): Document length normalization for BM25 [0, 1].
|
||||
Only used with custom corpus. Defaults to 0.75.
|
||||
k1 (float): Term frequency saturation for BM25.
|
||||
Only used with custom corpus. Defaults to 1.2.
|
||||
**kwargs: Additional DashText encoder parameters.
|
||||
|
||||
Raises:
|
||||
ValueError: If corpus is provided but empty or invalid.
|
||||
ImportError: If dashtext package is not installed.
|
||||
RuntimeError: If encoder initialization or training fails.
|
||||
"""
|
||||
# Validate corpus if provided
|
||||
if corpus is not None:
|
||||
if not corpus or not isinstance(corpus, list):
|
||||
raise ValueError("Corpus must be a non-empty list of strings")
|
||||
|
||||
if not all(isinstance(doc, str) for doc in corpus):
|
||||
raise ValueError("All corpus documents must be strings")
|
||||
|
||||
# Import dashtext
|
||||
self._dashtext = require_module("dashtext")
|
||||
|
||||
self._corpus = corpus
|
||||
self._encoding_type = encoding_type
|
||||
self._language = language
|
||||
self._b = b
|
||||
self._k1 = k1
|
||||
self._extra_params = kwargs
|
||||
|
||||
# Initialize the BM25 encoder
|
||||
self._build_encoder()
|
||||
|
||||
def _build_encoder(self):
|
||||
"""Build the BM25 sparse vector encoder.
|
||||
|
||||
Creates either a built-in encoder (pre-trained) or a custom encoder
|
||||
trained on the provided corpus.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If encoder initialization or training fails.
|
||||
ImportError: If dashtext package is not installed.
|
||||
"""
|
||||
try:
|
||||
if self._corpus is None:
|
||||
# Use built-in encoder (pre-trained on Wikipedia)
|
||||
# language: 'zh' for Chinese, 'en' for English
|
||||
self._encoder = self._dashtext.SparseVectorEncoder.default(
|
||||
name=self._language
|
||||
)
|
||||
else:
|
||||
# Create custom encoder with BM25 parameters
|
||||
self._encoder = self._dashtext.SparseVectorEncoder(
|
||||
b=self._b, k1=self._k1, **self._extra_params
|
||||
)
|
||||
|
||||
# Train encoder with the corpus
|
||||
self._encoder.train(self._corpus)
|
||||
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"dashtext package is required for BM25EmbeddingFunction. "
|
||||
"Install it with: pip install dashtext"
|
||||
) from e
|
||||
except Exception as e:
|
||||
if isinstance(e, (ValueError, RuntimeError)):
|
||||
raise
|
||||
raise RuntimeError(f"Failed to build BM25 encoder: {e!s}") from e
|
||||
|
||||
@property
|
||||
def corpus_size(self) -> int:
|
||||
"""int: Number of documents in the training corpus (0 if using built-in encoder)."""
|
||||
return len(self._corpus) if self._corpus is not None else 0
|
||||
|
||||
@property
|
||||
def encoding_type(self) -> str:
|
||||
"""str: The encoding type being used ("query" or "document")."""
|
||||
return self._encoding_type
|
||||
|
||||
@property
|
||||
def language(self) -> str:
|
||||
"""str: The language of the built-in encoder ("zh" or "en")."""
|
||||
return self._language
|
||||
|
||||
@property
|
||||
def extra_params(self) -> dict:
|
||||
"""dict: Extra parameters for DashText encoder customization."""
|
||||
return self._extra_params
|
||||
|
||||
def __call__(self, input: TEXT) -> SparseVectorType:
|
||||
"""Make the embedding function callable.
|
||||
|
||||
Args:
|
||||
input (TEXT): Input text to embed.
|
||||
|
||||
Returns:
|
||||
SparseVectorType: Sparse vector as dictionary.
|
||||
"""
|
||||
return self.embed(input)
|
||||
|
||||
@lru_cache(maxsize=10)
|
||||
def embed(self, input: TEXT) -> SparseVectorType:
|
||||
"""Generate BM25 sparse embedding for the input text.
|
||||
|
||||
This method computes BM25 scores for the input text using DashText's
|
||||
SparseVectorEncoder. The encoding behavior depends on the encoding_type:
|
||||
|
||||
- ``encoding_type="query"``: Uses ``encode_queries()`` for search queries
|
||||
- ``encoding_type="document"``: Uses ``encode_documents()`` for documents
|
||||
|
||||
The result is a sparse vector where keys are term indices in the
|
||||
vocabulary and values are BM25 scores.
|
||||
|
||||
Args:
|
||||
input (TEXT): Input text string to embed. Must be non-empty after
|
||||
stripping whitespace.
|
||||
|
||||
Returns:
|
||||
SparseVectorType: A dictionary mapping vocabulary term index to BM25 score.
|
||||
Only non-zero scores are included. The dictionary is sorted by indices
|
||||
(keys) in ascending order for consistent output.
|
||||
Example: ``{1169440797: 0.29, 2045788977: 0.70, ...}``
|
||||
|
||||
Raises:
|
||||
TypeError: If ``input`` is not a string.
|
||||
ValueError: If input is empty or whitespace-only.
|
||||
RuntimeError: If BM25 encoding fails.
|
||||
|
||||
Examples:
|
||||
>>> bm25 = BM25EmbeddingFunction(language="zh", encoding_type="query")
|
||||
>>> sparse_vec = bm25.embed("query text")
|
||||
>>> isinstance(sparse_vec, dict)
|
||||
True
|
||||
>>> all(isinstance(k, int) and isinstance(v, float) for k, v in sparse_vec.items())
|
||||
True
|
||||
|
||||
>>> # Verify sorted output
|
||||
>>> keys = list(sparse_vec.keys())
|
||||
>>> keys == sorted(keys)
|
||||
True
|
||||
|
||||
>>> # Error: empty input
|
||||
>>> bm25.embed(" ")
|
||||
ValueError: Input text cannot be empty or whitespace only
|
||||
|
||||
>>> # Error: non-string input
|
||||
>>> bm25.embed(123)
|
||||
TypeError: Expected 'input' to be str, got int
|
||||
|
||||
Note:
|
||||
- BM25 scores are relative to the vocabulary statistics
|
||||
- Output dictionary is always sorted by indices for consistency
|
||||
- Terms not in the vocabulary will have zero scores (not included)
|
||||
- This method is cached (maxsize=10) for performance
|
||||
- DashText automatically handles Chinese/English text segmentation
|
||||
"""
|
||||
if not isinstance(input, str):
|
||||
raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}")
|
||||
|
||||
input = input.strip()
|
||||
if not input:
|
||||
raise ValueError("Input text cannot be empty or whitespace only")
|
||||
|
||||
try:
|
||||
# Encode based on encoding_type
|
||||
if self._encoding_type == "query":
|
||||
sparse_vector = self._encoder.encode_queries(input)
|
||||
else: # encoding_type == "document"
|
||||
sparse_vector = self._encoder.encode_documents(input)
|
||||
|
||||
# DashText returns dict with int/long keys and float values
|
||||
# Convert to standard format: {int: float}
|
||||
sparse_dict: dict[int, float] = {}
|
||||
for key, value in sparse_vector.items():
|
||||
try:
|
||||
idx = int(key)
|
||||
val = float(value)
|
||||
if val > 0:
|
||||
sparse_dict[idx] = val
|
||||
except (ValueError, TypeError):
|
||||
# Skip invalid entries
|
||||
continue
|
||||
|
||||
# Sort by indices (keys) to ensure consistent ordering
|
||||
return dict(sorted(sparse_dict.items()))
|
||||
|
||||
except Exception as e:
|
||||
if isinstance(e, (TypeError, ValueError)):
|
||||
raise
|
||||
raise RuntimeError(f"Failed to generate BM25 embedding: {e!s}") from e
|
||||
@@ -0,0 +1,147 @@
|
||||
# Copyright 2025-present the zvec project
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import abstractmethod
|
||||
from typing import Protocol, runtime_checkable
|
||||
|
||||
from ..common.constants import MD, DenseVectorType, SparseVectorType
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class DenseEmbeddingFunction(Protocol[MD]):
|
||||
"""Protocol for dense vector embedding functions.
|
||||
|
||||
Dense embedding functions map multimodal input (text, image, or audio) to
|
||||
fixed-length real-valued vectors. This is a Protocol class that defines
|
||||
the interface - implementations should provide their own initialization
|
||||
and properties.
|
||||
|
||||
Type Parameters:
|
||||
MD: The type of input data (bound to Embeddable: TEXT, IMAGE, or AUDIO).
|
||||
|
||||
Note:
|
||||
- This is a Protocol class - it only defines the ``embed()`` interface.
|
||||
- Implementations are free to define their own ``__init__``, properties,
|
||||
and additional methods as needed.
|
||||
- The ``embed()`` method is the only required interface.
|
||||
|
||||
Examples:
|
||||
>>> # Custom text embedding implementation
|
||||
>>> class MyTextEmbedding:
|
||||
... def __init__(self, dimension: int, model_name: str):
|
||||
... self.dimension = dimension
|
||||
... self.model = load_model(model_name)
|
||||
...
|
||||
... def embed(self, input: str) -> list[float]:
|
||||
... return self.model.encode(input).tolist()
|
||||
|
||||
>>> # Custom image embedding implementation
|
||||
>>> class MyImageEmbedding:
|
||||
... def __init__(self, dimension: int = 512):
|
||||
... self.dimension = dimension
|
||||
... self.model = load_image_model()
|
||||
...
|
||||
... def embed(self, input: Union[str, bytes, np.ndarray]) -> list[float]:
|
||||
... if isinstance(input, str):
|
||||
... image = load_image_from_path(input)
|
||||
... else:
|
||||
... image = input
|
||||
... return self.model.extract_features(image).tolist()
|
||||
|
||||
>>> # Using built-in implementations
|
||||
>>> from zvec.extension import QwenDenseEmbedding
|
||||
>>> text_emb = QwenDenseEmbedding(dimension=768, api_key="sk-xxx")
|
||||
>>> vector = text_emb.embed("Hello world")
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def embed(self, input: MD) -> DenseVectorType:
|
||||
"""Generate a dense embedding vector for the input data.
|
||||
|
||||
Args:
|
||||
input (MD): Multimodal input data to embed. Can be:
|
||||
- TEXT (str): Text string
|
||||
- IMAGE (str | bytes | np.ndarray): Image file path, raw bytes, or array
|
||||
- AUDIO (str | bytes | np.ndarray): Audio file path, raw bytes, or array
|
||||
|
||||
Returns:
|
||||
DenseVectorType: A dense vector representing the embedding.
|
||||
Can be list[float], list[int], or np.ndarray.
|
||||
Length should match the implementation's dimension.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SparseEmbeddingFunction(Protocol[MD]):
|
||||
"""Abstract base class for sparse vector embedding functions.
|
||||
|
||||
Sparse embedding functions map multimodal input (text, image, or audio) to
|
||||
a dictionary of {index: weight}, where only non-zero dimensions are stored.
|
||||
You can inherit this class to create custom sparse embedding functions.
|
||||
|
||||
Type Parameters:
|
||||
MD: The type of input data (bound to Embeddable: TEXT, IMAGE, or AUDIO).
|
||||
|
||||
Note:
|
||||
Subclasses must implement the ``embed()`` method.
|
||||
|
||||
Examples:
|
||||
>>> # Using built-in text sparse embedding (e.g., BM25, TF-IDF)
|
||||
>>> sparse_emb = SomeSparseEmbedding()
|
||||
>>> vector = sparse_emb.embed("Hello world")
|
||||
>>> # Returns: {0: 0.5, 42: 1.2, 100: 0.8}
|
||||
|
||||
>>> # Custom BM25 sparse embedding function
|
||||
>>> class MyBM25Embedding(SparseEmbeddingFunction):
|
||||
... def __init__(self, vocab_size: int = 10000):
|
||||
... self.vocab_size = vocab_size
|
||||
... self.tokenizer = MyTokenizer()
|
||||
...
|
||||
... def embed(self, input: str) -> dict[int, float]:
|
||||
... tokens = self.tokenizer.tokenize(input)
|
||||
... sparse_vector = {}
|
||||
... for token_id, weight in self._calculate_bm25(tokens):
|
||||
... if weight > 0:
|
||||
... sparse_vector[token_id] = weight
|
||||
... return sparse_vector
|
||||
...
|
||||
... def _calculate_bm25(self, tokens):
|
||||
... # BM25 calculation logic
|
||||
... pass
|
||||
|
||||
>>> # Custom sparse image feature extractor
|
||||
>>> class MySparseImageEmbedding(SparseEmbeddingFunction):
|
||||
... def embed(self, input: Union[str, bytes, np.ndarray]) -> dict[int, float]:
|
||||
... image = self._load_image(input)
|
||||
... features = self._extract_sparse_features(image)
|
||||
... return {idx: val for idx, val in enumerate(features) if val != 0}
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def embed(self, input: MD) -> SparseVectorType:
|
||||
"""Generate a sparse embedding for the input data.
|
||||
|
||||
Args:
|
||||
input (MD): Multimodal input data to embed. Can be:
|
||||
- TEXT (str): Text string
|
||||
- IMAGE (str | bytes | np.ndarray): Image file path, raw bytes, or array
|
||||
- AUDIO (str | bytes | np.ndarray): Audio file path, raw bytes, or array
|
||||
|
||||
Returns:
|
||||
SparseVectorType: Mapping from dimension index to non-zero weight.
|
||||
Only dimensions with non-zero values are included.
|
||||
"""
|
||||
...
|
||||
@@ -0,0 +1,162 @@
|
||||
# Copyright 2025-present the zvec project
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import urllib.request
|
||||
from functools import lru_cache
|
||||
from typing import Optional
|
||||
|
||||
from ..common.constants import TEXT, DenseVectorType
|
||||
from .embedding_function import DenseEmbeddingFunction
|
||||
|
||||
|
||||
class HTTPDenseEmbedding(DenseEmbeddingFunction[TEXT]):
|
||||
"""Dense text embedding function using any OpenAI-compatible HTTP endpoint.
|
||||
|
||||
This class calls any server that implements the ``/v1/embeddings`` API
|
||||
(LM Studio, Ollama, vLLM, LocalAI, etc.) using only the Python standard
|
||||
library — no extra dependencies are required.
|
||||
|
||||
The embedding dimension is detected automatically from the first server
|
||||
response.
|
||||
|
||||
Args:
|
||||
base_url (str, optional): Base URL of the embedding server.
|
||||
Defaults to ``"http://localhost:1234"`` (LM Studio).
|
||||
Common values:
|
||||
|
||||
- ``"http://localhost:1234"`` — LM Studio
|
||||
- ``"http://localhost:11434"`` — Ollama
|
||||
model (str, optional): Model identifier as expected by the server.
|
||||
Defaults to ``"text-embedding-nomic-embed-text-v1.5@f16"``.
|
||||
api_key (Optional[str], optional): Bearer token for authenticated
|
||||
endpoints. Falls back to the ``OPENAI_API_KEY`` environment
|
||||
variable. Leave as ``None`` for local servers that do not
|
||||
require authentication.
|
||||
timeout (int, optional): HTTP request timeout in seconds.
|
||||
Defaults to 30.
|
||||
|
||||
Attributes:
|
||||
dimension (int): Embedding vector dimensionality (auto-detected).
|
||||
|
||||
Raises:
|
||||
TypeError: If ``embed()`` receives a non-string input.
|
||||
ValueError: If input is empty/whitespace-only or the server returns
|
||||
an unexpected response format.
|
||||
RuntimeError: If the HTTP request fails or the server is unreachable.
|
||||
|
||||
Examples:
|
||||
>>> from zvec.extension import HTTPDenseEmbedding
|
||||
>>>
|
||||
>>> # LM Studio (default)
|
||||
>>> emb = HTTPDenseEmbedding()
|
||||
>>> vector = emb.embed("Hello, world!")
|
||||
>>> len(vector)
|
||||
768
|
||||
>>>
|
||||
>>> # Ollama
|
||||
>>> emb = HTTPDenseEmbedding(
|
||||
... base_url="http://localhost:11434",
|
||||
... model="nomic-embed-text",
|
||||
... )
|
||||
>>> vector = emb.embed("Semantic search with local models")
|
||||
|
||||
See Also:
|
||||
- ``DenseEmbeddingFunction``: Protocol for dense embeddings.
|
||||
- ``OpenAIDenseEmbedding``: Cloud embedding via the OpenAI API.
|
||||
"""
|
||||
|
||||
ENDPOINT = "/v1/embeddings"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str = "http://localhost:1234",
|
||||
model: str = "text-embedding-nomic-embed-text-v1.5@f16",
|
||||
api_key: Optional[str] = None,
|
||||
timeout: int = 30,
|
||||
) -> None:
|
||||
self._base_url = base_url.rstrip("/")
|
||||
self._model = model
|
||||
self._api_key = api_key or os.environ.get("OPENAI_API_KEY", "")
|
||||
self._timeout = timeout
|
||||
self._dimension: Optional[int] = None
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
"""int: Embedding vector dimensionality (auto-detected on first call)."""
|
||||
if self._dimension is None:
|
||||
self._dimension = len(self.embed("dimension probe"))
|
||||
return self._dimension
|
||||
|
||||
def __call__(self, input: TEXT) -> DenseVectorType:
|
||||
"""Make the embedding function callable."""
|
||||
return self.embed(input)
|
||||
|
||||
@lru_cache(maxsize=256)
|
||||
def embed(self, input: TEXT) -> DenseVectorType:
|
||||
"""Generate a dense embedding vector for the input text.
|
||||
|
||||
Results are cached (LRU, up to 256 entries) so repeated strings
|
||||
do not trigger extra HTTP requests.
|
||||
|
||||
Args:
|
||||
input (TEXT): Input text string to embed. Must be non-empty
|
||||
after stripping whitespace.
|
||||
|
||||
Returns:
|
||||
DenseVectorType: A list of floats representing the embedding.
|
||||
|
||||
Raises:
|
||||
TypeError: If *input* is not a string.
|
||||
ValueError: If *input* is empty/whitespace-only or the server
|
||||
returns an unexpected response format.
|
||||
RuntimeError: If the HTTP request fails.
|
||||
"""
|
||||
if not isinstance(input, TEXT):
|
||||
raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}")
|
||||
|
||||
input = input.strip()
|
||||
if not input:
|
||||
raise ValueError("Input text cannot be empty or whitespace only")
|
||||
|
||||
url = self._base_url + self.ENDPOINT
|
||||
payload = json.dumps({"model": self._model, "input": input}).encode()
|
||||
|
||||
headers: dict[str, str] = {"Content-Type": "application/json"}
|
||||
if self._api_key:
|
||||
headers["Authorization"] = f"Bearer {self._api_key}"
|
||||
|
||||
req = urllib.request.Request(url, data=payload, headers=headers, method="POST")
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=self._timeout) as resp:
|
||||
body = json.loads(resp.read())
|
||||
except urllib.error.HTTPError as exc:
|
||||
raise RuntimeError(
|
||||
f"Embedding server returned HTTP {exc.code}: {exc.read().decode()}"
|
||||
) from exc
|
||||
except OSError as exc:
|
||||
raise RuntimeError(
|
||||
f"Could not reach embedding server at {url}: {exc}"
|
||||
) from exc
|
||||
|
||||
try:
|
||||
vector: list[float] = body["data"][0]["embedding"]
|
||||
except (KeyError, IndexError) as exc:
|
||||
raise ValueError(
|
||||
f"Unexpected response format from embedding server: {body}"
|
||||
) from exc
|
||||
|
||||
return vector
|
||||
@@ -0,0 +1,240 @@
|
||||
# Copyright 2025-present the zvec project
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from functools import lru_cache
|
||||
from typing import Optional
|
||||
|
||||
from ..common.constants import TEXT, DenseVectorType
|
||||
from .embedding_function import DenseEmbeddingFunction
|
||||
from .jina_function import JinaFunctionBase
|
||||
|
||||
|
||||
class JinaDenseEmbedding(JinaFunctionBase, DenseEmbeddingFunction[TEXT]):
|
||||
"""Dense text embedding function using Jina AI API.
|
||||
|
||||
This class provides text-to-vector embedding capabilities using Jina AI's
|
||||
embedding models. It inherits from ``DenseEmbeddingFunction`` and implements
|
||||
dense text embedding via the Jina Embeddings API (OpenAI-compatible).
|
||||
|
||||
Jina Embeddings v5 models support task-specific embedding through the
|
||||
``task`` parameter, which optimizes the embedding for different use cases
|
||||
such as retrieval, text matching, or classification. They also support
|
||||
Matryoshka Representation Learning, allowing flexible output dimensions.
|
||||
|
||||
Args:
|
||||
model (str, optional): Jina embedding model identifier.
|
||||
Defaults to ``"jina-embeddings-v5-text-nano"``. Available models:
|
||||
- ``"jina-embeddings-v5-text-nano"``: 768 dims, 239M params, 8K context
|
||||
- ``"jina-embeddings-v5-text-small"``: 1024 dims, 677M params, 32K context
|
||||
dimension (Optional[int], optional): Desired output embedding dimension.
|
||||
If ``None``, uses model's default dimension. Supports Matryoshka
|
||||
dimensions: 32, 64, 128, 256, 512, 768 (nano) / 1024 (small).
|
||||
Defaults to ``None``.
|
||||
api_key (Optional[str], optional): Jina API authentication key.
|
||||
If ``None``, reads from ``JINA_API_KEY`` environment variable.
|
||||
Obtain your key from: https://jina.ai/api-dashboard
|
||||
task (Optional[str], optional): Task type to optimize embeddings for.
|
||||
Defaults to ``None``. Valid values:
|
||||
- ``"retrieval.query"``: For search queries
|
||||
- ``"retrieval.passage"``: For documents/passages to be searched
|
||||
- ``"text-matching"``: For symmetric text similarity
|
||||
- ``"classification"``: For text classification
|
||||
- ``"separation"``: For clustering/separation tasks
|
||||
|
||||
Attributes:
|
||||
dimension (int): The embedding vector dimension.
|
||||
data_type (DataType): Always ``DataType.VECTOR_FP32`` for this implementation.
|
||||
model (str): The Jina model name being used.
|
||||
task (Optional[str]): The task type for embedding optimization.
|
||||
|
||||
Raises:
|
||||
ValueError: If API key is not provided and not found in environment,
|
||||
if task is not a valid task type, or if API returns an error response.
|
||||
TypeError: If input to ``embed()`` is not a string.
|
||||
RuntimeError: If network error or Jina service error occurs.
|
||||
|
||||
Note:
|
||||
- Requires Python 3.10, 3.11, or 3.12
|
||||
- Requires the ``openai`` package: ``pip install openai``
|
||||
- Jina API is OpenAI-compatible, so it uses the ``openai`` Python client
|
||||
- Embedding results are cached (LRU cache, maxsize=10) to reduce API calls
|
||||
- For retrieval tasks, use ``"retrieval.query"`` for queries and
|
||||
``"retrieval.passage"`` for documents
|
||||
- API usage requires a Jina API key from https://jina.ai/api-dashboard
|
||||
|
||||
Examples:
|
||||
>>> # Basic usage with default model
|
||||
>>> from zvec.extension import JinaDenseEmbedding
|
||||
>>> import os
|
||||
>>> os.environ["JINA_API_KEY"] = "jina_..."
|
||||
>>>
|
||||
>>> emb_func = JinaDenseEmbedding()
|
||||
>>> vector = emb_func.embed("Hello, world!")
|
||||
>>> len(vector)
|
||||
768
|
||||
|
||||
>>> # Retrieval use case: embed queries and documents differently
|
||||
>>> query_emb = JinaDenseEmbedding(task="retrieval.query")
|
||||
>>> doc_emb = JinaDenseEmbedding(task="retrieval.passage")
|
||||
>>>
|
||||
>>> query_vector = query_emb.embed("What is machine learning?")
|
||||
>>> doc_vector = doc_emb.embed("Machine learning is a subset of AI...")
|
||||
|
||||
>>> # Using larger model with custom dimension (Matryoshka)
|
||||
>>> emb_func = JinaDenseEmbedding(
|
||||
... model="jina-embeddings-v5-text-small",
|
||||
... dimension=256,
|
||||
... api_key="jina_...",
|
||||
... task="text-matching",
|
||||
... )
|
||||
>>> vector = emb_func.embed("Semantic similarity comparison")
|
||||
>>> len(vector)
|
||||
256
|
||||
|
||||
>>> # Using with zvec collection
|
||||
>>> import zvec
|
||||
>>> emb_func = JinaDenseEmbedding(task="retrieval.passage")
|
||||
>>> schema = zvec.CollectionSchema(
|
||||
... name="docs",
|
||||
... vectors=zvec.VectorSchema(
|
||||
... "embedding", zvec.DataType.VECTOR_FP32, emb_func.dimension
|
||||
... ),
|
||||
... )
|
||||
>>> collection = zvec.create_and_open(path="./my_docs", schema=schema)
|
||||
|
||||
See Also:
|
||||
- ``DenseEmbeddingFunction``: Base class for dense embeddings
|
||||
- ``OpenAIDenseEmbedding``: Alternative using OpenAI API
|
||||
- ``QwenDenseEmbedding``: Alternative using Qwen/DashScope API
|
||||
- ``DefaultLocalDenseEmbedding``: Local model without API calls
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str = "jina-embeddings-v5-text-nano",
|
||||
dimension: Optional[int] = None,
|
||||
api_key: Optional[str] = None,
|
||||
task: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize the Jina dense embedding function.
|
||||
|
||||
Args:
|
||||
model (str): Jina model name. Defaults to "jina-embeddings-v5-text-nano".
|
||||
dimension (Optional[int]): Target embedding dimension or None for default.
|
||||
api_key (Optional[str]): API key or None to use environment variable.
|
||||
task (Optional[str]): Task type for embedding optimization or None.
|
||||
**kwargs: Additional parameters for API calls.
|
||||
|
||||
Raises:
|
||||
ValueError: If API key is not provided and not in environment,
|
||||
or if task is not a valid task type.
|
||||
"""
|
||||
# Initialize base class for API connection
|
||||
JinaFunctionBase.__init__(self, model=model, api_key=api_key, task=task)
|
||||
|
||||
# Store dimension configuration
|
||||
self._custom_dimension = dimension
|
||||
|
||||
# Determine actual dimension
|
||||
if dimension is None:
|
||||
self._dimension = self._MODEL_DIMENSIONS.get(model, 768)
|
||||
else:
|
||||
self._dimension = dimension
|
||||
|
||||
# Store extra attributes
|
||||
self._extra_params = kwargs
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
"""int: The expected dimensionality of the embedding vector."""
|
||||
return self._dimension
|
||||
|
||||
@property
|
||||
def extra_params(self) -> dict:
|
||||
"""dict: Extra parameters for model-specific customization."""
|
||||
return self._extra_params
|
||||
|
||||
def __call__(self, input: TEXT) -> DenseVectorType:
|
||||
"""Make the embedding function callable."""
|
||||
return self.embed(input)
|
||||
|
||||
@lru_cache(maxsize=10)
|
||||
def embed(self, input: TEXT) -> DenseVectorType:
|
||||
"""Generate dense embedding vector for the input text.
|
||||
|
||||
This method calls the Jina Embeddings API to convert input text
|
||||
into a dense vector representation. Results are cached to improve
|
||||
performance for repeated inputs.
|
||||
|
||||
Args:
|
||||
input (TEXT): Input text string to embed. Must be non-empty after
|
||||
stripping whitespace. Maximum length depends on model:
|
||||
8192 tokens for v5-nano, 32768 tokens for v5-small.
|
||||
|
||||
Returns:
|
||||
DenseVectorType: A list of floats representing the embedding vector.
|
||||
Length equals ``self.dimension``. Example:
|
||||
``[0.123, -0.456, 0.789, ...]``
|
||||
|
||||
Raises:
|
||||
TypeError: If ``input`` is not a string.
|
||||
ValueError: If input is empty/whitespace-only, or if the API returns
|
||||
an error or malformed response.
|
||||
RuntimeError: If network connectivity issues or Jina service
|
||||
errors occur.
|
||||
|
||||
Examples:
|
||||
>>> emb = JinaDenseEmbedding(task="retrieval.query")
|
||||
>>> vector = emb.embed("What is deep learning?")
|
||||
>>> len(vector)
|
||||
768
|
||||
>>> isinstance(vector[0], float)
|
||||
True
|
||||
|
||||
>>> # Error: empty input
|
||||
>>> emb.embed(" ")
|
||||
ValueError: Input text cannot be empty or whitespace only
|
||||
|
||||
>>> # Error: non-string input
|
||||
>>> emb.embed(123)
|
||||
TypeError: Expected 'input' to be str, got int
|
||||
|
||||
Note:
|
||||
- This method is cached (maxsize=10). Identical inputs return cached results.
|
||||
- The cache is based on exact string match (case-sensitive).
|
||||
- Task type affects embedding optimization but not caching behavior.
|
||||
"""
|
||||
if not isinstance(input, TEXT):
|
||||
raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}")
|
||||
|
||||
input = input.strip()
|
||||
if not input:
|
||||
raise ValueError("Input text cannot be empty or whitespace only")
|
||||
|
||||
# Call API
|
||||
embedding_vector = self._call_text_embedding_api(
|
||||
input=input,
|
||||
dimension=self._custom_dimension,
|
||||
)
|
||||
|
||||
# Verify dimension
|
||||
if len(embedding_vector) != self.dimension:
|
||||
raise ValueError(
|
||||
f"Dimension mismatch: expected {self.dimension}, "
|
||||
f"got {len(embedding_vector)}"
|
||||
)
|
||||
|
||||
return embedding_vector
|
||||
@@ -0,0 +1,182 @@
|
||||
# Copyright 2025-present the zvec project
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import ClassVar, Optional
|
||||
|
||||
from ..common.constants import TEXT
|
||||
from ..tool import require_module
|
||||
|
||||
|
||||
class JinaFunctionBase:
|
||||
"""Base class for Jina AI functions.
|
||||
|
||||
This base class provides common functionality for calling Jina AI APIs
|
||||
and handling responses. It supports embeddings (dense) operations via
|
||||
the OpenAI-compatible Jina Embeddings API.
|
||||
|
||||
This class is not meant to be used directly. Use concrete implementations:
|
||||
- ``JinaDenseEmbedding`` for dense embeddings
|
||||
|
||||
Args:
|
||||
model (str): Jina embedding model identifier.
|
||||
api_key (Optional[str]): Jina API authentication key.
|
||||
task (Optional[str]): Task type for the embedding model.
|
||||
|
||||
Note:
|
||||
- This is an internal base class for code reuse across Jina features
|
||||
- Subclasses should inherit from appropriate Protocol
|
||||
- Provides unified API connection and response handling
|
||||
- Jina API is OpenAI-compatible, using the ``openai`` Python client
|
||||
"""
|
||||
|
||||
_BASE_URL: ClassVar[str] = "https://api.jina.ai/v1"
|
||||
|
||||
# Model default dimensions
|
||||
_MODEL_DIMENSIONS: ClassVar[dict[str, int]] = {
|
||||
"jina-embeddings-v5-text-nano": 768,
|
||||
"jina-embeddings-v5-text-small": 1024,
|
||||
}
|
||||
|
||||
# Model max tokens
|
||||
_MODEL_MAX_TOKENS: ClassVar[dict[str, int]] = {
|
||||
"jina-embeddings-v5-text-nano": 8192,
|
||||
"jina-embeddings-v5-text-small": 32768,
|
||||
}
|
||||
|
||||
# Valid task types
|
||||
_VALID_TASKS: ClassVar[tuple[str, ...]] = (
|
||||
"retrieval.query",
|
||||
"retrieval.passage",
|
||||
"text-matching",
|
||||
"classification",
|
||||
"separation",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str,
|
||||
api_key: Optional[str] = None,
|
||||
task: Optional[str] = None,
|
||||
):
|
||||
"""Initialize the base Jina functionality.
|
||||
|
||||
Args:
|
||||
model (str): Jina model name.
|
||||
api_key (Optional[str]): API key or None to use environment variable.
|
||||
task (Optional[str]): Task type for the embedding model.
|
||||
Valid values: "retrieval.query", "retrieval.passage",
|
||||
"text-matching", "classification", "separation".
|
||||
|
||||
Raises:
|
||||
ValueError: If API key is not provided and not in environment,
|
||||
or if task is not a valid task type.
|
||||
"""
|
||||
self._model = model
|
||||
self._api_key = api_key or os.environ.get("JINA_API_KEY")
|
||||
self._task = task
|
||||
|
||||
if not self._api_key:
|
||||
raise ValueError(
|
||||
"Jina API key is required. Please provide 'api_key' parameter "
|
||||
"or set the 'JINA_API_KEY' environment variable. "
|
||||
"Get your key from: https://jina.ai/api-dashboard"
|
||||
)
|
||||
|
||||
if task is not None and task not in self._VALID_TASKS:
|
||||
raise ValueError(
|
||||
f"Invalid task '{task}'. Valid tasks: {', '.join(self._VALID_TASKS)}"
|
||||
)
|
||||
|
||||
@property
|
||||
def model(self) -> str:
|
||||
"""str: The Jina model name currently in use."""
|
||||
return self._model
|
||||
|
||||
@property
|
||||
def task(self) -> Optional[str]:
|
||||
"""Optional[str]: The task type for the embedding model."""
|
||||
return self._task
|
||||
|
||||
def _get_client(self):
|
||||
"""Get OpenAI-compatible client instance configured for Jina API.
|
||||
|
||||
Returns:
|
||||
OpenAI: Configured OpenAI client pointing to Jina API.
|
||||
|
||||
Raises:
|
||||
ImportError: If openai package is not installed.
|
||||
"""
|
||||
openai = require_module("openai")
|
||||
return openai.OpenAI(api_key=self._api_key, base_url=self._BASE_URL)
|
||||
|
||||
def _call_text_embedding_api(
|
||||
self,
|
||||
input: TEXT,
|
||||
dimension: Optional[int] = None,
|
||||
) -> list:
|
||||
"""Call Jina Embeddings API.
|
||||
|
||||
Args:
|
||||
input (TEXT): Input text to embed.
|
||||
dimension (Optional[int]): Target dimension for Matryoshka embeddings.
|
||||
|
||||
Returns:
|
||||
list: Embedding vector as list of floats.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If API call fails.
|
||||
ValueError: If API returns error response.
|
||||
"""
|
||||
try:
|
||||
client = self._get_client()
|
||||
|
||||
# Prepare embedding parameters
|
||||
params = {"model": self.model, "input": input}
|
||||
|
||||
# Add dimension parameter for Matryoshka support
|
||||
if dimension is not None:
|
||||
params["dimensions"] = dimension
|
||||
|
||||
# Add task parameter via extra_body
|
||||
if self._task is not None:
|
||||
params["extra_body"] = {"task": self._task}
|
||||
|
||||
# Call Jina API (OpenAI-compatible)
|
||||
response = client.embeddings.create(**params)
|
||||
|
||||
except Exception as e:
|
||||
# Check if it's an OpenAI API error
|
||||
openai = require_module("openai")
|
||||
if isinstance(e, (openai.APIError, openai.APIConnectionError)):
|
||||
raise RuntimeError(f"Failed to call Jina API: {e!s}") from e
|
||||
raise RuntimeError(f"Unexpected error during API call: {e!s}") from e
|
||||
|
||||
# Extract embedding from response
|
||||
try:
|
||||
if not response.data:
|
||||
raise ValueError("Invalid API response: no embedding data returned")
|
||||
|
||||
embedding_vector = response.data[0].embedding
|
||||
|
||||
if not isinstance(embedding_vector, list):
|
||||
raise ValueError(
|
||||
"Invalid API response: embedding is not a list of numbers"
|
||||
)
|
||||
|
||||
return embedding_vector
|
||||
|
||||
except (AttributeError, IndexError, TypeError) as e:
|
||||
raise ValueError(f"Failed to parse API response: {e!s}") from e
|
||||
@@ -0,0 +1,197 @@
|
||||
# Copyright 2025-present the zvec project
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from zvec._zvec import (
|
||||
_CallbackParams,
|
||||
_Doc,
|
||||
_reranker_rerank,
|
||||
_RrfParams,
|
||||
_WeightedParams,
|
||||
)
|
||||
|
||||
from ..model.doc import Doc, DocList
|
||||
from .rerank_function import RerankFunction
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..model.schema import FieldSchema, VectorSchema
|
||||
|
||||
|
||||
def _to_cpp_doc_lists(
|
||||
query_results: list[list[Doc]],
|
||||
) -> tuple[list[list], dict[str, Doc]]:
|
||||
"""Convert Python Doc lists to C++ _Doc lists for reranker input."""
|
||||
id_to_doc: dict[str, Doc] = {}
|
||||
cpp_results: list[list] = []
|
||||
for query_result in query_results:
|
||||
cpp_list: list = []
|
||||
for doc in query_result:
|
||||
_doc = _Doc()
|
||||
_doc.set_pk(doc.id)
|
||||
_doc.set_score(doc.score if doc.score is not None else 0.0)
|
||||
cpp_list.append(_doc)
|
||||
if doc.id not in id_to_doc:
|
||||
id_to_doc[doc.id] = doc
|
||||
cpp_results.append(cpp_list)
|
||||
return cpp_results, id_to_doc
|
||||
|
||||
|
||||
def _from_cpp_docs(cpp_docs: list, id_to_doc: dict[str, Doc]) -> DocList:
|
||||
"""Convert C++ rerank result _Doc list back to Python DocList."""
|
||||
results: DocList = []
|
||||
for _doc in cpp_docs:
|
||||
doc_id = _doc.pk()
|
||||
new_score = _doc.score()
|
||||
original = id_to_doc.get(doc_id)
|
||||
if original is not None:
|
||||
results.append(original._replace(score=new_score))
|
||||
else:
|
||||
results.append(Doc(id=doc_id, score=new_score))
|
||||
return results
|
||||
|
||||
|
||||
class RrfReRanker(RerankFunction):
|
||||
"""Re-ranker using Reciprocal Rank Fusion (RRF) for multi-vector search.
|
||||
|
||||
RRF combines results from multiple vector queries without requiring
|
||||
relevance scores. The RRF score for a document at rank r is:
|
||||
score = 1 / (k + r + 1)
|
||||
where k is the rank constant.
|
||||
|
||||
Args:
|
||||
rank_constant: RRF smoothing constant (default: 60).
|
||||
Higher values reduce the influence of rank position.
|
||||
|
||||
Example:
|
||||
>>> reranker = RrfReRanker(rank_constant=60)
|
||||
>>> merged = reranker.rerank([results_a, results_b], topn=10)
|
||||
"""
|
||||
|
||||
def __init__(self, rank_constant: int = 60):
|
||||
self._rank_constant = rank_constant
|
||||
|
||||
@property
|
||||
def rank_constant(self) -> int:
|
||||
"""int: RRF rank constant."""
|
||||
return self._rank_constant
|
||||
|
||||
def _to_cpp_params(self):
|
||||
return _RrfParams(self._rank_constant)
|
||||
|
||||
def rerank(
|
||||
self,
|
||||
query_results: list[list[Doc]],
|
||||
topn: int = 10,
|
||||
*,
|
||||
fields: list[FieldSchema | VectorSchema] | None = None, # noqa: ARG002
|
||||
) -> DocList:
|
||||
"""Apply RRF to combine multiple query results via C++ reranker."""
|
||||
cpp_results, id_to_doc = _to_cpp_doc_lists(query_results)
|
||||
cpp_docs = _reranker_rerank(self._to_cpp_params(), cpp_results, [], topn)
|
||||
return _from_cpp_docs(cpp_docs, id_to_doc)
|
||||
|
||||
|
||||
class WeightedReRanker(RerankFunction):
|
||||
"""Re-ranker that combines scores using per-sub-query weights.
|
||||
|
||||
Each sub-query's score is normalized by metric type (automatic when used
|
||||
via collection.multi_query), then multiplied by the corresponding weight.
|
||||
|
||||
Args:
|
||||
weights: Per-sub-query weights. Length must match the number of
|
||||
sub-queries.
|
||||
|
||||
Example:
|
||||
>>> reranker = WeightedReRanker([0.7, 0.3])
|
||||
>>> merged = reranker.rerank([results_a, results_b], topn=10,
|
||||
... fields=field_schemas)
|
||||
"""
|
||||
|
||||
def __init__(self, weights: list[float]):
|
||||
self._weights = list(weights)
|
||||
|
||||
@property
|
||||
def weights(self) -> list[float]:
|
||||
"""list[float]: Per-sub-query weights."""
|
||||
return self._weights
|
||||
|
||||
def _to_cpp_params(self):
|
||||
return _WeightedParams(self._weights)
|
||||
|
||||
def rerank(
|
||||
self,
|
||||
query_results: list[list[Doc]],
|
||||
topn: int = 10,
|
||||
*,
|
||||
fields: list[FieldSchema | VectorSchema] | None = None,
|
||||
) -> DocList:
|
||||
"""Combine scores from multiple sub-queries using weighted sum via C++ reranker.
|
||||
|
||||
Args:
|
||||
query_results: Per-sub-query document lists.
|
||||
topn: Maximum results to return.
|
||||
fields: Per-sub-query Python FieldSchema/VectorSchema objects
|
||||
(required for score normalization by metric type).
|
||||
|
||||
Raises:
|
||||
ValueError: If fields is None (required for normalization).
|
||||
"""
|
||||
if not fields:
|
||||
raise ValueError(
|
||||
"WeightedReRanker.rerank() requires 'fields' for score normalization. "
|
||||
"Pass field schemas via fields= parameter."
|
||||
)
|
||||
cpp_fields = [f._get_object() for f in fields]
|
||||
cpp_results, id_to_doc = _to_cpp_doc_lists(query_results)
|
||||
cpp_docs = _reranker_rerank(
|
||||
self._to_cpp_params(), cpp_results, cpp_fields, topn
|
||||
)
|
||||
return _from_cpp_docs(cpp_docs, id_to_doc)
|
||||
|
||||
|
||||
class CallbackReRanker(RerankFunction):
|
||||
"""Re-ranker that delegates to a user-provided callback.
|
||||
|
||||
The callback receives sub-query results, field schemas, and topn.
|
||||
|
||||
Args:
|
||||
callback: A callable with signature
|
||||
(results: list[list[Doc]], fields: list, topn: int) -> list[Doc]
|
||||
|
||||
Example:
|
||||
>>> def my_rerank(results, fields, topn):
|
||||
... # custom logic
|
||||
... return merged[:topn]
|
||||
>>> reranker = CallbackReRanker(my_rerank)
|
||||
>>> merged = reranker.rerank([results_a, results_b], topn=10)
|
||||
"""
|
||||
|
||||
def __init__(self, callback: Callable):
|
||||
self._callback = callback
|
||||
|
||||
def _to_cpp_params(self):
|
||||
return _CallbackParams(self._callback)
|
||||
|
||||
def rerank(
|
||||
self,
|
||||
query_results: list[list[Doc]],
|
||||
topn: int = 10,
|
||||
*,
|
||||
fields: list[FieldSchema | VectorSchema] | None = None,
|
||||
) -> DocList:
|
||||
"""Invoke the callback to re-rank documents."""
|
||||
return self._callback(query_results, fields, topn)
|
||||
@@ -0,0 +1,238 @@
|
||||
# Copyright 2025-present the zvec project
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from functools import lru_cache
|
||||
from typing import Optional
|
||||
|
||||
from ..common.constants import TEXT, DenseVectorType
|
||||
from .embedding_function import DenseEmbeddingFunction
|
||||
from .openai_function import OpenAIFunctionBase
|
||||
|
||||
|
||||
class OpenAIDenseEmbedding(OpenAIFunctionBase, DenseEmbeddingFunction[TEXT]):
|
||||
"""Dense text embedding function using OpenAI API.
|
||||
|
||||
This class provides text-to-vector embedding capabilities using OpenAI's
|
||||
embedding models. It inherits from ``DenseEmbeddingFunction`` and implements
|
||||
dense text embedding via the OpenAI API.
|
||||
|
||||
The implementation supports various OpenAI embedding models with different
|
||||
dimensions and includes automatic result caching for improved performance.
|
||||
|
||||
Args:
|
||||
model (str, optional): OpenAI embedding model identifier.
|
||||
Defaults to ``"text-embedding-3-small"``. Common options:
|
||||
- ``"text-embedding-3-small"``: 1536 dims, cost-efficient, good performance
|
||||
- ``"text-embedding-3-large"``: 3072 dims, highest quality
|
||||
- ``"text-embedding-ada-002"``: 1536 dims, legacy model
|
||||
dimension (Optional[int], optional): Desired output embedding dimension.
|
||||
If ``None``, uses model's default dimension. For text-embedding-3 models,
|
||||
you can specify custom dimensions (e.g., 256, 512, 1024, 1536).
|
||||
Defaults to ``None``.
|
||||
api_key (Optional[str], optional): OpenAI API authentication key.
|
||||
If ``None``, reads from ``OPENAI_API_KEY`` environment variable.
|
||||
Obtain your key from: https://platform.openai.com/api-keys
|
||||
base_url (Optional[str], optional): Custom API base URL for OpenAI-compatible
|
||||
services. Defaults to ``None`` (uses official OpenAI endpoint).
|
||||
|
||||
Attributes:
|
||||
dimension (int): The embedding vector dimension.
|
||||
data_type (DataType): Always ``DataType.VECTOR_FP32`` for this implementation.
|
||||
model (str): The OpenAI model name being used.
|
||||
|
||||
Raises:
|
||||
ValueError: If API key is not provided and not found in environment,
|
||||
or if API returns an error response.
|
||||
TypeError: If input to ``embed()`` is not a string.
|
||||
RuntimeError: If network error or OpenAI service error occurs.
|
||||
|
||||
Note:
|
||||
- Requires Python 3.10, 3.11, or 3.12
|
||||
- Requires the ``openai`` package: ``pip install openai``
|
||||
- Embedding results are cached (LRU cache, maxsize=10) to reduce API calls
|
||||
- Network connectivity to OpenAI API endpoints is required
|
||||
- API usage incurs costs based on your OpenAI subscription plan
|
||||
- Rate limits apply based on your OpenAI account tier
|
||||
|
||||
Examples:
|
||||
>>> # Basic usage with default model
|
||||
>>> from zvec.extension import OpenAIDenseEmbedding
|
||||
>>> import os
|
||||
>>> os.environ["OPENAI_API_KEY"] = "sk-..."
|
||||
>>>
|
||||
>>> emb_func = OpenAIDenseEmbedding()
|
||||
>>> vector = emb_func.embed("Hello, world!")
|
||||
>>> len(vector)
|
||||
1536
|
||||
|
||||
>>> # Using specific model with custom dimension
|
||||
>>> emb_func = OpenAIDenseEmbedding(
|
||||
... model="text-embedding-3-large",
|
||||
... dimension=1024,
|
||||
... api_key="sk-..."
|
||||
... )
|
||||
>>> vector = emb_func.embed("Machine learning is fascinating")
|
||||
>>> len(vector)
|
||||
1024
|
||||
|
||||
>>> # Using with custom base URL (e.g., Azure OpenAI)
|
||||
>>> emb_func = OpenAIDenseEmbedding(
|
||||
... model="text-embedding-ada-002",
|
||||
... api_key="your-azure-key",
|
||||
... base_url="https://your-resource.openai.azure.com/"
|
||||
... )
|
||||
>>> vector = emb_func("Natural language processing")
|
||||
>>> isinstance(vector, list)
|
||||
True
|
||||
|
||||
>>> # Batch processing with caching benefit
|
||||
>>> texts = ["First text", "Second text", "First text"]
|
||||
>>> vectors = [emb_func.embed(text) for text in texts]
|
||||
>>> # Third call uses cached result for "First text"
|
||||
|
||||
>>> # Error handling
|
||||
>>> try:
|
||||
... emb_func.embed("") # Empty string
|
||||
... except ValueError as e:
|
||||
... print(f"Error: {e}")
|
||||
Error: Input text cannot be empty or whitespace only
|
||||
|
||||
See Also:
|
||||
- ``DenseEmbeddingFunction``: Base class for dense embeddings
|
||||
- ``QwenDenseEmbedding``: Alternative using Qwen/DashScope API
|
||||
- ``DefaultDenseEmbedding``: Local model without API calls
|
||||
- ``SparseEmbeddingFunction``: Base class for sparse embeddings
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str = "text-embedding-3-small",
|
||||
dimension: Optional[int] = None,
|
||||
api_key: Optional[str] = None,
|
||||
base_url: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize the OpenAI dense embedding function.
|
||||
|
||||
Args:
|
||||
model (str): OpenAI model name. Defaults to "text-embedding-3-small".
|
||||
dimension (Optional[int]): Target embedding dimension or None for default.
|
||||
api_key (Optional[str]): API key or None to use environment variable.
|
||||
base_url (Optional[str]): Custom API base URL or None for default.
|
||||
**kwargs: Additional parameters for API calls. Examples:
|
||||
- ``encoding_format`` (str): Format of embeddings, "float" or "base64".
|
||||
- ``user`` (str): User identifier for tracking.
|
||||
|
||||
Raises:
|
||||
ValueError: If API key is not provided and not in environment.
|
||||
"""
|
||||
# Initialize base class for API connection
|
||||
OpenAIFunctionBase.__init__(
|
||||
self, model=model, api_key=api_key, base_url=base_url
|
||||
)
|
||||
|
||||
# Store dimension configuration
|
||||
self._custom_dimension = dimension
|
||||
|
||||
# Determine actual dimension
|
||||
if dimension is None:
|
||||
# Use model default dimension
|
||||
self._dimension = self._MODEL_DIMENSIONS.get(model, 1536)
|
||||
else:
|
||||
self._dimension = dimension
|
||||
|
||||
# Store dense-specific attributes
|
||||
self._extra_params = kwargs
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
"""int: The expected dimensionality of the embedding vector."""
|
||||
return self._dimension
|
||||
|
||||
@property
|
||||
def extra_params(self) -> dict:
|
||||
"""dict: Extra parameters for model-specific customization."""
|
||||
return self._extra_params
|
||||
|
||||
def __call__(self, input: TEXT) -> DenseVectorType:
|
||||
"""Make the embedding function callable."""
|
||||
return self.embed(input)
|
||||
|
||||
@lru_cache(maxsize=10)
|
||||
def embed(self, input: TEXT) -> DenseVectorType:
|
||||
"""Generate dense embedding vector for the input text.
|
||||
|
||||
This method calls the OpenAI Embeddings API to convert input text
|
||||
into a dense vector representation. Results are cached to improve
|
||||
performance for repeated inputs.
|
||||
|
||||
Args:
|
||||
input (TEXT): Input text string to embed. Must be non-empty after
|
||||
stripping whitespace. Maximum length is 8191 tokens for most models.
|
||||
|
||||
Returns:
|
||||
DenseVectorType: A list of floats representing the embedding vector.
|
||||
Length equals ``self.dimension``. Example:
|
||||
``[0.123, -0.456, 0.789, ...]``
|
||||
|
||||
Raises:
|
||||
TypeError: If ``input`` is not a string.
|
||||
ValueError: If input is empty/whitespace-only, or if the API returns
|
||||
an error or malformed response.
|
||||
RuntimeError: If network connectivity issues or OpenAI service
|
||||
errors occur.
|
||||
|
||||
Examples:
|
||||
>>> emb = OpenAIDenseEmbedding()
|
||||
>>> vector = emb.embed("Natural language processing")
|
||||
>>> len(vector)
|
||||
1536
|
||||
>>> isinstance(vector[0], float)
|
||||
True
|
||||
|
||||
>>> # Error: empty input
|
||||
>>> emb.embed(" ")
|
||||
ValueError: Input text cannot be empty or whitespace only
|
||||
|
||||
>>> # Error: non-string input
|
||||
>>> emb.embed(123)
|
||||
TypeError: Expected 'input' to be str, got int
|
||||
|
||||
Note:
|
||||
- This method is cached (maxsize=10). Identical inputs return cached results.
|
||||
- The cache is based on exact string match (case-sensitive).
|
||||
- Consider pre-processing text (lowercasing, normalization) for better caching.
|
||||
"""
|
||||
if not isinstance(input, TEXT):
|
||||
raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}")
|
||||
|
||||
input = input.strip()
|
||||
if not input:
|
||||
raise ValueError("Input text cannot be empty or whitespace only")
|
||||
|
||||
# Call API
|
||||
embedding_vector = self._call_text_embedding_api(
|
||||
input=input,
|
||||
dimension=self._custom_dimension,
|
||||
)
|
||||
|
||||
# Verify dimension
|
||||
if len(embedding_vector) != self.dimension:
|
||||
raise ValueError(
|
||||
f"Dimension mismatch: expected {self.dimension}, "
|
||||
f"got {len(embedding_vector)}"
|
||||
)
|
||||
|
||||
return embedding_vector
|
||||
@@ -0,0 +1,149 @@
|
||||
# Copyright 2025-present the zvec project
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import ClassVar, Optional
|
||||
|
||||
from ..common.constants import TEXT
|
||||
from ..tool import require_module
|
||||
|
||||
|
||||
class OpenAIFunctionBase:
|
||||
"""Base class for OpenAI functions.
|
||||
|
||||
This base class provides common functionality for calling OpenAI APIs
|
||||
and handling responses. It supports embeddings (dense) operations.
|
||||
|
||||
This class is not meant to be used directly. Use concrete implementations:
|
||||
- ``OpenAIDenseEmbedding`` for dense embeddings
|
||||
|
||||
Args:
|
||||
model (str): OpenAI model identifier.
|
||||
api_key (Optional[str]): OpenAI API authentication key.
|
||||
base_url (Optional[str]): Custom API base URL.
|
||||
|
||||
Note:
|
||||
- This is an internal base class for code reuse across OpenAI features
|
||||
- Subclasses should inherit from appropriate Protocol
|
||||
- Provides unified API connection and response handling
|
||||
"""
|
||||
|
||||
# Model default dimensions
|
||||
_MODEL_DIMENSIONS: ClassVar[dict[str, int]] = {
|
||||
"text-embedding-3-small": 1536,
|
||||
"text-embedding-3-large": 3072,
|
||||
"text-embedding-ada-002": 1536,
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str,
|
||||
api_key: Optional[str] = None,
|
||||
base_url: Optional[str] = None,
|
||||
):
|
||||
"""Initialize the base OpenAI functionality.
|
||||
|
||||
Args:
|
||||
model (str): OpenAI model name.
|
||||
api_key (Optional[str]): API key or None to use environment variable.
|
||||
base_url (Optional[str]): Custom API base URL or None for default.
|
||||
|
||||
Raises:
|
||||
ValueError: If API key is not provided and not in environment.
|
||||
"""
|
||||
self._model = model
|
||||
self._api_key = api_key or os.environ.get("OPENAI_API_KEY")
|
||||
self._base_url = base_url
|
||||
|
||||
if not self._api_key:
|
||||
raise ValueError(
|
||||
"OpenAI API key is required. Please provide 'api_key' parameter "
|
||||
"or set the 'OPENAI_API_KEY' environment variable."
|
||||
)
|
||||
|
||||
@property
|
||||
def model(self) -> str:
|
||||
"""str: The OpenAI model name currently in use."""
|
||||
return self._model
|
||||
|
||||
def _get_client(self):
|
||||
"""Get OpenAI client instance.
|
||||
|
||||
Returns:
|
||||
OpenAI: Configured OpenAI client.
|
||||
|
||||
Raises:
|
||||
ImportError: If openai package is not installed.
|
||||
"""
|
||||
openai = require_module("openai")
|
||||
|
||||
if self._base_url:
|
||||
return openai.OpenAI(api_key=self._api_key, base_url=self._base_url)
|
||||
return openai.OpenAI(api_key=self._api_key)
|
||||
|
||||
def _call_text_embedding_api(
|
||||
self,
|
||||
input: TEXT,
|
||||
dimension: Optional[int] = None,
|
||||
) -> list:
|
||||
"""Call OpenAI Embeddings API.
|
||||
|
||||
Args:
|
||||
input (TEXT): Input text to embed.
|
||||
dimension (Optional[int]): Target dimension (for models that support it).
|
||||
|
||||
Returns:
|
||||
list: Embedding vector as list of floats.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If API call fails.
|
||||
ValueError: If API returns error response.
|
||||
"""
|
||||
try:
|
||||
client = self._get_client()
|
||||
|
||||
# Prepare embedding parameters
|
||||
params = {"model": self.model, "input": input}
|
||||
|
||||
# Add dimension parameter for models that support it
|
||||
if dimension is not None:
|
||||
params["dimensions"] = dimension
|
||||
|
||||
# Call OpenAI API
|
||||
response = client.embeddings.create(**params)
|
||||
|
||||
except Exception as e:
|
||||
# Check if it's an OpenAI API error
|
||||
openai = require_module("openai")
|
||||
if isinstance(e, (openai.APIError, openai.APIConnectionError)):
|
||||
raise RuntimeError(f"Failed to call OpenAI API: {e!s}") from e
|
||||
raise RuntimeError(f"Unexpected error during API call: {e!s}") from e
|
||||
|
||||
# Extract embedding from response
|
||||
try:
|
||||
if not response.data:
|
||||
raise ValueError("Invalid API response: no embedding data returned")
|
||||
|
||||
embedding_vector = response.data[0].embedding
|
||||
|
||||
if not isinstance(embedding_vector, list):
|
||||
raise ValueError(
|
||||
"Invalid API response: embedding is not a list of numbers"
|
||||
)
|
||||
|
||||
return embedding_vector
|
||||
|
||||
except (AttributeError, IndexError, TypeError) as e:
|
||||
raise ValueError(f"Failed to parse API response: {e!s}") from e
|
||||
@@ -0,0 +1,537 @@
|
||||
# Copyright 2025-present the zvec project
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from functools import lru_cache
|
||||
from typing import Optional
|
||||
|
||||
from ..common.constants import TEXT, DenseVectorType, SparseVectorType
|
||||
from .embedding_function import DenseEmbeddingFunction, SparseEmbeddingFunction
|
||||
from .qwen_function import QwenFunctionBase
|
||||
|
||||
|
||||
class QwenDenseEmbedding(QwenFunctionBase, DenseEmbeddingFunction[TEXT]):
|
||||
"""Dense text embedding function using Qwen (DashScope) API.
|
||||
|
||||
This class provides text-to-vector embedding capabilities using Alibaba Cloud's
|
||||
DashScope service and Qwen embedding models. It inherits from
|
||||
``DenseEmbeddingFunction`` and implements dense text embedding.
|
||||
|
||||
The implementation supports various Qwen embedding models with configurable
|
||||
dimensions and includes automatic result caching for improved performance.
|
||||
|
||||
Args:
|
||||
dimension (int): Desired output embedding dimension. Common values:
|
||||
- 512: Balanced performance and accuracy
|
||||
- 1024: Higher accuracy, larger storage
|
||||
- 1536: Maximum accuracy for supported models
|
||||
model (str, optional): DashScope embedding model identifier.
|
||||
Defaults to ``"text-embedding-v4"``. Other options include:
|
||||
- ``"text-embedding-v3"``
|
||||
- ``"text-embedding-v2"``
|
||||
- ``"text-embedding-v1"``
|
||||
api_key (Optional[str], optional): DashScope API authentication key.
|
||||
If ``None``, reads from ``DASHSCOPE_API_KEY`` environment variable.
|
||||
Obtain your key from: https://dashscope.console.aliyun.com/
|
||||
**kwargs: Additional DashScope API parameters. Supported options:
|
||||
- ``text_type`` (str): Specifies the text role in retrieval tasks.
|
||||
Options: ``"query"`` (search query) or ``"document"`` (indexed content).
|
||||
This parameter optimizes embeddings for asymmetric search scenarios.
|
||||
|
||||
Reference: https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api
|
||||
|
||||
Attributes:
|
||||
dimension (int): The embedding vector dimension.
|
||||
data_type (DataType): Always ``DataType.VECTOR_FP32`` for this implementation.
|
||||
model (str): The DashScope model name being used.
|
||||
|
||||
Raises:
|
||||
ValueError: If API key is not provided and not found in environment,
|
||||
or if API returns an error response.
|
||||
TypeError: If input to ``embed()`` is not a string.
|
||||
RuntimeError: If network error or DashScope service error occurs.
|
||||
|
||||
Note:
|
||||
- Requires Python 3.10, 3.11, or 3.12
|
||||
- Requires the ``dashscope`` package: ``pip install dashscope``
|
||||
- Embedding results are cached (LRU cache, maxsize=10) to reduce API calls
|
||||
- Network connectivity to DashScope API endpoints is required
|
||||
- API usage may incur costs based on your DashScope subscription plan
|
||||
|
||||
**Parameter Guidelines:**
|
||||
|
||||
- Use ``text_type="query"`` for search queries and ``text_type="document"``
|
||||
for indexed content to optimize asymmetric retrieval tasks.
|
||||
- For detailed API specifications and parameter usage, refer to:
|
||||
https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api
|
||||
|
||||
Examples:
|
||||
>>> # Basic usage with default model
|
||||
>>> from zvec.extension import QwenDenseEmbedding
|
||||
>>> import os
|
||||
>>> os.environ["DASHSCOPE_API_KEY"] = "your-api-key"
|
||||
>>>
|
||||
>>> emb_func = QwenDenseEmbedding(dimension=1024)
|
||||
>>> vector = emb_func.embed("Hello, world!")
|
||||
>>> len(vector)
|
||||
1024
|
||||
|
||||
>>> # Using specific model with explicit API key
|
||||
>>> emb_func = QwenDenseEmbedding(
|
||||
... dimension=512,
|
||||
... model="text-embedding-v3",
|
||||
... api_key="sk-xxxxx"
|
||||
... )
|
||||
>>> vector = emb_func("Machine learning is fascinating")
|
||||
>>> isinstance(vector, list)
|
||||
True
|
||||
|
||||
>>> # Using with custom parameters (text_type)
|
||||
>>> # For search queries - optimize for query-document matching
|
||||
>>> emb_func = QwenDenseEmbedding(
|
||||
... dimension=1024,
|
||||
... text_type="query"
|
||||
... )
|
||||
>>> query_vector = emb_func.embed("What is machine learning?")
|
||||
>>>
|
||||
>>> # For document embeddings - optimize for being matched by queries
|
||||
>>> doc_emb_func = QwenDenseEmbedding(
|
||||
... dimension=1024,
|
||||
... text_type="document"
|
||||
... )
|
||||
>>> doc_vector = doc_emb_func.embed(
|
||||
... "Machine learning is a subset of artificial intelligence..."
|
||||
... )
|
||||
|
||||
>>> # Batch processing with caching benefit
|
||||
>>> texts = ["First text", "Second text", "First text"]
|
||||
>>> vectors = [emb_func.embed(text) for text in texts]
|
||||
>>> # Third call uses cached result for "First text"
|
||||
|
||||
>>> # Error handling
|
||||
>>> try:
|
||||
... emb_func.embed("") # Empty string
|
||||
... except ValueError as e:
|
||||
... print(f"Error: {e}")
|
||||
Error: Input text cannot be empty or whitespace only
|
||||
|
||||
See Also:
|
||||
- ``DenseEmbeddingFunction``: Base class for dense embeddings
|
||||
- ``SparseEmbeddingFunction``: Base class for sparse embeddings
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dimension: int,
|
||||
model: str = "text-embedding-v4",
|
||||
api_key: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize the Qwen dense embedding function.
|
||||
|
||||
Args:
|
||||
dimension (int): Target embedding dimension.
|
||||
model (str): DashScope model name. Defaults to "text-embedding-v4".
|
||||
api_key (Optional[str]): API key or None to use environment variable.
|
||||
**kwargs: Additional DashScope API parameters. Supported options:
|
||||
- ``text_type`` (str): Text role in asymmetric retrieval.
|
||||
* ``"query"``: Optimize for search queries (short, question-like).
|
||||
* ``"document"``: Optimize for indexed documents (longer content).
|
||||
Using appropriate text_type improves retrieval accuracy by
|
||||
optimizing the embedding space for query-document matching.
|
||||
|
||||
For detailed API documentation, see:
|
||||
https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api
|
||||
|
||||
Raises:
|
||||
ValueError: If API key is not provided and not in environment.
|
||||
"""
|
||||
# Initialize base class for API connection
|
||||
QwenFunctionBase.__init__(self, model=model, api_key=api_key)
|
||||
|
||||
# Store dense-specific attributes
|
||||
self._dimension = dimension
|
||||
self._extra_params = kwargs
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
"""int: The expected dimensionality of the embedding vector."""
|
||||
return self._dimension
|
||||
|
||||
@property
|
||||
def extra_params(self) -> dict:
|
||||
"""dict: Extra parameters for model-specific customization."""
|
||||
return self._extra_params
|
||||
|
||||
def __call__(self, input: TEXT) -> DenseVectorType:
|
||||
"""Make the embedding function callable."""
|
||||
return self.embed(input)
|
||||
|
||||
@lru_cache(maxsize=10)
|
||||
def embed(self, input: TEXT) -> DenseVectorType:
|
||||
"""Generate dense embedding vector for the input text.
|
||||
|
||||
This method calls the DashScope TextEmbedding API to convert input text
|
||||
into a dense vector representation. Results are cached to improve
|
||||
performance for repeated inputs.
|
||||
|
||||
Args:
|
||||
input (TEXT): Input text string to embed. Must be non-empty after
|
||||
stripping whitespace. Maximum length depends on the model used
|
||||
(typically 2048-8192 tokens).
|
||||
|
||||
Returns:
|
||||
DenseVectorType: A list of floats representing the embedding vector.
|
||||
Length equals ``self.dimension``. Example:
|
||||
``[0.123, -0.456, 0.789, ...]``
|
||||
|
||||
Raises:
|
||||
TypeError: If ``input`` is not a string.
|
||||
ValueError: If input is empty/whitespace-only, or if the API returns
|
||||
an error or malformed response.
|
||||
RuntimeError: If network connectivity issues or DashScope service
|
||||
errors occur.
|
||||
|
||||
Examples:
|
||||
>>> emb = QwenDenseEmbedding(dimension=1024)
|
||||
>>> vector = emb.embed("Natural language processing")
|
||||
>>> len(vector)
|
||||
1024
|
||||
>>> isinstance(vector[0], float)
|
||||
True
|
||||
|
||||
>>> # Error: empty input
|
||||
>>> emb.embed(" ")
|
||||
ValueError: Input text cannot be empty or whitespace only
|
||||
|
||||
>>> # Error: non-string input
|
||||
>>> emb.embed(123)
|
||||
TypeError: Expected 'input' to be str, got int
|
||||
|
||||
Note:
|
||||
- This method is cached (maxsize=10). Identical inputs return cached results.
|
||||
- The cache is based on exact string match (case-sensitive).
|
||||
- Consider pre-processing text (lowercasing, normalization) for better caching.
|
||||
"""
|
||||
if not isinstance(input, TEXT):
|
||||
raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}")
|
||||
|
||||
input = input.strip()
|
||||
if not input:
|
||||
raise ValueError("Input text cannot be empty or whitespace only")
|
||||
|
||||
# Call API with dense output type
|
||||
output = self._call_text_embedding_api(
|
||||
input=input,
|
||||
dimension=self.dimension,
|
||||
output_type="dense",
|
||||
text_type=self.extra_params.get("text_type"),
|
||||
)
|
||||
|
||||
embeddings = output.get("embeddings")
|
||||
if not isinstance(embeddings, list):
|
||||
raise ValueError(
|
||||
"Invalid API response: 'embeddings' field is missing or not a list"
|
||||
)
|
||||
|
||||
if len(embeddings) != 1:
|
||||
raise ValueError(
|
||||
f"Expected exactly 1 embedding in response, got {len(embeddings)}"
|
||||
)
|
||||
|
||||
first_emb = embeddings[0]
|
||||
if not isinstance(first_emb, dict):
|
||||
raise ValueError("Invalid API response: embedding item is not a dictionary")
|
||||
|
||||
embedding_vector = first_emb.get("embedding")
|
||||
if not isinstance(embedding_vector, list):
|
||||
raise ValueError(
|
||||
"Invalid API response: 'embedding' field is missing or not a list"
|
||||
)
|
||||
|
||||
if len(embedding_vector) != self.dimension:
|
||||
raise ValueError(
|
||||
f"Dimension mismatch: expected {self.dimension}, "
|
||||
f"got {len(embedding_vector)}"
|
||||
)
|
||||
|
||||
return list(embedding_vector)
|
||||
|
||||
|
||||
class QwenSparseEmbedding(QwenFunctionBase, SparseEmbeddingFunction[TEXT]):
|
||||
"""Sparse text embedding function using Qwen (DashScope) API.
|
||||
|
||||
This class provides text-to-sparse-vector embedding capabilities using
|
||||
Alibaba Cloud's DashScope service and Qwen embedding models. It generates
|
||||
sparse keyword-weighted vectors suitable for lexical matching and BM25-style
|
||||
retrieval scenarios.
|
||||
|
||||
Sparse embeddings are particularly useful for:
|
||||
- Keyword-based search and exact matching
|
||||
- Hybrid retrieval (combining with dense embeddings)
|
||||
- Interpretable search results (weights show term importance)
|
||||
|
||||
Args:
|
||||
dimension (int): Desired output embedding dimension. Common values:
|
||||
- 512: Balanced performance and accuracy
|
||||
- 1024: Higher accuracy, larger storage
|
||||
- 1536: Maximum accuracy for supported models
|
||||
model (str, optional): DashScope embedding model identifier.
|
||||
Defaults to ``"text-embedding-v4"``. Other options include:
|
||||
- ``"text-embedding-v3"``
|
||||
- ``"text-embedding-v2"``
|
||||
api_key (Optional[str], optional): DashScope API authentication key.
|
||||
If ``None``, reads from ``DASHSCOPE_API_KEY`` environment variable.
|
||||
Obtain your key from: https://dashscope.console.aliyun.com/
|
||||
**kwargs: Additional DashScope API parameters. Supported options:
|
||||
- ``encoding_type`` (Literal["query", "document"]): Encoding type.
|
||||
* ``"query"``: Optimize for search queries (default).
|
||||
* ``"document"``: Optimize for indexed documents.
|
||||
This distinction is important for asymmetric retrieval tasks.
|
||||
|
||||
Attributes:
|
||||
model (str): The DashScope model name being used.
|
||||
encoding_type (str): The encoding type ("query" or "document").
|
||||
|
||||
Raises:
|
||||
ValueError: If API key is not provided and not found in environment,
|
||||
or if API returns an error response.
|
||||
TypeError: If input to ``embed()`` is not a string.
|
||||
RuntimeError: If network error or DashScope service error occurs.
|
||||
|
||||
Note:
|
||||
- Requires Python 3.10, 3.11, or 3.12
|
||||
- Requires the ``dashscope`` package: ``pip install dashscope``
|
||||
- Embedding results are cached (LRU cache, maxsize=10) to reduce API calls
|
||||
- Network connectivity to DashScope API endpoints is required
|
||||
- API usage may incur costs based on your DashScope subscription plan
|
||||
- Sparse vectors have only non-zero dimensions stored as dict
|
||||
- Output is sorted by indices (keys) in ascending order
|
||||
|
||||
**Parameter Guidelines:**
|
||||
|
||||
- Use ``encoding_type="query"`` for search queries and
|
||||
``encoding_type="document"`` for indexed content to optimize
|
||||
asymmetric retrieval tasks.
|
||||
- For detailed API specifications, refer to:
|
||||
https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api
|
||||
|
||||
Examples:
|
||||
>>> # Basic usage for query embedding
|
||||
>>> from zvec.extension import QwenSparseEmbedding
|
||||
>>> import os
|
||||
>>> os.environ["DASHSCOPE_API_KEY"] = "your-api-key"
|
||||
>>>
|
||||
>>> query_emb = QwenSparseEmbedding(dimension=1024, encoding_type="query")
|
||||
>>> query_vec = query_emb.embed("machine learning")
|
||||
>>> type(query_vec)
|
||||
<class 'dict'>
|
||||
>>> len(query_vec) # Only non-zero dimensions
|
||||
156
|
||||
|
||||
>>> # Document embedding
|
||||
>>> doc_emb = QwenSparseEmbedding(dimension=1024, encoding_type="document")
|
||||
>>> doc_vec = doc_emb.embed("Machine learning is a subset of AI")
|
||||
>>> isinstance(doc_vec, dict)
|
||||
True
|
||||
|
||||
>>> # Asymmetric retrieval example
|
||||
>>> query_vec = query_emb.embed("what causes aging fast")
|
||||
>>> doc_vec = doc_emb.embed(
|
||||
... "UV-A light causes tanning, skin aging, and cataracts..."
|
||||
... )
|
||||
>>>
|
||||
>>> # Calculate similarity (dot product for sparse vectors)
|
||||
>>> similarity = sum(
|
||||
... query_vec.get(k, 0) * doc_vec.get(k, 0)
|
||||
... for k in set(query_vec) | set(doc_vec)
|
||||
... )
|
||||
|
||||
>>> # Output is sorted by indices
|
||||
>>> list(query_vec.items())[:5] # First 5 dimensions (by index)
|
||||
[(10, 0.45), (23, 0.87), (56, 0.32), (89, 1.12), (120, 0.65)]
|
||||
|
||||
>>> # Hybrid retrieval (combining dense + sparse)
|
||||
>>> from zvec.extension import QwenDenseEmbedding
|
||||
>>> dense_emb = QwenDenseEmbedding(dimension=1024)
|
||||
>>> sparse_emb = QwenSparseEmbedding(dimension=1024)
|
||||
>>>
|
||||
>>> query = "deep learning neural networks"
|
||||
>>> dense_vec = dense_emb.embed(query) # [0.1, -0.3, 0.5, ...]
|
||||
>>> sparse_vec = sparse_emb.embed(query) # {12: 0.8, 45: 1.2, ...}
|
||||
|
||||
>>> # Error handling
|
||||
>>> try:
|
||||
... sparse_emb.embed("") # Empty string
|
||||
... except ValueError as e:
|
||||
... print(f"Error: {e}")
|
||||
Error: Input text cannot be empty or whitespace only
|
||||
|
||||
See Also:
|
||||
- ``SparseEmbeddingFunction``: Base class for sparse embeddings
|
||||
- ``QwenDenseEmbedding``: Dense embedding using Qwen API
|
||||
- ``DefaultSparseEmbedding``: Sparse embedding with SPLADE model
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dimension: int,
|
||||
model: str = "text-embedding-v4",
|
||||
api_key: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize the Qwen sparse embedding function.
|
||||
|
||||
Args:
|
||||
dimension (int): Target embedding dimension.
|
||||
model (str): DashScope model name. Defaults to "text-embedding-v4".
|
||||
api_key (Optional[str]): API key or None to use environment variable.
|
||||
**kwargs: Additional DashScope API parameters. Supported options:
|
||||
- ``encoding_type`` (Literal["query", "document"]): Encoding type.
|
||||
* ``"query"``: Optimize for search queries (default).
|
||||
* ``"document"``: Optimize for indexed documents.
|
||||
This distinction is important for asymmetric retrieval tasks.
|
||||
|
||||
Raises:
|
||||
ValueError: If API key is not provided and not in environment.
|
||||
"""
|
||||
# Initialize base class for API connection
|
||||
QwenFunctionBase.__init__(self, model=model, api_key=api_key)
|
||||
|
||||
self._dimension = dimension
|
||||
self._extra_params = kwargs
|
||||
|
||||
@property
|
||||
def extra_params(self) -> dict:
|
||||
"""dict: Extra parameters for model-specific customization."""
|
||||
return self._extra_params
|
||||
|
||||
def __call__(self, input: TEXT) -> SparseVectorType:
|
||||
"""Make the embedding function callable."""
|
||||
return self.embed(input)
|
||||
|
||||
@lru_cache(maxsize=10)
|
||||
def embed(self, input: TEXT) -> SparseVectorType:
|
||||
"""Generate sparse embedding vector for the input text.
|
||||
|
||||
This method calls the DashScope TextEmbedding API with sparse output type
|
||||
to convert input text into a sparse vector representation. The result is
|
||||
a dictionary where keys are dimension indices and values are importance
|
||||
weights (only non-zero values included).
|
||||
|
||||
The embedding is optimized based on the ``encoding_type`` specified during
|
||||
initialization: "query" for search queries or "document" for indexed content.
|
||||
|
||||
Args:
|
||||
input (TEXT): Input text string to embed. Must be non-empty after
|
||||
stripping whitespace. Maximum length depends on the model used
|
||||
(typically 2048-8192 tokens).
|
||||
|
||||
Returns:
|
||||
SparseVectorType: A dictionary mapping dimension index to weight.
|
||||
Only non-zero dimensions are included. The dictionary is sorted
|
||||
by indices (keys) in ascending order for consistent output.
|
||||
Example: ``{10: 0.5, 245: 0.8, 1023: 1.2, 5678: 0.5}``
|
||||
|
||||
Raises:
|
||||
TypeError: If ``input`` is not a string.
|
||||
ValueError: If input is empty/whitespace-only, or if the API returns
|
||||
an error or malformed response.
|
||||
RuntimeError: If network connectivity issues or DashScope service
|
||||
errors occur.
|
||||
|
||||
Examples:
|
||||
>>> emb = QwenSparseEmbedding(dimension=1024, encoding_type="query")
|
||||
>>> sparse_vec = emb.embed("machine learning")
|
||||
>>> isinstance(sparse_vec, dict)
|
||||
True
|
||||
>>>
|
||||
>>> # Verify sorted output
|
||||
>>> keys = list(sparse_vec.keys())
|
||||
>>> keys == sorted(keys)
|
||||
True
|
||||
|
||||
>>> # Error: empty input
|
||||
>>> emb.embed(" ")
|
||||
ValueError: Input text cannot be empty or whitespace only
|
||||
|
||||
>>> # Error: non-string input
|
||||
>>> emb.embed(123)
|
||||
TypeError: Expected 'input' to be str, got int
|
||||
|
||||
Note:
|
||||
- This method is cached (maxsize=10). Identical inputs return cached results.
|
||||
- The cache is based on exact string match (case-sensitive).
|
||||
- Output dictionary is always sorted by indices for consistency.
|
||||
"""
|
||||
if not isinstance(input, TEXT):
|
||||
raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}")
|
||||
|
||||
input = input.strip()
|
||||
if not input:
|
||||
raise ValueError("Input text cannot be empty or whitespace only")
|
||||
|
||||
# Call API with sparse output type
|
||||
output = self._call_text_embedding_api(
|
||||
input=input,
|
||||
dimension=self._dimension,
|
||||
output_type="sparse",
|
||||
text_type=self.extra_params.get("encoding_type", "query"),
|
||||
)
|
||||
|
||||
embeddings = output.get("embeddings")
|
||||
if not isinstance(embeddings, list):
|
||||
raise ValueError(
|
||||
"Invalid API response: 'embeddings' field is missing or not a list"
|
||||
)
|
||||
|
||||
if len(embeddings) != 1:
|
||||
raise ValueError(
|
||||
f"Expected exactly 1 embedding in response, got {len(embeddings)}"
|
||||
)
|
||||
|
||||
first_emb = embeddings[0]
|
||||
if not isinstance(first_emb, dict):
|
||||
raise ValueError("Invalid API response: embedding item is not a dictionary")
|
||||
|
||||
sparse_embedding = first_emb.get("sparse_embedding")
|
||||
if not isinstance(sparse_embedding, list):
|
||||
raise ValueError(
|
||||
"Invalid API response: 'sparse_embedding' field is missing or not a list"
|
||||
)
|
||||
|
||||
# Parse sparse embedding: convert array of {index, value, token} to dict
|
||||
sparse_dict = {}
|
||||
for item in sparse_embedding:
|
||||
if not isinstance(item, dict):
|
||||
raise ValueError(
|
||||
"Invalid API response: sparse_embedding item is not a dictionary"
|
||||
)
|
||||
|
||||
index = item.get("index")
|
||||
value = item.get("value")
|
||||
|
||||
if index is None or value is None:
|
||||
raise ValueError(
|
||||
"Invalid API response: sparse_embedding item missing 'index' or 'value'"
|
||||
)
|
||||
|
||||
# Convert to int and float, filter positive values
|
||||
idx = int(index)
|
||||
val = float(value)
|
||||
if val > 0:
|
||||
sparse_dict[idx] = val
|
||||
|
||||
# Sort by indices (keys) to ensure consistent ordering
|
||||
return dict(sorted(sparse_dict.items()))
|
||||
@@ -0,0 +1,186 @@
|
||||
# Copyright 2025-present the zvec project
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from http import HTTPStatus
|
||||
from typing import Optional
|
||||
|
||||
from ..common.constants import TEXT
|
||||
from ..tool import require_module
|
||||
|
||||
|
||||
class QwenFunctionBase:
|
||||
"""Base class for Qwen (DashScope) functions.
|
||||
|
||||
This base class provides common functionality for calling DashScope APIs
|
||||
and handling responses. It supports embeddings (dense and sparse) and
|
||||
re-ranking operations.
|
||||
|
||||
This class is not meant to be used directly. Use concrete implementations:
|
||||
- ``QwenDenseEmbedding`` for dense embeddings
|
||||
- ``QwenSparseEmbedding`` for sparse embeddings
|
||||
- ``QwenReRanker`` for semantic re-ranking
|
||||
|
||||
Args:
|
||||
model (str): DashScope model identifier.
|
||||
api_key (Optional[str]): DashScope API authentication key.
|
||||
|
||||
Note:
|
||||
- This is an internal base class for code reuse across Qwen features
|
||||
- Subclasses should inherit from appropriate Protocol/ABC
|
||||
- Provides unified API connection and response handling
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str,
|
||||
api_key: Optional[str] = None,
|
||||
):
|
||||
"""Initialize the base Qwen embedding functionality.
|
||||
|
||||
Args:
|
||||
model (str): DashScope model name.
|
||||
api_key (Optional[str]): API key or None to use environment variable.
|
||||
|
||||
Raises:
|
||||
ValueError: If API key is not provided and not in environment.
|
||||
"""
|
||||
self._model = model
|
||||
self._api_key = api_key or os.environ.get("DASHSCOPE_API_KEY")
|
||||
if not self._api_key:
|
||||
raise ValueError(
|
||||
"DashScope API key is required. Please provide 'api_key' parameter "
|
||||
"or set the 'DASHSCOPE_API_KEY' environment variable."
|
||||
)
|
||||
|
||||
@property
|
||||
def model(self) -> str:
|
||||
"""str: The DashScope embedding model name currently in use."""
|
||||
return self._model
|
||||
|
||||
def _get_connection(self):
|
||||
"""Establish connection to DashScope API.
|
||||
|
||||
Returns:
|
||||
module: The dashscope module with API key configured.
|
||||
|
||||
Raises:
|
||||
ImportError: If dashscope package is not installed.
|
||||
"""
|
||||
dashscope = require_module("dashscope")
|
||||
dashscope.api_key = self._api_key
|
||||
return dashscope
|
||||
|
||||
def _call_text_embedding_api(
|
||||
self,
|
||||
input: TEXT,
|
||||
dimension: int,
|
||||
output_type: str,
|
||||
text_type: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""Call DashScope TextEmbedding API.
|
||||
|
||||
Args:
|
||||
input (TEXT): Input text to embed.
|
||||
dimension (int): Target embedding dimension.
|
||||
output_type (str): Output type ("dense" or "sparse").
|
||||
text_type (Optional[str]): Text type ("query" or "document").
|
||||
|
||||
Returns:
|
||||
dict: API response output field.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If API call fails.
|
||||
ValueError: If API returns error response.
|
||||
"""
|
||||
try:
|
||||
# Prepare API call parameters
|
||||
call_params = {
|
||||
"model": self.model,
|
||||
"input": input,
|
||||
"dimension": dimension,
|
||||
"output_type": output_type,
|
||||
}
|
||||
|
||||
# Add optional text_type parameter if provided
|
||||
if text_type is not None:
|
||||
call_params["text_type"] = text_type
|
||||
|
||||
resp = self._get_connection().TextEmbedding.call(**call_params)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to call DashScope API: {e!s}") from e
|
||||
|
||||
if resp.status_code != HTTPStatus.OK:
|
||||
error_msg = getattr(resp, "message", "Unknown error")
|
||||
error_code = getattr(resp, "code", "N/A")
|
||||
raise ValueError(
|
||||
f"DashScope API error: [Code={error_code}, "
|
||||
f"Status={resp.status_code}] {error_msg}"
|
||||
)
|
||||
|
||||
output = getattr(resp, "output", None)
|
||||
if not isinstance(output, dict):
|
||||
raise ValueError(
|
||||
"Invalid API response: missing or malformed 'output' field"
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
def _call_rerank_api(
|
||||
self,
|
||||
query: str,
|
||||
documents: list[str],
|
||||
top_n: int,
|
||||
) -> dict:
|
||||
"""Call DashScope TextReRank API.
|
||||
|
||||
Args:
|
||||
query (str): Query text for semantic matching.
|
||||
documents (list[str]): List of document texts to re-rank.
|
||||
top_n (int): Maximum number of documents to return.
|
||||
|
||||
Returns:
|
||||
dict: API response output field containing re-ranked results.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If API call fails.
|
||||
ValueError: If API returns error response.
|
||||
"""
|
||||
try:
|
||||
resp = self._get_connection().TextReRank.call(
|
||||
model=self.model,
|
||||
query=query,
|
||||
documents=documents,
|
||||
top_n=top_n,
|
||||
return_documents=False,
|
||||
)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to call DashScope API: {e!s}") from e
|
||||
|
||||
if resp.status_code != HTTPStatus.OK:
|
||||
error_msg = getattr(resp, "message", "Unknown error")
|
||||
error_code = getattr(resp, "code", "N/A")
|
||||
raise ValueError(
|
||||
f"DashScope API error: [Code={error_code}, "
|
||||
f"Status={resp.status_code}] {error_msg}"
|
||||
)
|
||||
|
||||
output = getattr(resp, "output", None)
|
||||
if not isinstance(output, dict):
|
||||
raise ValueError(
|
||||
"Invalid API response: missing or malformed 'output' field"
|
||||
)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,177 @@
|
||||
# Copyright 2025-present the zvec project
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
from ..model.doc import Doc, DocList
|
||||
from .qwen_function import QwenFunctionBase
|
||||
from .rerank_function import RerankFunction
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..model.schema import FieldSchema, VectorSchema
|
||||
|
||||
|
||||
class QwenReRanker(QwenFunctionBase, RerankFunction):
|
||||
"""Re-ranker using Qwen (DashScope) cross-encoder API for semantic re-ranking.
|
||||
|
||||
This re-ranker leverages DashScope's TextReRank service to perform
|
||||
cross-encoder style re-ranking. It sends query and document pairs to the
|
||||
API and receives relevance scores based on deep semantic understanding.
|
||||
|
||||
The re-ranker is suitable for single-vector or multi-vector search scenarios
|
||||
where semantic relevance to a specific query is required.
|
||||
|
||||
Args:
|
||||
query (str): Query text for semantic re-ranking. **Required**.
|
||||
rerank_field (str): Document field name to use as re-ranking input text.
|
||||
**Required** (e.g., "content", "title", "body").
|
||||
model (str, optional): DashScope re-ranking model identifier.
|
||||
Defaults to ``"gte-rerank-v2"``.
|
||||
api_key (Optional[str], optional): DashScope API authentication key.
|
||||
If not provided, reads from ``DASHSCOPE_API_KEY`` environment variable.
|
||||
|
||||
Raises:
|
||||
ValueError: If ``query`` is empty/None, ``rerank_field`` is None,
|
||||
or API key is not available.
|
||||
|
||||
Note:
|
||||
- Requires ``dashscope`` Python package installed
|
||||
- Documents without valid content in ``rerank_field`` are skipped
|
||||
- API rate limits and quotas apply per DashScope subscription
|
||||
|
||||
Example:
|
||||
>>> reranker = QwenReRanker(
|
||||
... query="machine learning algorithms",
|
||||
... rerank_field="content",
|
||||
... model="gte-rerank-v2",
|
||||
... api_key="your-api-key"
|
||||
... )
|
||||
>>> # Use in collection.query(reranker=reranker)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
query: Optional[str] = None,
|
||||
rerank_field: Optional[str] = None,
|
||||
model: str = "gte-rerank-v2",
|
||||
api_key: Optional[str] = None,
|
||||
):
|
||||
"""Initialize QwenReRanker with query and configuration.
|
||||
|
||||
Args:
|
||||
query (Optional[str]): Query text for semantic matching. Required.
|
||||
rerank_field (Optional[str]): Document field for re-ranking input.
|
||||
model (str): DashScope model name.
|
||||
api_key (Optional[str]): API key or None to use environment variable.
|
||||
|
||||
Raises:
|
||||
ValueError: If query is empty or API key is unavailable.
|
||||
"""
|
||||
QwenFunctionBase.__init__(self, model=model, api_key=api_key)
|
||||
self._rerank_field = rerank_field
|
||||
|
||||
if not query:
|
||||
raise ValueError("Query is required for QwenReRanker")
|
||||
self._query = query
|
||||
|
||||
@property
|
||||
def rerank_field(self) -> Optional[str]:
|
||||
"""Optional[str]: Field name used as re-ranking input."""
|
||||
return self._rerank_field
|
||||
|
||||
@property
|
||||
def query(self) -> str:
|
||||
"""str: Query text used for semantic re-ranking."""
|
||||
return self._query
|
||||
|
||||
def rerank(
|
||||
self,
|
||||
query_results: list[list[Doc]],
|
||||
topn: int = 10,
|
||||
*,
|
||||
fields: list[FieldSchema | VectorSchema] | None = None, # noqa: ARG002
|
||||
) -> DocList:
|
||||
"""Re-rank documents using Qwen's TextReRank API.
|
||||
|
||||
Sends document texts to DashScope TextReRank service along with the query.
|
||||
Returns documents sorted by relevance scores from the cross-encoder model.
|
||||
|
||||
Args:
|
||||
query_results (list[list[Doc]]): Per-sub-query lists of retrieved
|
||||
documents. Documents from all lists are deduplicated and
|
||||
re-ranked together.
|
||||
topn (int): Maximum number of documents to return.
|
||||
fields: Unused; present for interface compatibility.
|
||||
|
||||
Returns:
|
||||
list[Doc]: Re-ranked documents (up to ``topn``) with updated ``score``
|
||||
fields containing relevance scores from the API.
|
||||
|
||||
Raises:
|
||||
ValueError: If no valid documents are found or API call fails.
|
||||
|
||||
Note:
|
||||
- Duplicate documents (same ID) across lists are processed once
|
||||
- Documents with empty/missing ``rerank_field`` content are skipped
|
||||
- Returned scores are relevance scores from the cross-encoder model
|
||||
"""
|
||||
if not query_results:
|
||||
return []
|
||||
|
||||
# Accept both dict (legacy) and list formats
|
||||
if isinstance(query_results, dict):
|
||||
query_results = list(query_results.values())
|
||||
|
||||
# Collect and deduplicate documents
|
||||
id_to_doc: dict[str, Doc] = {}
|
||||
doc_ids: list[str] = []
|
||||
contents: list[str] = []
|
||||
|
||||
for query_result in query_results:
|
||||
for doc in query_result:
|
||||
doc_id = doc.id
|
||||
if doc_id in id_to_doc:
|
||||
continue
|
||||
|
||||
# Extract text content from specified field
|
||||
field_value = doc.field(self.rerank_field)
|
||||
rank_content = str(field_value).strip() if field_value else ""
|
||||
if not rank_content:
|
||||
continue
|
||||
|
||||
id_to_doc[doc_id] = doc
|
||||
doc_ids.append(doc_id)
|
||||
contents.append(rank_content)
|
||||
|
||||
if not contents:
|
||||
raise ValueError("No documents to rerank")
|
||||
|
||||
# Call DashScope TextReRank API
|
||||
output = self._call_rerank_api(
|
||||
query=self.query,
|
||||
documents=contents,
|
||||
top_n=topn,
|
||||
)
|
||||
|
||||
# Build result list with updated scores
|
||||
results: DocList = []
|
||||
for item in output["results"]:
|
||||
idx = item["index"]
|
||||
doc_id = doc_ids[idx]
|
||||
doc = id_to_doc[doc_id]
|
||||
new_doc = doc._replace(score=item["relevance_score"])
|
||||
results.append(new_doc)
|
||||
|
||||
return results
|
||||
@@ -0,0 +1,56 @@
|
||||
# Copyright 2025-present the zvec project
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ..model.doc import Doc, DocList
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..model.schema import FieldSchema, VectorSchema
|
||||
|
||||
|
||||
class RerankFunction(ABC):
|
||||
"""Abstract base class for reranker parameter containers.
|
||||
|
||||
Subclasses define rerank parameters and implement _to_cpp_params()
|
||||
for conversion to C++ parameter structs (used by collection fast path).
|
||||
Each subclass also provides a standalone rerank() implementation.
|
||||
"""
|
||||
|
||||
def _to_cpp_params(self):
|
||||
"""Return C++ reranker params. Override in subclasses that use C++ path."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def rerank(
|
||||
self,
|
||||
query_results: list[list[Doc]],
|
||||
topn: int = 10,
|
||||
*,
|
||||
fields: list[FieldSchema | VectorSchema] | None = None,
|
||||
) -> DocList:
|
||||
"""Execute rerank on sub-query results.
|
||||
|
||||
Args:
|
||||
query_results: List of per-sub-query document lists.
|
||||
topn: Maximum number of results to return.
|
||||
fields: Per-sub-query Python FieldSchema/VectorSchema objects
|
||||
(required for WeightedReRanker score normalization).
|
||||
|
||||
Returns:
|
||||
Re-ranked document list.
|
||||
"""
|
||||
...
|
||||
@@ -0,0 +1,839 @@
|
||||
# Copyright 2025-present the zvec project
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import ClassVar, Literal, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ..common.constants import TEXT, DenseVectorType, SparseVectorType
|
||||
from .embedding_function import DenseEmbeddingFunction, SparseEmbeddingFunction
|
||||
from .sentence_transformer_function import SentenceTransformerFunctionBase
|
||||
|
||||
|
||||
class DefaultLocalDenseEmbedding(
|
||||
SentenceTransformerFunctionBase, DenseEmbeddingFunction[TEXT]
|
||||
):
|
||||
"""Default local dense embedding using all-MiniLM-L6-v2 model.
|
||||
|
||||
This is the default implementation for dense text embedding that uses the
|
||||
``all-MiniLM-L6-v2`` model from Hugging Face by default. This model provides
|
||||
a good balance between speed and quality for general-purpose text embedding.
|
||||
|
||||
The class provides text-to-vector dense embedding capabilities using the
|
||||
sentence-transformers library. It supports models from Hugging Face Hub and
|
||||
ModelScope, runs locally without API calls, and supports CPU/GPU acceleration.
|
||||
|
||||
The model produces 384-dimensional embeddings and is optimized for semantic
|
||||
similarity tasks. It runs locally without requiring API keys.
|
||||
|
||||
Args:
|
||||
model_source (Literal["huggingface", "modelscope"], optional): Model source.
|
||||
- ``"huggingface"``: Use Hugging Face Hub (default, for international users)
|
||||
- ``"modelscope"``: Use ModelScope (recommended for users in China)
|
||||
Defaults to ``"huggingface"``.
|
||||
device (Optional[str], optional): Device to run the model on.
|
||||
Options: ``"cpu"``, ``"cuda"``, ``"mps"`` (for Apple Silicon), or ``None``
|
||||
for automatic detection. Defaults to ``None``.
|
||||
normalize_embeddings (bool, optional): Whether to normalize embeddings to
|
||||
unit length (L2 normalization). Useful for cosine similarity.
|
||||
Defaults to ``True``.
|
||||
batch_size (int, optional): Batch size for encoding. Defaults to ``32``.
|
||||
**kwargs: Additional parameters for future extension.
|
||||
|
||||
Attributes:
|
||||
dimension (int): Always 384 for both models.
|
||||
model_name (str): "all-MiniLM-L6-v2" (HF) or "iic/nlp_gte_sentence-embedding_chinese-small" (MS).
|
||||
model_source (str): The model source being used.
|
||||
device (str): The device the model is running on.
|
||||
|
||||
Raises:
|
||||
ValueError: If the model cannot be loaded or input is invalid.
|
||||
TypeError: If input to ``embed()`` is not a string.
|
||||
RuntimeError: If model inference fails.
|
||||
|
||||
Note:
|
||||
- Requires Python 3.10, 3.11, or 3.12
|
||||
- Requires the ``sentence-transformers`` package:
|
||||
``pip install sentence-transformers``
|
||||
- For ModelScope, also requires: ``pip install modelscope``
|
||||
- First run downloads the model (~50-80MB) from chosen source
|
||||
- Hugging Face cache: ``~/.cache/torch/sentence_transformers/``
|
||||
- ModelScope cache: ``~/.cache/modelscope/hub/``
|
||||
- No API keys or network required after initial download
|
||||
- Inference speed: ~1000 sentences/sec on CPU, ~10000 on GPU
|
||||
|
||||
**For users in China:**
|
||||
|
||||
If you encounter Hugging Face access issues, use ModelScope instead:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Recommended for users in China
|
||||
emb = DefaultLocalDenseEmbedding(model_source="modelscope")
|
||||
|
||||
Alternatively, use Hugging Face mirror:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
export HF_ENDPOINT=https://hf-mirror.com
|
||||
# Then use default Hugging Face mode
|
||||
|
||||
Examples:
|
||||
>>> # Basic usage with Hugging Face (default)
|
||||
>>> from zvec.extension import DefaultLocalDenseEmbedding
|
||||
>>>
|
||||
>>> emb_func = DefaultLocalDenseEmbedding()
|
||||
>>> vector = emb_func.embed("Hello, world!")
|
||||
>>> len(vector)
|
||||
384
|
||||
>>> isinstance(vector, list)
|
||||
True
|
||||
|
||||
>>> # Recommended for users in China (uses ModelScope)
|
||||
>>> emb_func = DefaultLocalDenseEmbedding(model_source="modelscope")
|
||||
>>> vector = emb_func.embed("你好，世界！") # Works well with Chinese text
|
||||
>>> len(vector)
|
||||
384
|
||||
|
||||
>>> # Alternative for China users: Use Hugging Face mirror
|
||||
>>> import os
|
||||
>>> os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
|
||||
>>> emb_func = DefaultLocalDenseEmbedding() # Uses HF mirror
|
||||
>>> vector = emb_func.embed("Hello, world!")
|
||||
|
||||
>>> # Using GPU for faster inference
|
||||
>>> emb_func = DefaultLocalDenseEmbedding(device="cuda")
|
||||
>>> vector = emb_func("Machine learning is fascinating")
|
||||
>>> # Normalized vector has unit length
|
||||
>>> import numpy as np
|
||||
>>> np.linalg.norm(vector)
|
||||
1.0
|
||||
|
||||
>>> # Batch processing
|
||||
>>> texts = ["First text", "Second text", "Third text"]
|
||||
>>> vectors = [emb_func.embed(text) for text in texts]
|
||||
>>> len(vectors)
|
||||
3
|
||||
>>> all(len(v) == 384 for v in vectors)
|
||||
True
|
||||
|
||||
>>> # Semantic similarity
|
||||
>>> v1 = emb_func.embed("The cat sits on the mat")
|
||||
>>> v2 = emb_func.embed("A feline rests on a rug")
|
||||
>>> v3 = emb_func.embed("Python programming")
|
||||
>>> similarity_high = np.dot(v1, v2) # Similar sentences
|
||||
>>> similarity_low = np.dot(v1, v3) # Different topics
|
||||
>>> similarity_high > similarity_low
|
||||
True
|
||||
|
||||
>>> # Error handling
|
||||
>>> try:
|
||||
... emb_func.embed("") # Empty string
|
||||
... except ValueError as e:
|
||||
... print(f"Error: {e}")
|
||||
Error: Input text cannot be empty or whitespace only
|
||||
|
||||
See Also:
|
||||
- ``DenseEmbeddingFunction``: Base class for dense embeddings
|
||||
- ``DefaultLocalSparseEmbedding``: Sparse embedding with SPLADE
|
||||
- ``QwenDenseEmbedding``: Alternative using Qwen API
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_source: Literal["huggingface", "modelscope"] = "huggingface",
|
||||
device: Optional[str] = None,
|
||||
normalize_embeddings: bool = True,
|
||||
batch_size: int = 32,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize with all-MiniLM-L6-v2 model.
|
||||
|
||||
Args:
|
||||
model_source (Literal["huggingface", "modelscope"]): Model source.
|
||||
Defaults to "huggingface".
|
||||
device (Optional[str]): Target device ("cpu", "cuda", "mps", or None).
|
||||
Defaults to None (automatic detection).
|
||||
normalize_embeddings (bool): Whether to L2-normalize output vectors.
|
||||
Defaults to True.
|
||||
batch_size (int): Batch size for encoding. Defaults to 32.
|
||||
**kwargs: Additional parameters for future extension.
|
||||
|
||||
Raises:
|
||||
ImportError: If sentence-transformers or modelscope is not installed.
|
||||
ValueError: If model cannot be loaded.
|
||||
"""
|
||||
# Use different models based on source
|
||||
if model_source == "modelscope":
|
||||
# Use Chinese-optimized model for ModelScope (better for Chinese text)
|
||||
model_name = "iic/nlp_gte_sentence-embedding_chinese-small"
|
||||
else:
|
||||
model_name = "all-MiniLM-L6-v2"
|
||||
|
||||
# Initialize base class for model loading
|
||||
SentenceTransformerFunctionBase.__init__(
|
||||
self, model_name=model_name, model_source=model_source, device=device
|
||||
)
|
||||
|
||||
self._normalize_embeddings = normalize_embeddings
|
||||
self._batch_size = batch_size
|
||||
|
||||
# Load model and get dimension
|
||||
model = self._get_model()
|
||||
self._dimension = model.get_sentence_embedding_dimension()
|
||||
|
||||
# Store extra parameters
|
||||
self._extra_params = kwargs
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
"""int: The expected dimensionality of the embedding vector."""
|
||||
return self._dimension
|
||||
|
||||
@property
|
||||
def extra_params(self) -> dict:
|
||||
"""dict: Extra parameters for model-specific customization."""
|
||||
return self._extra_params
|
||||
|
||||
def __call__(self, input: str) -> DenseVectorType:
|
||||
"""Make the embedding function callable."""
|
||||
return self.embed(input)
|
||||
|
||||
def embed(self, input: str) -> DenseVectorType:
|
||||
"""Generate dense embedding vector for the input text.
|
||||
|
||||
This method uses the Sentence Transformer model to convert input text
|
||||
into a dense vector representation. The model runs locally without
|
||||
requiring API calls.
|
||||
|
||||
Args:
|
||||
input (str): Input text string to embed. Must be non-empty after
|
||||
stripping whitespace. Maximum length depends on the model used
|
||||
(typically 128-512 tokens for most models).
|
||||
|
||||
Returns:
|
||||
DenseVectorType: A list of floats representing the embedding vector.
|
||||
Length equals ``self.dimension``. If ``normalize_embeddings=True``,
|
||||
the vector has unit length. Example:
|
||||
``[0.123, -0.456, 0.789, ...]``
|
||||
|
||||
Raises:
|
||||
TypeError: If ``input`` is not a string.
|
||||
ValueError: If input is empty or whitespace-only.
|
||||
RuntimeError: If model inference fails.
|
||||
|
||||
Examples:
|
||||
>>> emb = DefaultLocalDenseEmbedding()
|
||||
>>> vector = emb.embed("Natural language processing")
|
||||
>>> len(vector)
|
||||
384
|
||||
>>> isinstance(vector[0], float)
|
||||
True
|
||||
|
||||
>>> # Normalized vectors have unit length
|
||||
>>> import numpy as np
|
||||
>>> emb = DefaultLocalDenseEmbedding(normalize_embeddings=True)
|
||||
>>> vector = emb.embed("Test sentence")
|
||||
>>> np.linalg.norm(vector)
|
||||
1.0
|
||||
|
||||
>>> # Error: empty input
|
||||
>>> emb.embed(" ")
|
||||
ValueError: Input text cannot be empty or whitespace only
|
||||
|
||||
>>> # Error: non-string input
|
||||
>>> emb.embed(123)
|
||||
TypeError: Expected 'input' to be str, got int
|
||||
|
||||
>>> # Semantic similarity example
|
||||
>>> v1 = emb.embed("The cat sits on the mat")
|
||||
>>> v2 = emb.embed("A feline rests on a rug")
|
||||
>>> similarity = np.dot(v1, v2) # High similarity due to semantic meaning
|
||||
>>> similarity > 0.7
|
||||
True
|
||||
|
||||
Note:
|
||||
- First call may be slower due to model loading
|
||||
- Subsequent calls are much faster as the model stays in memory
|
||||
- For batch processing, consider encoding multiple texts together
|
||||
(though this method handles single texts only)
|
||||
- GPU acceleration provides 5-10x speedup over CPU
|
||||
"""
|
||||
if not isinstance(input, str):
|
||||
raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}")
|
||||
|
||||
input = input.strip()
|
||||
if not input:
|
||||
raise ValueError("Input text cannot be empty or whitespace only")
|
||||
|
||||
try:
|
||||
model = self._get_model()
|
||||
embedding = model.encode(
|
||||
input,
|
||||
convert_to_numpy=True,
|
||||
normalize_embeddings=self._normalize_embeddings,
|
||||
batch_size=self._batch_size,
|
||||
)
|
||||
|
||||
# Convert numpy array to list
|
||||
if isinstance(embedding, np.ndarray):
|
||||
embedding_list = embedding.tolist()
|
||||
else:
|
||||
embedding_list = list(embedding)
|
||||
|
||||
# Validate dimension
|
||||
if len(embedding_list) != self.dimension:
|
||||
raise ValueError(
|
||||
f"Dimension mismatch: expected {self.dimension}, "
|
||||
f"got {len(embedding_list)}"
|
||||
)
|
||||
|
||||
return embedding_list
|
||||
|
||||
except Exception as e:
|
||||
if isinstance(e, (TypeError, ValueError)):
|
||||
raise
|
||||
raise RuntimeError(f"Failed to generate embedding: {e!s}") from e
|
||||
|
||||
|
||||
class DefaultLocalSparseEmbedding(
|
||||
SentenceTransformerFunctionBase, SparseEmbeddingFunction[TEXT]
|
||||
):
|
||||
"""Default local sparse embedding using SPLADE model.
|
||||
|
||||
This class provides sparse vector embedding using the SPLADE (SParse Lexical
|
||||
AnD Expansion) model. SPLADE generates sparse, interpretable representations
|
||||
where each dimension corresponds to a vocabulary term with learned importance
|
||||
weights. It's ideal for lexical matching, BM25-style retrieval, and hybrid
|
||||
search scenarios.
|
||||
|
||||
The default model is ``naver/splade-cocondenser-ensembledistil``, which is
|
||||
publicly available without authentication. It produces sparse vectors with
|
||||
thousands of dimensions but only hundreds of non-zero values, making them
|
||||
efficient for storage and retrieval while maintaining strong lexical matching.
|
||||
|
||||
**Model Caching:**
|
||||
|
||||
This class uses class-level caching to share the SPLADE model across all instances
|
||||
with the same configuration (model_source, device). This significantly reduces
|
||||
memory usage when creating multiple instances for different encoding types
|
||||
(query vs document).
|
||||
|
||||
**Cache Management:**
|
||||
|
||||
The class provides methods to manage the model cache:
|
||||
|
||||
- ``clear_cache()``: Clear all cached models to free memory
|
||||
- ``get_cache_info()``: Get information about cached models
|
||||
- ``remove_from_cache(model_source, device)``: Remove a specific model from cache
|
||||
|
||||
.. note::
|
||||
**Why not use splade-v3?**
|
||||
|
||||
The newer ``naver/splade-v3`` model is gated (requires access approval).
|
||||
We use ``naver/splade-cocondenser-ensembledistil`` instead.
|
||||
|
||||
**To use splade-v3 (if you have access):**
|
||||
|
||||
1. Request access at https://huggingface.co/naver/splade-v3
|
||||
2. Get your Hugging Face token from https://huggingface.co/settings/tokens
|
||||
3. Set environment variable:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
export HF_TOKEN="your_huggingface_token"
|
||||
|
||||
4. Or login programmatically:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from huggingface_hub import login
|
||||
login(token="your_huggingface_token")
|
||||
|
||||
5. To use a custom SPLADE model, you can subclass this class and override
|
||||
the model_name in ``__init__``, or create your own implementation
|
||||
inheriting from ``SentenceTransformerFunctionBase`` and
|
||||
``SparseEmbeddingFunction``.
|
||||
|
||||
Args:
|
||||
model_source (Literal["huggingface", "modelscope"], optional): Model source.
|
||||
Defaults to ``"huggingface"``. ModelScope support may vary for SPLADE models.
|
||||
device (Optional[str], optional): Device to run the model on.
|
||||
Options: ``"cpu"``, ``"cuda"``, ``"mps"`` (for Apple Silicon), or ``None``
|
||||
for automatic detection. Defaults to ``None``.
|
||||
encoding_type (Literal["query", "document"], optional): Encoding type.
|
||||
- ``"query"``: Optimize for search queries (default)
|
||||
- ``"document"``: Optimize for indexed documents
|
||||
**kwargs: Additional parameters (currently unused, for future extension).
|
||||
|
||||
Attributes:
|
||||
model_name (str): Model identifier.
|
||||
model_source (str): The model source being used.
|
||||
device (str): The device the model is running on.
|
||||
|
||||
Raises:
|
||||
ValueError: If the model cannot be loaded or input is invalid.
|
||||
TypeError: If input to ``embed()`` is not a string.
|
||||
RuntimeError: If model inference fails.
|
||||
|
||||
Note:
|
||||
- Requires Python 3.10, 3.11, or 3.12
|
||||
- Requires the ``sentence-transformers`` package:
|
||||
``pip install sentence-transformers``
|
||||
- First run downloads the model (~100MB) from Hugging Face
|
||||
- Cache location: ``~/.cache/torch/sentence_transformers/``
|
||||
- No API keys or authentication required
|
||||
- Sparse vectors have ~30k dimensions but only ~100-200 non-zero values
|
||||
- Best combined with dense embeddings for hybrid retrieval
|
||||
|
||||
**SPLADE vs Dense Embeddings:**
|
||||
|
||||
- **Dense**: Continuous semantic vectors, good for semantic similarity
|
||||
- **Sparse**: Lexical keyword-based, interpretable, good for exact matching
|
||||
- **Hybrid**: Combine both for best retrieval performance
|
||||
|
||||
Examples:
|
||||
>>> # Memory-efficient: both instances share the same model (~200MB)
|
||||
>>> from zvec.extension import DefaultLocalSparseEmbedding
|
||||
>>>
|
||||
>>> # Query embedding
|
||||
>>> query_emb = DefaultLocalSparseEmbedding(encoding_type="query")
|
||||
>>> query_vec = query_emb.embed("machine learning algorithms")
|
||||
>>> type(query_vec)
|
||||
<class 'dict'>
|
||||
>>> len(query_vec) # Only non-zero dimensions
|
||||
156
|
||||
|
||||
>>> # Document embedding (shares model with query_emb)
|
||||
>>> doc_emb = DefaultLocalSparseEmbedding(encoding_type="document")
|
||||
>>> doc_vec = doc_emb.embed("Machine learning is a subset of AI")
|
||||
>>> # Total memory: ~200MB (not 400MB) thanks to model caching
|
||||
|
||||
>>> # Asymmetric retrieval example
|
||||
>>> query_vec = query_emb.embed("what causes aging fast")
|
||||
>>> doc_vec = doc_emb.embed(
|
||||
... "UV-A light causes tanning, skin aging, and cataracts..."
|
||||
... )
|
||||
>>>
|
||||
>>> # Calculate similarity (dot product for sparse vectors)
|
||||
>>> similarity = sum(
|
||||
... query_vec.get(k, 0) * doc_vec.get(k, 0)
|
||||
... for k in set(query_vec) | set(doc_vec)
|
||||
... )
|
||||
|
||||
>>> # Batch processing
|
||||
>>> queries = ["query 1", "query 2", "query 3"]
|
||||
>>> query_vecs = [query_emb.embed(q) for q in queries]
|
||||
>>>
|
||||
>>> documents = ["doc 1", "doc 2", "doc 3"]
|
||||
>>> doc_vecs = [doc_emb.embed(d) for d in documents]
|
||||
|
||||
>>> # Inspecting sparse dimensions (output is sorted by indices)
|
||||
>>> query_vec = query_emb.embed("machine learning")
|
||||
>>> list(query_vec.items())[:5] # First 5 dimensions (by index)
|
||||
[(10, 0.45), (23, 0.87), (56, 0.32), (89, 1.12), (120, 0.65)]
|
||||
>>>
|
||||
>>> # Sort by weight to find most important terms
|
||||
>>> sorted_by_weight = sorted(query_vec.items(), key=lambda x: x[1], reverse=True)
|
||||
>>> top_5 = sorted_by_weight[:5] # Top 5 most important terms
|
||||
>>> top_5
|
||||
[(1023, 1.45), (245, 1.23), (8901, 0.98), (5678, 0.87), (12034, 0.76)]
|
||||
|
||||
>>> # Using GPU for faster inference
|
||||
>>> sparse_emb = DefaultLocalSparseEmbedding(device="cuda")
|
||||
>>> vector = sparse_emb.embed("natural language processing")
|
||||
|
||||
>>> # Hybrid retrieval example (combining dense + sparse)
|
||||
>>> from zvec.extension import DefaultDenseEmbedding
|
||||
>>> dense_emb = DefaultDenseEmbedding()
|
||||
>>> sparse_emb = DefaultLocalSparseEmbedding()
|
||||
>>>
|
||||
>>> query = "deep learning neural networks"
|
||||
>>> dense_vec = dense_emb.embed(query) # [0.1, -0.3, 0.5, ...]
|
||||
>>> sparse_vec = sparse_emb.embed(query) # {12: 0.8, 45: 1.2, ...}
|
||||
|
||||
>>> # Error handling
|
||||
>>> try:
|
||||
... sparse_emb.embed("") # Empty string
|
||||
... except ValueError as e:
|
||||
... print(f"Error: {e}")
|
||||
Error: Input text cannot be empty or whitespace only
|
||||
|
||||
>>> # Cache management
|
||||
>>> # Check cache status
|
||||
>>> info = DefaultLocalSparseEmbedding.get_cache_info()
|
||||
>>> print(f"Cached models: {info['cached_models']}")
|
||||
Cached models: 1
|
||||
>>>
|
||||
>>> # Clear cache to free memory
|
||||
>>> DefaultLocalSparseEmbedding.clear_cache()
|
||||
>>> info = DefaultLocalSparseEmbedding.get_cache_info()
|
||||
>>> print(f"Cached models: {info['cached_models']}")
|
||||
Cached models: 0
|
||||
>>>
|
||||
>>> # Remove specific model from cache
|
||||
>>> query_emb = DefaultLocalSparseEmbedding() # Creates CPU model
|
||||
>>> cuda_emb = DefaultLocalSparseEmbedding(device="cuda") # Creates CUDA model
|
||||
>>> info = DefaultLocalSparseEmbedding.get_cache_info()
|
||||
>>> print(f"Cached models: {info['cached_models']}")
|
||||
Cached models: 2
|
||||
>>>
|
||||
>>> # Remove only CPU model
|
||||
>>> removed = DefaultLocalSparseEmbedding.remove_from_cache(device=None)
|
||||
>>> print(f"Removed: {removed}")
|
||||
True
|
||||
>>> info = DefaultLocalSparseEmbedding.get_cache_info()
|
||||
>>> print(f"Cached models: {info['cached_models']}")
|
||||
Cached models: 1
|
||||
|
||||
See Also:
|
||||
- ``SparseEmbeddingFunction``: Base class for sparse embeddings
|
||||
- ``DefaultDenseEmbedding``: Dense embedding with all-MiniLM-L6-v2
|
||||
- ``QwenDenseEmbedding``: Alternative using Qwen API
|
||||
|
||||
References:
|
||||
- SPLADE Paper: https://arxiv.org/abs/2109.10086
|
||||
- Model: https://huggingface.co/naver/splade-cocondenser-ensembledistil
|
||||
"""
|
||||
|
||||
# Class-level model cache: {(model_name, model_source, device): model}
|
||||
# Shared across all DefaultLocalSparseEmbedding instances to save memory
|
||||
_model_cache: ClassVar[dict] = {}
|
||||
|
||||
@classmethod
|
||||
def clear_cache(cls) -> None:
|
||||
"""Clear all cached SPLADE models from memory.
|
||||
|
||||
This is useful for:
|
||||
- Freeing memory when models are no longer needed
|
||||
- Forcing a fresh model reload
|
||||
- Testing and debugging
|
||||
Examples:
|
||||
>>> # Clear cache to free memory
|
||||
>>> DefaultLocalSparseEmbedding.clear_cache()
|
||||
|
||||
>>> # Or in tests to ensure fresh model loading
|
||||
>>> def test_something():
|
||||
... DefaultLocalSparseEmbedding.clear_cache()
|
||||
... emb = DefaultLocalSparseEmbedding()
|
||||
... # Test with fresh model
|
||||
"""
|
||||
cls._model_cache.clear()
|
||||
|
||||
@classmethod
|
||||
def get_cache_info(cls) -> dict:
|
||||
"""Get information about currently cached models.
|
||||
|
||||
Returns:
|
||||
dict: Dictionary with cache statistics:
|
||||
- cached_models (int): Number of cached model instances
|
||||
- cache_keys (list): List of cache keys (model_name, model_source, device)
|
||||
|
||||
Examples:
|
||||
>>> info = DefaultLocalSparseEmbedding.get_cache_info()
|
||||
>>> print(f"Cached models: {info['cached_models']}")
|
||||
Cached models: 2
|
||||
>>> print(f"Cache keys: {info['cache_keys']}")
|
||||
Cache keys: [('naver/splade-cocondenser-ensembledistil', 'huggingface', None),
|
||||
('naver/splade-cocondenser-ensembledistil', 'huggingface', 'cuda')]
|
||||
"""
|
||||
return {
|
||||
"cached_models": len(cls._model_cache),
|
||||
"cache_keys": list(cls._model_cache.keys()),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def remove_from_cache(
|
||||
cls, model_source: str = "huggingface", device: Optional[str] = None
|
||||
) -> bool:
|
||||
"""Remove a specific model from cache.
|
||||
|
||||
Args:
|
||||
model_source (str): Model source ("huggingface" or "modelscope").
|
||||
Defaults to "huggingface".
|
||||
device (Optional[str]): Device identifier. Defaults to None.
|
||||
|
||||
Returns:
|
||||
bool: True if model was found and removed, False otherwise.
|
||||
|
||||
Examples:
|
||||
>>> # Remove CPU model from cache
|
||||
>>> removed = DefaultLocalSparseEmbedding.remove_from_cache()
|
||||
>>> print(f"Removed: {removed}")
|
||||
True
|
||||
|
||||
>>> # Remove CUDA model from cache
|
||||
>>> removed = DefaultLocalSparseEmbedding.remove_from_cache(device="cuda")
|
||||
>>> print(f"Removed: {removed}")
|
||||
True
|
||||
"""
|
||||
model_name = "naver/splade-cocondenser-ensembledistil"
|
||||
cache_key = (model_name, model_source, device)
|
||||
|
||||
if cache_key in cls._model_cache:
|
||||
del cls._model_cache[cache_key]
|
||||
return True
|
||||
return False
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_source: Literal["huggingface", "modelscope"] = "huggingface",
|
||||
device: Optional[str] = None,
|
||||
encoding_type: Literal["query", "document"] = "query",
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize with SPLADE model.
|
||||
|
||||
Args:
|
||||
model_source (Literal["huggingface", "modelscope"]): Model source.
|
||||
Defaults to "huggingface".
|
||||
device (Optional[str]): Target device ("cpu", "cuda", "mps", or None).
|
||||
Defaults to None (automatic detection).
|
||||
encoding_type (Literal["query", "document"]): Encoding type for embeddings.
|
||||
- "query": Optimize for search queries (default)
|
||||
- "document": Optimize for indexed documents
|
||||
This distinction is important for asymmetric retrieval tasks.
|
||||
**kwargs: Additional parameters (reserved for future use).
|
||||
|
||||
Raises:
|
||||
ImportError: If sentence-transformers is not installed.
|
||||
ValueError: If model cannot be loaded.
|
||||
|
||||
Note:
|
||||
Multiple instances with the same (model_source, device) configuration
|
||||
will share the same underlying model to save memory. Different
|
||||
instances can use different encoding_type settings while sharing
|
||||
the model.
|
||||
|
||||
**Model Selection:**
|
||||
|
||||
Uses ``naver/splade-cocondenser-ensembledistil`` instead of the newer
|
||||
``naver/splade-v3`` because splade-v3 is a gated model requiring
|
||||
Hugging Face authentication. The cocondenser-ensembledistil variant:
|
||||
|
||||
- Does not require authentication or API tokens
|
||||
- Is immediately available for all users
|
||||
- Provides comparable retrieval performance (~2% difference)
|
||||
- Avoids "Access to model is restricted" errors
|
||||
|
||||
If you need splade-v3 and have obtained access, you can subclass
|
||||
this class and override the model_name parameter.
|
||||
|
||||
Examples:
|
||||
>>> # Both instances share the same model (saves memory)
|
||||
>>> query_emb = DefaultLocalSparseEmbedding(encoding_type="query")
|
||||
>>> doc_emb = DefaultLocalSparseEmbedding(encoding_type="document")
|
||||
>>> # Only one model is loaded in memory
|
||||
"""
|
||||
# Use publicly available SPLADE model (no gated access required)
|
||||
# Note: naver/splade-v3 requires authentication, so we use the
|
||||
# cocondenser-ensembledistil variant which is publicly accessible
|
||||
model_name = "naver/splade-cocondenser-ensembledistil"
|
||||
|
||||
# Initialize base class for model loading
|
||||
SentenceTransformerFunctionBase.__init__(
|
||||
self, model_name=model_name, model_source=model_source, device=device
|
||||
)
|
||||
|
||||
self._encoding_type = encoding_type
|
||||
self._extra_params = kwargs
|
||||
|
||||
# Create cache key for this model configuration
|
||||
self._cache_key = (model_name, model_source, device)
|
||||
|
||||
# Load model to ensure it's available (will use cache if exists)
|
||||
self._get_model()
|
||||
|
||||
@property
|
||||
def extra_params(self) -> dict:
|
||||
"""dict: Extra parameters for model-specific customization."""
|
||||
return self._extra_params
|
||||
|
||||
def __call__(self, input: str) -> SparseVectorType:
|
||||
"""Make the embedding function callable."""
|
||||
return self.embed(input)
|
||||
|
||||
def embed(self, input: str) -> SparseVectorType:
|
||||
"""Generate sparse embedding vector for the input text.
|
||||
|
||||
This method uses the SPLADE model to convert input text into a sparse
|
||||
vector representation. The result is a dictionary where keys are dimension
|
||||
indices and values are importance weights (only non-zero values included).
|
||||
|
||||
The embedding is optimized based on the ``encoding_type`` specified during
|
||||
initialization: "query" for search queries or "document" for indexed content.
|
||||
|
||||
Args:
|
||||
input (str): Input text string to embed. Must be non-empty after
|
||||
stripping whitespace.
|
||||
|
||||
Returns:
|
||||
SparseVectorType: A dictionary mapping dimension index to weight.
|
||||
Only non-zero dimensions are included. The dictionary is sorted
|
||||
by indices (keys) in ascending order for consistent output.
|
||||
Example: ``{10: 0.5, 245: 0.8, 1023: 1.2, 5678: 0.5}``
|
||||
|
||||
Raises:
|
||||
TypeError: If ``input`` is not a string.
|
||||
ValueError: If input is empty or whitespace-only.
|
||||
RuntimeError: If model inference fails.
|
||||
|
||||
Examples:
|
||||
>>> # Query embedding
|
||||
>>> query_emb = DefaultLocalSparseEmbedding(encoding_type="query")
|
||||
>>> query_vec = query_emb.embed("machine learning")
|
||||
>>> isinstance(query_vec, dict)
|
||||
True
|
||||
|
||||
Note:
|
||||
- First call may be slower due to model loading
|
||||
- Subsequent calls are much faster as the model stays in memory
|
||||
- GPU acceleration provides significant speedup
|
||||
- Sparse vectors are memory-efficient (only store non-zero values)
|
||||
"""
|
||||
if not isinstance(input, str):
|
||||
raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}")
|
||||
|
||||
input = input.strip()
|
||||
if not input:
|
||||
raise ValueError("Input text cannot be empty or whitespace only")
|
||||
|
||||
try:
|
||||
model = self._get_model()
|
||||
|
||||
# Use appropriate encoding method based on type
|
||||
if self._encoding_type == "document" and hasattr(model, "encode_document"):
|
||||
# Use document encoding
|
||||
sparse_matrix = model.encode_document([input])
|
||||
elif hasattr(model, "encode_query"):
|
||||
# Use query encoding (default)
|
||||
sparse_matrix = model.encode_query([input])
|
||||
else:
|
||||
# Fallback: manual implementation for older sentence-transformers
|
||||
return self._manual_sparse_encode(input)
|
||||
|
||||
# Convert sparse matrix to dictionary
|
||||
# SPLADE returns shape [1, vocab_size] for single input
|
||||
|
||||
# Check if it's a sparse matrix (duck typing - has toarray method)
|
||||
if hasattr(sparse_matrix, "toarray"):
|
||||
# Sparse matrix (CSR/CSC/etc.) - convert to dense array
|
||||
sparse_array = sparse_matrix[0].toarray().flatten()
|
||||
sparse_dict = {
|
||||
int(idx): float(val)
|
||||
for idx, val in enumerate(sparse_array)
|
||||
if val > 0
|
||||
}
|
||||
else:
|
||||
# Dense array format (numpy array or similar)
|
||||
if isinstance(sparse_matrix, np.ndarray):
|
||||
sparse_array = sparse_matrix[0]
|
||||
else:
|
||||
sparse_array = sparse_matrix
|
||||
|
||||
sparse_dict = {
|
||||
int(idx): float(val)
|
||||
for idx, val in enumerate(sparse_array)
|
||||
if val > 0
|
||||
}
|
||||
|
||||
# Sort by indices (keys) to ensure consistent ordering
|
||||
return dict(sorted(sparse_dict.items()))
|
||||
|
||||
except Exception as e:
|
||||
if isinstance(e, (TypeError, ValueError)):
|
||||
raise
|
||||
raise RuntimeError(f"Failed to generate sparse embedding: {e!s}") from e
|
||||
|
||||
def _manual_sparse_encode(self, input: str) -> SparseVectorType:
|
||||
"""Fallback manual SPLADE encoding for older sentence-transformers.
|
||||
|
||||
Args:
|
||||
input (str): Input text to encode.
|
||||
|
||||
Returns:
|
||||
SparseVectorType: Sparse vector as dictionary.
|
||||
"""
|
||||
import torch
|
||||
|
||||
model = self._get_model()
|
||||
|
||||
# Tokenize input
|
||||
features = model.tokenize([input])
|
||||
|
||||
# Move to correct device
|
||||
features = {k: v.to(model.device) for k, v in features.items()}
|
||||
|
||||
# Forward pass with no gradient
|
||||
with torch.no_grad():
|
||||
embeddings = model.forward(features)
|
||||
|
||||
# Get logits from model output
|
||||
# SPLADE models typically output 'token_embeddings'
|
||||
if isinstance(embeddings, dict) and "token_embeddings" in embeddings:
|
||||
logits = embeddings["token_embeddings"][0] # First batch item
|
||||
elif hasattr(embeddings, "token_embeddings"):
|
||||
logits = embeddings.token_embeddings[0]
|
||||
# Fallback: try to get first value
|
||||
elif isinstance(embeddings, dict):
|
||||
logits = next(iter(embeddings.values()))[0]
|
||||
else:
|
||||
logits = embeddings[0]
|
||||
|
||||
# Apply SPLADE activation: log(1 + relu(x))
|
||||
relu_log = torch.log(1 + torch.relu(logits))
|
||||
|
||||
# Max pooling over token dimension (reduce to vocab size)
|
||||
if relu_log.dim() > 1:
|
||||
sparse_vec, _ = torch.max(relu_log, dim=0)
|
||||
else:
|
||||
sparse_vec = relu_log
|
||||
|
||||
# Convert to sparse dictionary (only non-zero values)
|
||||
sparse_vec_np = sparse_vec.cpu().numpy()
|
||||
sparse_dict = {
|
||||
int(idx): float(val) for idx, val in enumerate(sparse_vec_np) if val > 0
|
||||
}
|
||||
|
||||
# Sort by indices (keys) to ensure consistent ordering
|
||||
return dict(sorted(sparse_dict.items()))
|
||||
|
||||
def _get_model(self):
|
||||
"""Load or retrieve the SPLADE model from class-level cache.
|
||||
|
||||
Returns:
|
||||
SentenceTransformer: The loaded SPLADE model instance.
|
||||
|
||||
Raises:
|
||||
ImportError: If required packages are not installed.
|
||||
ValueError: If model cannot be loaded.
|
||||
|
||||
Note:
|
||||
Models are cached at class level and shared across all instances
|
||||
with the same (model_name, model_source, device) configuration.
|
||||
This allows memory-efficient usage when creating multiple instances
|
||||
with different encoding_type settings.
|
||||
"""
|
||||
# Check class-level cache first
|
||||
if self._cache_key in self._model_cache:
|
||||
return self._model_cache[self._cache_key]
|
||||
|
||||
# Use parent class method to load model
|
||||
model = super()._get_model()
|
||||
|
||||
# Cache the model at class level
|
||||
self._model_cache[self._cache_key] = model
|
||||
|
||||
return model
|
||||
@@ -0,0 +1,150 @@
|
||||
# Copyright 2025-present the zvec project
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Literal, Optional
|
||||
|
||||
from ..tool import require_module
|
||||
|
||||
|
||||
class SentenceTransformerFunctionBase:
|
||||
"""Base class for Sentence Transformer functions (both dense and sparse).
|
||||
|
||||
This base class provides common functionality for loading and managing
|
||||
sentence-transformers models from Hugging Face or ModelScope. It supports
|
||||
both dense models (e.g., all-MiniLM-L6-v2) and sparse models (e.g., SPLADE).
|
||||
|
||||
This class is not meant to be used directly. Use concrete implementations:
|
||||
- ``SentenceTransformerEmbeddingFunction`` for dense embeddings
|
||||
- ``SentenceTransformerSparseEmbeddingFunction`` for sparse embeddings
|
||||
- ``DefaultDenseEmbedding`` for default dense embeddings
|
||||
- ``DefaultSparseEmbedding`` for default sparse embeddings
|
||||
|
||||
Args:
|
||||
model_name (str): Model identifier or local path.
|
||||
model_source (Literal["huggingface", "modelscope"]): Model source.
|
||||
device (Optional[str]): Device to run the model on.
|
||||
|
||||
Note:
|
||||
- This is an internal base class for code reuse
|
||||
- Subclasses should inherit from appropriate Protocol (Dense/Sparse)
|
||||
- Provides model loading and management functionality
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str,
|
||||
model_source: Literal["huggingface", "modelscope"] = "huggingface",
|
||||
device: Optional[str] = None,
|
||||
):
|
||||
"""Initialize the base Sentence Transformer functionality.
|
||||
|
||||
Args:
|
||||
model_name (str): Model identifier or local path.
|
||||
model_source (Literal["huggingface", "modelscope"]): Model source.
|
||||
device (Optional[str]): Device to run the model on.
|
||||
|
||||
Raises:
|
||||
ValueError: If model_source is invalid.
|
||||
"""
|
||||
# Validate model_source
|
||||
if model_source not in ("huggingface", "modelscope"):
|
||||
raise ValueError(
|
||||
f"Invalid model_source: '{model_source}'. "
|
||||
"Must be 'huggingface' or 'modelscope'."
|
||||
)
|
||||
|
||||
self._model_name = model_name
|
||||
self._model_source = model_source
|
||||
self._device = device
|
||||
self._model = None
|
||||
|
||||
@property
|
||||
def model_name(self) -> str:
|
||||
"""str: The Sentence Transformer model name currently in use."""
|
||||
return self._model_name
|
||||
|
||||
@property
|
||||
def model_source(self) -> str:
|
||||
"""str: The model source being used ("huggingface" or "modelscope")."""
|
||||
return self._model_source
|
||||
|
||||
@property
|
||||
def device(self) -> str:
|
||||
"""str: The device the model is running on."""
|
||||
model = self._get_model()
|
||||
if model is not None:
|
||||
return str(model.device)
|
||||
return self._device or "cpu"
|
||||
|
||||
def _get_model(self):
|
||||
"""Load or retrieve the Sentence Transformer model.
|
||||
|
||||
Returns:
|
||||
SentenceTransformer or SparseEncoder: The loaded model instance.
|
||||
|
||||
Raises:
|
||||
ImportError: If required packages are not installed.
|
||||
ValueError: If model cannot be loaded.
|
||||
"""
|
||||
# Return cached model if exists
|
||||
if self._model is not None:
|
||||
return self._model
|
||||
|
||||
# Load model
|
||||
try:
|
||||
sentence_transformers = require_module("sentence_transformers")
|
||||
|
||||
if self._model_source == "modelscope":
|
||||
# Load from ModelScope
|
||||
require_module("modelscope")
|
||||
from modelscope.hub.snapshot_download import snapshot_download
|
||||
|
||||
# Download model to cache
|
||||
model_dir = snapshot_download(self._model_name)
|
||||
|
||||
# Load from local path
|
||||
self._model = sentence_transformers.SentenceTransformer(
|
||||
model_dir, device=self._device, trust_remote_code=True
|
||||
)
|
||||
else:
|
||||
# Load from Hugging Face (default)
|
||||
self._model = sentence_transformers.SentenceTransformer(
|
||||
self._model_name, device=self._device, trust_remote_code=True
|
||||
)
|
||||
|
||||
return self._model
|
||||
|
||||
except ImportError as e:
|
||||
if "modelscope" in str(e) and self._model_source == "modelscope":
|
||||
raise ImportError(
|
||||
"ModelScope support requires the 'modelscope' package. "
|
||||
"Please install it with: pip install modelscope"
|
||||
) from e
|
||||
raise
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Failed to load Sentence Transformer model '{self._model_name}' "
|
||||
f"from {self._model_source}: {e!s}"
|
||||
) from e
|
||||
|
||||
def _is_sparse_model(self) -> bool:
|
||||
"""Check if the loaded model is a sparse encoder (e.g., SPLADE).
|
||||
|
||||
Returns:
|
||||
bool: True if model supports sparse encoding.
|
||||
"""
|
||||
model = self._get_model()
|
||||
# Check if model has sparse encoding methods
|
||||
return hasattr(model, "encode_query") or hasattr(model, "encode_document")
|
||||
@@ -0,0 +1,396 @@
|
||||
# Copyright 2025-present the zvec project
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Literal, Optional
|
||||
|
||||
from ..model.doc import Doc, DocList
|
||||
from ..tool import require_module
|
||||
from .rerank_function import RerankFunction
|
||||
from .sentence_transformer_function import SentenceTransformerFunctionBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..model.schema import FieldSchema, VectorSchema
|
||||
|
||||
|
||||
class DefaultLocalReRanker(SentenceTransformerFunctionBase, RerankFunction):
|
||||
"""Re-ranker using Sentence Transformer cross-encoder models for semantic re-ranking.
|
||||
|
||||
This re-ranker leverages pre-trained cross-encoder models to perform deep semantic
|
||||
re-ranking of search results. It runs locally without API calls, supports GPU
|
||||
acceleration, and works with models from Hugging Face or ModelScope.
|
||||
|
||||
Cross-encoder models evaluate query-document pairs jointly, providing more
|
||||
accurate relevance scores than bi-encoder (embedding-based) similarity.
|
||||
|
||||
Args:
|
||||
query (str): Query text for semantic re-ranking. **Required**.
|
||||
rerank_field (Optional[str], optional): Document field name to use as
|
||||
re-ranking input text. **Required** (e.g., "content", "title", "body").
|
||||
model_name (str, optional): Cross-encoder model identifier or local path.
|
||||
Defaults to ``"cross-encoder/ms-marco-MiniLM-L6-v2"`` (MS MARCO MiniLM).
|
||||
Common options:
|
||||
- ``"cross-encoder/ms-marco-MiniLM-L6-v2"``: Lightweight, fast (~80MB, recommended)
|
||||
- ``"cross-encoder/ms-marco-MiniLM-L12-v2"``: Better accuracy (~120MB)
|
||||
- ``"BAAI/bge-reranker-base"``: BGE Reranker Base (~280MB)
|
||||
- ``"BAAI/bge-reranker-large"``: BGE Reranker Large (highest quality, ~560MB)
|
||||
model_source (Literal["huggingface", "modelscope"], optional): Model source.
|
||||
Defaults to ``"huggingface"``.
|
||||
- ``"huggingface"``: Load from Hugging Face Hub
|
||||
- ``"modelscope"``: Load from ModelScope (recommended for users in China)
|
||||
device (Optional[str], optional): Device to run the model on.
|
||||
Options: ``"cpu"``, ``"cuda"``, ``"mps"`` (for Apple Silicon), or ``None``
|
||||
for automatic detection. Defaults to ``None``.
|
||||
batch_size (int, optional): Batch size for processing query-document pairs.
|
||||
Larger values speed up processing but use more memory. Defaults to ``32``.
|
||||
|
||||
Attributes:
|
||||
query (str): The query text used for re-ranking.
|
||||
rerank_field (Optional[str]): Field name used for re-ranking input.
|
||||
model_name (str): The cross-encoder model being used.
|
||||
model_source (str): The model source ("huggingface" or "modelscope").
|
||||
device (str): The device the model is running on.
|
||||
|
||||
Raises:
|
||||
ValueError: If ``query`` is empty/None, ``rerank_field`` is None,
|
||||
or model cannot be loaded.
|
||||
TypeError: If input types are invalid.
|
||||
RuntimeError: If model inference fails.
|
||||
|
||||
Note:
|
||||
- Requires Python 3.10, 3.11, or 3.12
|
||||
- Requires ``sentence-transformers`` package: ``pip install sentence-transformers``
|
||||
- For ModelScope support, also requires: ``pip install modelscope``
|
||||
- First run downloads the model (~80-560MB depending on model) from chosen source
|
||||
- No API keys or network required after initial download
|
||||
- Cross-encoders are slower than bi-encoders but more accurate
|
||||
- GPU acceleration provides significant speedup (5-10x)
|
||||
|
||||
**MS MARCO MiniLM-L6-v2 Model (Default):**
|
||||
|
||||
The default model ``cross-encoder/ms-marco-MiniLM-L6-v2`` is a lightweight and
|
||||
efficient cross-encoder trained on MS MARCO dataset. It provides:
|
||||
|
||||
- Fast inference speed (suitable for real-time applications)
|
||||
- Small model size (~80MB, quick to download)
|
||||
- Good balance between speed and accuracy
|
||||
- Trained on 500K+ query-document pairs
|
||||
- Public availability without authentication
|
||||
|
||||
**For users in China:**
|
||||
|
||||
If you encounter Hugging Face access issues, use ModelScope instead:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Recommended for users in China
|
||||
reranker = SentenceTransformerReRanker(
|
||||
query="机器学习算法",
|
||||
rerank_field="content",
|
||||
model_source="modelscope"
|
||||
)
|
||||
|
||||
Alternatively, use Hugging Face mirror:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
export HF_ENDPOINT=https://hf-mirror.com
|
||||
|
||||
Examples:
|
||||
>>> # Basic usage with default MS MARCO MiniLM model
|
||||
>>> from zvec.extension import SentenceTransformerReRanker
|
||||
>>>
|
||||
>>> reranker = SentenceTransformerReRanker(
|
||||
... query="machine learning algorithms",
|
||||
... rerank_field="content"
|
||||
... )
|
||||
>>>
|
||||
>>> # Use in collection.query()
|
||||
>>> results = collection.query(
|
||||
... data={"vector_field": query_vector},
|
||||
... reranker=reranker,
|
||||
... topk=20
|
||||
... )
|
||||
|
||||
>>> # Using ModelScope for users in China
|
||||
>>> reranker = SentenceTransformerReRanker(
|
||||
... query="深度学习",
|
||||
... rerank_field="content",
|
||||
... model_source="modelscope"
|
||||
... )
|
||||
|
||||
>>> # Using larger model for better quality
|
||||
>>> reranker = SentenceTransformerReRanker(
|
||||
... query="neural networks",
|
||||
... rerank_field="content",
|
||||
... model_name="BAAI/bge-reranker-large",
|
||||
... device="cuda",
|
||||
... batch_size=64
|
||||
... )
|
||||
|
||||
>>> # Direct rerank call (for testing)
|
||||
>>> query_results = {
|
||||
... "vector1": [
|
||||
... Doc(id="1", score=0.9, fields={"content": "Machine learning is..."}),
|
||||
... Doc(id="2", score=0.8, fields={"content": "Deep learning is..."}),
|
||||
... ]
|
||||
... }
|
||||
>>> reranked = reranker.rerank(query_results)
|
||||
>>> for doc in reranked:
|
||||
... print(f"ID: {doc.id}, Score: {doc.score:.4f}")
|
||||
ID: 2, Score: 0.9234
|
||||
ID: 1, Score: 0.8567
|
||||
|
||||
See Also:
|
||||
- ``RerankFunction``: Abstract base class for re-rankers
|
||||
- ``QwenReRanker``: Re-ranker using Qwen API
|
||||
- ``RrfReRanker``: Multi-vector re-ranker using RRF
|
||||
- ``WeightedReRanker``: Multi-vector re-ranker using weighted scores
|
||||
|
||||
References:
|
||||
- MS MARCO Cross-Encoder: https://huggingface.co/cross-encoder/ms-marco-MiniLM-L6-v2
|
||||
- BGE Reranker: https://huggingface.co/BAAI/bge-reranker-base
|
||||
- Cross-Encoder vs Bi-Encoder: https://www.sbert.net/examples/applications/cross-encoder/README.html
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
query: Optional[str] = None,
|
||||
rerank_field: Optional[str] = None,
|
||||
model_name: str = "cross-encoder/ms-marco-MiniLM-L6-v2",
|
||||
model_source: Literal["huggingface", "modelscope"] = "huggingface",
|
||||
device: Optional[str] = None,
|
||||
batch_size: int = 32,
|
||||
):
|
||||
"""Initialize SentenceTransformerReRanker with query and configuration.
|
||||
|
||||
Args:
|
||||
query (Optional[str]): Query text for semantic matching. Required.
|
||||
rerank_field (Optional[str]): Document field for re-ranking input.
|
||||
model_name (str): Cross-encoder model identifier.
|
||||
model_source (Literal["huggingface", "modelscope"]): Model source.
|
||||
device (Optional[str]): Target device ("cpu", "cuda", "mps", or None).
|
||||
batch_size (int): Batch size for processing query-document pairs.
|
||||
|
||||
Raises:
|
||||
ValueError: If query is empty or model cannot be loaded.
|
||||
"""
|
||||
# Initialize base class for model loading
|
||||
SentenceTransformerFunctionBase.__init__(
|
||||
self, model_name=model_name, model_source=model_source, device=device
|
||||
)
|
||||
|
||||
# Initialize rerank parameters
|
||||
self._rerank_field = rerank_field
|
||||
|
||||
# Validate query
|
||||
if not query:
|
||||
raise ValueError("Query is required for DefaultLocalReRanker")
|
||||
self._query = query
|
||||
self._batch_size = batch_size
|
||||
|
||||
# Load and validate cross-encoder model
|
||||
model = self._get_model()
|
||||
if not hasattr(model, "predict"):
|
||||
raise ValueError(
|
||||
f"Model '{model_name}' does not appear to be a cross-encoder model. "
|
||||
"Cross-encoder models should have a 'predict' method."
|
||||
)
|
||||
self._model = model
|
||||
|
||||
def _get_model(self):
|
||||
"""Load or retrieve the CrossEncoder model.
|
||||
|
||||
This overrides the base class method to load CrossEncoder instead of
|
||||
SentenceTransformer, as reranking requires cross-encoder models.
|
||||
|
||||
Returns:
|
||||
CrossEncoder: The loaded cross-encoder model instance.
|
||||
|
||||
Raises:
|
||||
ImportError: If required packages are not installed.
|
||||
ValueError: If model cannot be loaded.
|
||||
"""
|
||||
# Return cached model if exists
|
||||
if self._model is not None:
|
||||
return self._model
|
||||
|
||||
# Load cross-encoder model
|
||||
try:
|
||||
sentence_transformers = require_module("sentence_transformers")
|
||||
|
||||
if self._model_source == "modelscope":
|
||||
# Load from ModelScope
|
||||
require_module("modelscope")
|
||||
from modelscope.hub.snapshot_download import snapshot_download
|
||||
|
||||
# Download model to cache
|
||||
model_dir = snapshot_download(self._model_name)
|
||||
|
||||
# Load CrossEncoder from local path
|
||||
model = sentence_transformers.CrossEncoder(
|
||||
model_dir, device=self._device
|
||||
)
|
||||
else:
|
||||
# Load CrossEncoder from Hugging Face (default)
|
||||
model = sentence_transformers.CrossEncoder(
|
||||
self._model_name, device=self._device
|
||||
)
|
||||
|
||||
return model
|
||||
|
||||
except ImportError as e:
|
||||
if "modelscope" in str(e) and self._model_source == "modelscope":
|
||||
raise ImportError(
|
||||
"ModelScope support requires the 'modelscope' package. "
|
||||
"Please install it with: pip install modelscope"
|
||||
) from e
|
||||
raise
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Failed to load CrossEncoder model '{self._model_name}' "
|
||||
f"from {self._model_source}: {e!s}"
|
||||
) from e
|
||||
|
||||
@property
|
||||
def rerank_field(self) -> Optional[str]:
|
||||
"""Optional[str]: Field name used as re-ranking input."""
|
||||
return self._rerank_field
|
||||
|
||||
@property
|
||||
def query(self) -> str:
|
||||
"""str: Query text used for semantic re-ranking."""
|
||||
return self._query
|
||||
|
||||
@property
|
||||
def batch_size(self) -> int:
|
||||
"""int: Batch size for processing query-document pairs."""
|
||||
return self._batch_size
|
||||
|
||||
def rerank(
|
||||
self,
|
||||
query_results: list[list[Doc]],
|
||||
topn: int = 10,
|
||||
*,
|
||||
fields: list[FieldSchema | VectorSchema] | None = None, # noqa: ARG002
|
||||
) -> DocList:
|
||||
"""Re-rank documents using Sentence Transformer cross-encoder model.
|
||||
|
||||
Evaluates each query-document pair using the cross-encoder model to compute
|
||||
relevance scores. Documents are then sorted by these scores and the top-k
|
||||
results are returned.
|
||||
|
||||
Args:
|
||||
query_results (list[list[Doc]]): Per-sub-query lists of retrieved
|
||||
documents. Documents from all lists are deduplicated and
|
||||
re-ranked together.
|
||||
topn (int): Maximum number of documents to return.
|
||||
fields: Unused; present for interface compatibility.
|
||||
|
||||
Returns:
|
||||
list[Doc]: Re-ranked documents (up to ``topn``) with updated ``score``
|
||||
fields containing relevance scores from the cross-encoder model.
|
||||
|
||||
Raises:
|
||||
ValueError: If no valid documents are found or model inference fails.
|
||||
|
||||
Note:
|
||||
- Duplicate documents (same ID) across fields are processed once
|
||||
- Documents with empty/missing ``rerank_field`` content are skipped
|
||||
- Returned scores are logits from the cross-encoder model
|
||||
- Higher scores indicate higher relevance
|
||||
- Processing time is O(n) where n is the number of documents
|
||||
|
||||
Examples:
|
||||
>>> reranker = SentenceTransformerReRanker(
|
||||
... query="machine learning",
|
||||
... topn=3,
|
||||
... rerank_field="content"
|
||||
... )
|
||||
>>> query_results = {
|
||||
... "vector1": [
|
||||
... Doc(id="1", score=0.9, fields={"content": "ML basics"}),
|
||||
... Doc(id="2", score=0.8, fields={"content": "DL tutorial"}),
|
||||
... ]
|
||||
... }
|
||||
>>> reranked = reranker.rerank(query_results)
|
||||
>>> len(reranked) <= 3
|
||||
True
|
||||
"""
|
||||
if not query_results:
|
||||
return []
|
||||
|
||||
# Accept both dict (legacy) and list formats
|
||||
if isinstance(query_results, dict):
|
||||
query_results = list(query_results.values())
|
||||
|
||||
# Collect and deduplicate documents
|
||||
id_to_doc: dict[str, Doc] = {}
|
||||
doc_ids: list[str] = []
|
||||
contents: list[str] = []
|
||||
|
||||
for query_result in query_results:
|
||||
for doc in query_result:
|
||||
doc_id = doc.id
|
||||
if doc_id in id_to_doc:
|
||||
continue
|
||||
|
||||
# Extract text content from specified field
|
||||
field_value = doc.field(self.rerank_field)
|
||||
rank_content = str(field_value).strip() if field_value else ""
|
||||
if not rank_content:
|
||||
continue
|
||||
|
||||
id_to_doc[doc_id] = doc
|
||||
doc_ids.append(doc_id)
|
||||
contents.append(rank_content)
|
||||
|
||||
if not contents:
|
||||
raise ValueError("No documents to rerank")
|
||||
|
||||
try:
|
||||
# Use standard cross-encoder predict method
|
||||
pairs = [[self.query, content] for content in contents]
|
||||
scores = self._model.predict(
|
||||
pairs,
|
||||
batch_size=self.batch_size,
|
||||
show_progress_bar=False,
|
||||
convert_to_numpy=True,
|
||||
)
|
||||
|
||||
# Convert to float list if needed
|
||||
if hasattr(scores, "tolist"):
|
||||
scores = scores.tolist()
|
||||
else:
|
||||
scores = [float(s) for s in scores]
|
||||
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to compute rerank scores: {e!s}") from e
|
||||
|
||||
# Create scored documents
|
||||
scored_docs = [
|
||||
(doc_ids[i], id_to_doc[doc_ids[i]], scores[i]) for i in range(len(doc_ids))
|
||||
]
|
||||
|
||||
# Sort by score (descending) and take top-k
|
||||
scored_docs.sort(key=lambda x: x[2], reverse=True)
|
||||
top_scored_docs = scored_docs[:topn]
|
||||
|
||||
# Build result list with updated scores
|
||||
results: DocList = []
|
||||
for _, doc, score in top_scored_docs:
|
||||
new_doc = doc._replace(score=score)
|
||||
results.append(new_doc)
|
||||
|
||||
return results
|
||||
Reference in New Issue
Block a user