chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 12:47:42 +08:00
commit be3ef883e1
1214 changed files with 431743 additions and 0 deletions
@@ -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()))