xiaoxue_admin
/
zvec
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: support ai extension (#88)

This commit is contained in:
Cuiys 2026-02-12 18:33:41 +08:00 committed by GitHub
parent e7ad7cc31e
commit 1fa15ce96b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
24 changed files with 6358 additions and 665 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
.github/workflows/linux_arm64_docker_ci.yml vendored
View File

@ -69,9 +69,9 @@ jobs:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
shell: bash
- name: Install Ruff
- name: Install dependencies
run: |
${{ env.PIP_BIN }} install --upgrade pip ruff
${{ env.PIP_BIN }} install --upgrade pip ruff==v0.14.4 clang-format==18.1.8 pybind11==3.0 pytest pytest-cov
shell: bash
- name: Run Ruff Linter
@ -88,7 +88,6 @@ jobs:
- name: Run clang-format Check
run: |
${{ env.PIP_BIN }} install clang-format==18.1.8
cd "$CLEAN_WORKSPACE"
@ -120,11 +119,6 @@ jobs:
${{ env.PIP_BIN }} install -v . --config-settings='cmake.define.BUILD_TOOLS="ON"'
shell: bash
- name: Install test dependencies
run: |
${{ env.PIP_BIN }} install pytest pytest-cov
shell: bash
- name: Run Python Tests with Coverage
run: |
cd "$CLEAN_WORKSPACE"
@ -133,7 +127,6 @@ jobs:
- name: Run Cpp Tests
run: |
${{ env.PIP_BIN }} install pybind11==3.0
cd "$CLEAN_WORKSPACE/build"
make unittest -j$(nproc)
shell: bash

11
.github/workflows/linux_x64_docker_ci.yml vendored
View File

@ -69,9 +69,9 @@ jobs:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
shell: bash
- name: Install Ruff
- name: Install dependencies
run: |
${{ env.PIP_BIN }} install --upgrade pip ruff
${{ env.PIP_BIN }} install --upgrade pip ruff==v0.14.4 clang-format==18.1.8 pybind11==3.0 pytest pytest-cov
shell: bash
- name: Run Ruff Linter
@ -88,7 +88,6 @@ jobs:
- name: Run clang-format Check
run: |
${{ env.PIP_BIN }} install clang-format==18.1.8
cd "$CLEAN_WORKSPACE"
@ -120,11 +119,6 @@ jobs:
${{ env.PIP_BIN }} install -v . --config-settings='cmake.define.BUILD_TOOLS="ON"'
shell: bash
- name: Install test dependencies
run: |
${{ env.PIP_BIN }} install pytest pytest-cov
shell: bash
- name: Run Python Tests with Coverage
run: |
cd "$CLEAN_WORKSPACE"
@ -133,7 +127,6 @@ jobs:
- name: Run Cpp Tests
run: |
${{ env.PIP_BIN }} install pybind11==3.0
cd "$CLEAN_WORKSPACE/build"
make unittest -j$(nproc)
shell: bash

14
pyproject.toml
View File

@ -217,10 +217,21 @@ ignore = [
"E731", # Lambda assignment (used in callbacks)
"B019", # `functools.lru_cache` on methods (handled manually)
"PLR0912", # Too many branches
"PLC0105", # Ignore contravariant
"RUF002", # Ignore Unicode
]
fixable = ["ALL"]
unfixable = []
# Ignore all errors in docstrings
[tool.ruff.lint.pydocstyle]
convention = "google" # or "numpy", "pep257"
ignore-decorators = ["typing.overload"]
[tool.ruff.lint.flake8-type-checking]
# Don't check code examples in docstrings
quote-annotations = true
[tool.ruff.lint.isort]
required-imports = ["from __future__ import annotations"]
known-first-party = ["zvec"]
@ -237,6 +248,9 @@ known-first-party = ["zvec"]
"python/zvec/model/doc.py" = [
"RUF023", # Unused sort (for __slot__)
]
"python/zvec/extension/**" = [
"PLC0415", # Import outside top-level (dynamic imports in _get_model)
]
[tool.ruff.format]
indent-style = "space"

2026
python/tests/test_embedding.py
View File

File diff suppressed because it is too large Load Diff

934
python/tests/test_reranker.py
View File

@ -13,11 +13,23 @@
# limitations under the License.
from __future__ import annotations
from unittest.mock import patch
from unittest.mock import patch, MagicMock
import pytest
import math
import os
from zvec import RrfReRanker, WeightedReRanker, Doc, MetricType
from zvec import Doc, MetricType
from zvec.extension.multi_vector_reranker import (
RrfReRanker,
WeightedReRanker,
)
from zvec.extension.sentence_transformer_rerank_function import (
DefaultLocalReRanker,
)
from zvec.extension.qwen_rerank_function import QwenReRanker
# Set ZVEC_RUN_INTEGRATION_TESTS=1 to run real API tests
RUN_INTEGRATION_TESTS = os.environ.get("ZVEC_RUN_INTEGRATION_TESTS", "0") == "1"
# ----------------------------
@ -25,23 +37,20 @@ from zvec import RrfReRanker, WeightedReRanker, Doc, MetricType
# ----------------------------
class TestRrfReRanker:
def test_init(self):
reranker = RrfReRanker(
query="test", topn=5, rerank_field="content", rank_constant=100
)
assert reranker.query == "test"
reranker = RrfReRanker(topn=5, rerank_field="content", rank_constant=100)
assert reranker.topn == 5
assert reranker.rerank_field == "content"
assert reranker.rank_constant == 100
def test_rrf_score(self):
reranker = RrfReRanker(query="test", rank_constant=60)
reranker = RrfReRanker(rank_constant=60)
# 根据公式 1.0 / (k + rank + 1)，其中k=60
assert reranker._rrf_score(0) == 1.0 / (60 + 0 + 1)
assert reranker._rrf_score(1) == 1.0 / (60 + 1 + 1)
assert reranker._rrf_score(10) == 1.0 / (60 + 10 + 1)
def test_rerank(self):
reranker = RrfReRanker(query="test", topn=3)
reranker = RrfReRanker(topn=3)
doc1 = Doc(id="1", score=0.8)
doc2 = Doc(id="2", score=0.7)
@ -68,20 +77,18 @@ class TestWeightedReRanker:
def test_init(self):
weights = {"vector1": 0.7, "vector2": 0.3}
reranker = WeightedReRanker(
query="test",
topn=5,
rerank_field="content",
metric=MetricType.L2,
weights=weights,
)
assert reranker.query == "test"
assert reranker.topn == 5
assert reranker.rerank_field == "content"
assert reranker.metric == MetricType.L2
assert reranker.weights == weights
def test_normalize_score(self):
reranker = WeightedReRanker(query="test")
reranker = WeightedReRanker()
score = reranker._normalize_score(1.0, MetricType.L2)
expected = 1.0 - 2 * math.atan(1.0) / math.pi
@ -100,9 +107,7 @@ class TestWeightedReRanker:
def test_rerank(self):
weights = {"vector1": 0.7, "vector2": 0.3}
reranker = WeightedReRanker(
query="test", topn=3, weights=weights, metric=MetricType.L2
)
reranker = WeightedReRanker(topn=3, weights=weights, metric=MetricType.L2)
doc1 = Doc(id="1", score=0.8)
doc2 = Doc(id="2", score=0.7)
@ -121,64 +126,843 @@ class TestWeightedReRanker:
assert scores == sorted(scores, reverse=True)
# # ----------------------------
# # QwenReRanker Test Case
# # ----------------------------
# class TestQwenReRanker:
# def test_init_without_query(self):
# with pytest.raises(ValueError):
# QwenReRanker()
#
# def test_init_without_api_key(self):
# with patch.dict(os.environ, {"DASHSCOPE_API_KEY": ""}):
# with pytest.raises(ValueError, match="DashScope API key is required"):
# QwenReRanker(query="test")
#
# @patch.dict(os.environ, {"DASHSCOPE_API_KEY": "test_key"})
# def test_init_with_env_api_key(self):
# reranker = QwenReRanker(query="test")
# assert reranker.query == "test"
# assert reranker._api_key == "test_key"
#
# def test_model_property(self):
# reranker = QwenReRanker(query="test", api_key="test_key")
# assert reranker.model == "gte-rerank-v2"
#
# reranker = QwenReRanker(query="test", model="custom-model", api_key="test_key")
# assert reranker.model == "custom-model"
#
# def test_rerank_empty_results(self):
# reranker = QwenReRanker(query="test", api_key="test_key")
# results = reranker.rerank({})
# assert results == []
#
# def test_rerank_no_documents(self):
# reranker = QwenReRanker(query="test", api_key="test_key")
# query_results = {"vector1": [Doc(id="1")]}
# with pytest.raises(ValueError, match="No documents to rerank"):
# reranker.rerank(query_results)
#
# @pytest.mark.skip(reason="Qwen ReRanker is not available in CI")
# def test_rerank_success(self):
# reranker = QwenReRanker(
# topn=3,
# query="test",
# api_key="*",
# rerank_field="content",
# )
# query_results = {
# "vector1": [
# Doc(id="1", fields={"content": "This is a test document."}),
# Doc(id="2", fields={"content": "Another test document."}),
# Doc(id="3", fields={"content": "Yet another test document."}),
# Doc(id="4", fields={"content": "One more test document."}),
# ],
# "vector2": [
# Doc(id="5", fields={"content": "This is a test document2."}),
# Doc(id="6", fields={"content": "Another test document2."}),
# Doc(id="7", fields={"content": "Yet another test document2."}),
# Doc(id="8", fields={"content": "One more test document2."}),
# ],
# }
# results = reranker.rerank(query_results)
# assert len(results) == 3
# ----------------------------
# QwenReRanker Test Case
# ----------------------------
class TestQwenReRanker:
def test_init_without_query(self):
with pytest.raises(ValueError, match="Query is required for QwenReRanker"):
QwenReRanker(api_key="test_key")
def test_init_without_api_key(self):
with patch.dict(os.environ, {}, clear=True):
with pytest.raises(ValueError, match="DashScope API key is required"):
QwenReRanker(query="test")
@patch.dict(os.environ, {"DASHSCOPE_API_KEY": "test_key"})
def test_init_with_env_api_key(self):
reranker = QwenReRanker(query="test", rerank_field="content")
assert reranker.query == "test"
assert reranker._api_key == "test_key"
assert reranker.rerank_field == "content"
def test_init_with_explicit_api_key(self):
reranker = QwenReRanker(
query="test", api_key="explicit_key", rerank_field="content"
)
assert reranker.query == "test"
assert reranker._api_key == "explicit_key"
def test_model_property(self):
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
assert reranker.model == "gte-rerank-v2"
reranker = QwenReRanker(
query="test",
model="custom-model",
api_key="test_key",
rerank_field="content",
)
assert reranker.model == "custom-model"
def test_query_property(self):
reranker = QwenReRanker(
query="test query", api_key="test_key", rerank_field="content"
)
assert reranker.query == "test query"
def test_topn_property(self):
reranker = QwenReRanker(
query="test", topn=5, api_key="test_key", rerank_field="content"
)
assert reranker.topn == 5
def test_rerank_field_property(self):
reranker = QwenReRanker(query="test", api_key="test_key", rerank_field="title")
assert reranker.rerank_field == "title"
def test_rerank_empty_results(self):
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
results = reranker.rerank({})
assert results == []
def test_rerank_no_valid_documents(self):
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
# Document without the rerank_field
query_results = {"vector1": [Doc(id="1")]}
with pytest.raises(ValueError, match="No documents to rerank"):
reranker.rerank(query_results)
def test_rerank_skip_empty_content(self):
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
query_results = {
"vector1": [
Doc(id="1", fields={"content": ""}),
Doc(id="2", fields={"content": " "}),
]
}
with pytest.raises(ValueError, match="No documents to rerank"):
reranker.rerank(query_results)
@patch("zvec.extension.qwen_function.require_module")
def test_rerank_success(self, mock_require_module):
# Mock dashscope module
mock_dashscope = MagicMock()
mock_require_module.return_value = mock_dashscope
# Mock API response
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.output = {
"results": [
{"index": 0, "relevance_score": 0.95},
{"index": 1, "relevance_score": 0.85},
]
}
mock_dashscope.TextReRank.call.return_value = mock_response
reranker = QwenReRanker(
query="test query", topn=2, api_key="test_key", rerank_field="content"
)
query_results = {
"vector1": [
Doc(id="1", fields={"content": "Document 1"}),
Doc(id="2", fields={"content": "Document 2"}),
]
}
results = reranker.rerank(query_results)
assert len(results) == 2
assert results[0].id == "1"
assert results[0].score == 0.95
assert results[1].id == "2"
assert results[1].score == 0.85
# Verify API call
mock_dashscope.TextReRank.call.assert_called_once_with(
model="gte-rerank-v2",
query="test query",
documents=["Document 1", "Document 2"],
top_n=2,
return_documents=False,
)
@patch("zvec.extension.qwen_function.require_module")
def test_rerank_deduplicate_documents(self, mock_require_module):
# Mock dashscope module
mock_dashscope = MagicMock()
mock_require_module.return_value = mock_dashscope
# Mock API response
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.output = {
"results": [
{"index": 0, "relevance_score": 0.9},
]
}
mock_dashscope.TextReRank.call.return_value = mock_response
reranker = QwenReRanker(
query="test", topn=5, api_key="test_key", rerank_field="content"
)
# Same document in multiple vector results
doc1 = Doc(id="1", fields={"content": "Document 1"})
query_results = {"vector1": [doc1], "vector2": [doc1]}
results = reranker.rerank(query_results)
# Should only call API with document once
call_args = mock_dashscope.TextReRank.call.call_args
assert len(call_args[1]["documents"]) == 1
@patch("zvec.extension.qwen_function.require_module")
def test_rerank_api_error(self, mock_require_module):
# Mock dashscope module
mock_dashscope = MagicMock()
mock_require_module.return_value = mock_dashscope
# Mock API error response
mock_response = MagicMock()
mock_response.status_code = 400
mock_response.message = "Invalid request"
mock_response.code = "InvalidParameter"
mock_dashscope.TextReRank.call.return_value = mock_response
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
query_results = {"vector1": [Doc(id="1", fields={"content": "Document 1"})]}
with pytest.raises(ValueError, match="DashScope API error"):
reranker.rerank(query_results)
@patch("zvec.extension.qwen_function.require_module")
def test_rerank_runtime_error(self, mock_require_module):
# Mock dashscope module that raises exception
mock_dashscope = MagicMock()
mock_require_module.return_value = mock_dashscope
mock_dashscope.TextReRank.call.side_effect = Exception("Network error")
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
query_results = {"vector1": [Doc(id="1", fields={"content": "Document 1"})]}
with pytest.raises(RuntimeError, match="Failed to call DashScope API"):
reranker.rerank(query_results)
@pytest.mark.skipif(
not RUN_INTEGRATION_TESTS,
reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
)
def test_real_qwen_rerank(self):
"""Integration test with real DashScope TextReRank API.
To run this test, set environment variables:
export ZVEC_RUN_INTEGRATION_TESTS=1
export DASHSCOPE_API_KEY=your-api-key
"""
# Create reranker with real API
reranker = QwenReRanker(
query="What is machine learning?",
topn=3,
rerank_field="content",
model="gte-rerank-v2",
)
# Prepare test documents
query_results = {
"vector1": [
Doc(
id="1",
score=0.8,
fields={
"content": "Machine learning is a subset of artificial intelligence that focuses on building systems that can learn from data."
},
),
Doc(
id="2",
score=0.7,
fields={
"content": "The weather is nice today with clear skies and sunshine."
},
),
Doc(
id="3",
score=0.75,
fields={
"content": "Deep learning is a specialized branch of machine learning using neural networks with multiple layers."
},
),
],
"vector2": [
Doc(
id="4",
score=0.6,
fields={
"content": "Python is a popular programming language for data science and machine learning applications."
},
),
Doc(
id="5",
score=0.65,
fields={
"content": "A recipe for chocolate cake includes flour, sugar, eggs, and cocoa powder."
},
),
],
}
# Call real API
results = reranker.rerank(query_results)
# Verify results
assert len(results) <= 3, "Should return at most topn documents"
assert len(results) > 0, "Should return at least one document"
# All results should have valid scores
for doc in results:
assert hasattr(doc, "score"), "Each document should have a score"
assert isinstance(doc.score, (int, float)), "Score should be numeric"
assert doc.score > 0, "Score should be positive"
# Verify scores are in descending order
scores = [doc.score for doc in results]
assert scores == sorted(scores, reverse=True), (
"Results should be sorted by score in descending order"
)
# Verify relevant documents are ranked higher
# Document 1 and 3 are about machine learning, should rank higher than weather/recipe docs
result_ids = [doc.id for doc in results]
# At least one of the ML-related documents should be in top results
ml_related_docs = {"1", "3", "4"}
assert any(doc_id in ml_related_docs for doc_id in result_ids[:2]), (
"ML-related documents should rank higher"
)
# Print results for manual verification (useful during development)
print("\nReranking results:")
for i, doc in enumerate(results, 1):
print(f"{i}. ID={doc.id}, Score={doc.score:.4f}")
if doc.fields:
content = doc.field("content")
if content:
print(f" Content: {content[:80]}...")
# ----------------------------
# DefaultLocalReRanker Test Case
# ----------------------------
class TestDefaultLocalReRanker:
"""Test cases for DefaultLocalReRanker."""
def test_init_without_query(self):
"""Test initialization fails without query."""
with pytest.raises(
ValueError, match="Query is required for DefaultLocalReRanker"
):
DefaultLocalReRanker(rerank_field="content")
def test_init_with_empty_query(self):
"""Test initialization fails with empty query."""
with pytest.raises(
ValueError, match="Query is required for DefaultLocalReRanker"
):
DefaultLocalReRanker(query="", rerank_field="content")
@patch("zvec.extension.sentence_transformer_rerank_function.require_module")
def test_init_success(self, mock_require_module):
"""Test successful initialization with mocked model."""
# Mock sentence_transformers module
mock_st = MagicMock()
mock_model = MagicMock()
mock_model.predict = MagicMock() # Cross-encoder has predict method
mock_model.device = "cpu"
mock_st.CrossEncoder.return_value = mock_model
mock_require_module.return_value = mock_st
reranker = DefaultLocalReRanker(
query="test query",
topn=5,
rerank_field="content",
model_name="cross-encoder/ms-marco-MiniLM-L6-v2",
)
assert reranker.query == "test query"
assert reranker.topn == 5
assert reranker.rerank_field == "content"
assert reranker.model_name == "cross-encoder/ms-marco-MiniLM-L6-v2"
assert reranker.model_source == "huggingface"
assert reranker.batch_size == 32
@pytest.mark.skipif(
not RUN_INTEGRATION_TESTS,
reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
)
@patch("zvec.extension.sentence_transformer_rerank_function.require_module")
def test_init_with_custom_params(self, mock_require_module):
"""Test initialization with custom parameters."""
mock_st = MagicMock()
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_model.device = "cuda"
mock_st.CrossEncoder.return_value = mock_model
mock_require_module.return_value = mock_st
reranker = DefaultLocalReRanker(
query="custom query",
topn=10,
rerank_field="title",
model_name="cross-encoder/ms-marco-MiniLM-L12-v2",
model_source="modelscope",
device="cuda",
batch_size=64,
)
assert reranker.query == "custom query"
assert reranker.topn == 10
assert reranker.rerank_field == "title"
assert reranker.model_name == "cross-encoder/ms-marco-MiniLM-L12-v2"
assert reranker.model_source == "modelscope"
assert reranker.batch_size == 64
@patch("zvec.extension.sentence_transformer_rerank_function.require_module")
def test_init_invalid_model(self, mock_require_module):
"""Test initialization fails with non-cross-encoder model."""
# Mock a model without predict method (not a cross-encoder)
mock_st = MagicMock()
mock_model = MagicMock(spec=[]) # No predict method
mock_st.CrossEncoder.return_value = mock_model
mock_require_module.return_value = mock_st
with pytest.raises(ValueError, match="does not appear to be a cross-encoder"):
DefaultLocalReRanker(query="test", rerank_field="content")
def test_query_property(self):
"""Test query property."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test query", rerank_field="content")
assert reranker.query == "test query"
def test_topn_property(self):
"""Test topn property."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(
query="test", topn=15, rerank_field="content"
)
assert reranker.topn == 15
def test_rerank_field_property(self):
"""Test rerank_field property."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="title")
assert reranker.rerank_field == "title"
def test_batch_size_property(self):
"""Test batch_size property."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(
query="test", rerank_field="content", batch_size=128
)
assert reranker.batch_size == 128
def test_rerank_empty_results(self):
"""Test rerank with empty query_results."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
results = reranker.rerank({})
assert results == []
def test_rerank_no_valid_documents(self):
"""Test rerank with documents missing rerank_field."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
# Document without the rerank_field
query_results = {"vector1": [Doc(id="1")]}
with pytest.raises(ValueError, match="No documents to rerank"):
reranker.rerank(query_results)
def test_rerank_skip_empty_content(self):
"""Test rerank skips documents with empty content."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
query_results = {
"vector1": [
Doc(id="1", fields={"content": ""}),
Doc(id="2", fields={"content": " "}),
]
}
with pytest.raises(ValueError, match="No documents to rerank"):
reranker.rerank(query_results)
def test_rerank_success(self):
"""Test successful rerank with mocked model."""
# Mock standard cross-encoder model
mock_model = MagicMock()
# Mock predict method to return scores
import numpy as np
mock_scores = np.array([0.95, 0.85, 0.75])
mock_model.predict.return_value = mock_scores
mock_model.device = "cpu"
# Mock sentence_transformers module
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(
query="test query", topn=3, rerank_field="content"
)
query_results = {
"vector1": [
Doc(id="1", score=0.8, fields={"content": "Document 1"}),
Doc(id="2", score=0.7, fields={"content": "Document 2"}),
Doc(id="3", score=0.6, fields={"content": "Document 3"}),
]
}
results = reranker.rerank(query_results)
# Verify results
assert len(results) == 3
assert results[0].id == "1"
assert results[0].score == 0.95
assert results[1].id == "2"
assert results[1].score == 0.85
assert results[2].id == "3"
assert results[2].score == 0.75
# Verify model.predict was called correctly
assert mock_model.predict.called
call_args = mock_model.predict.call_args
pairs = call_args[0][0]
assert len(pairs) == 3
assert pairs[0] == ["test query", "Document 1"]
assert pairs[1] == ["test query", "Document 2"]
assert pairs[2] == ["test query", "Document 3"]
assert call_args[1]["batch_size"] == 32
assert call_args[1]["show_progress_bar"] is False
def test_rerank_with_topn_limit(self):
"""Test rerank respects topn limit."""
mock_model = MagicMock()
import numpy as np
mock_scores = np.array([0.9, 0.8, 0.7, 0.6, 0.5])
mock_model.predict.return_value = mock_scores
# Mock sentence_transformers module
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(
query="test", topn=2, rerank_field="content"
)
query_results = {
"vector1": [
Doc(id="1", fields={"content": "Doc 1"}),
Doc(id="2", fields={"content": "Doc 2"}),
Doc(id="3", fields={"content": "Doc 3"}),
Doc(id="4", fields={"content": "Doc 4"}),
Doc(id="5", fields={"content": "Doc 5"}),
]
}
results = reranker.rerank(query_results)
# Should only return top 2
assert len(results) == 2
assert results[0].id == "1"
assert results[0].score == 0.9
assert results[1].id == "2"
assert results[1].score == 0.8
def test_rerank_deduplicate_documents(self):
"""Test rerank deduplicates documents across multiple vectors."""
mock_model = MagicMock()
import numpy as np
mock_scores = np.array([0.95, 0.85])
mock_model.predict.return_value = mock_scores
# Mock sentence_transformers module
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(
query="test", topn=5, rerank_field="content"
)
# Same document in multiple vector results
doc1 = Doc(id="1", fields={"content": "Document 1"})
doc2 = Doc(id="2", fields={"content": "Document 2"})
query_results = {
"vector1": [doc1, doc2],
"vector2": [doc1], # doc1 appears in both
}
results = reranker.rerank(query_results)
# Should only process each document once
assert len(results) == 2
assert mock_model.predict.call_count == 1
call_args = mock_model.predict.call_args
pairs = call_args[0][0]
assert len(pairs) == 2 # Only 2 unique documents
def test_rerank_sorting(self):
"""Test rerank sorts documents by score in descending order."""
mock_model = MagicMock()
import numpy as np
# Return scores in non-sorted order
mock_scores = np.array([0.6, 0.9, 0.7])
mock_model.predict.return_value = mock_scores
# Mock sentence_transformers module
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(
query="test", topn=3, rerank_field="content"
)
query_results = {
"vector1": [
Doc(id="1", fields={"content": "Doc 1"}),
Doc(id="2", fields={"content": "Doc 2"}),
Doc(id="3", fields={"content": "Doc 3"}),
]
}
results = reranker.rerank(query_results)
# Should be sorted by score (descending)
assert len(results) == 3
assert results[0].id == "2" # score 0.9
assert results[0].score == 0.9
assert results[1].id == "3" # score 0.7
assert results[1].score == 0.7
assert results[2].id == "1" # score 0.6
assert results[2].score == 0.6
def test_rerank_model_error(self):
"""Test rerank handles model prediction errors."""
mock_model = MagicMock()
# Mock predict to raise exception
mock_model.predict.side_effect = Exception("Model inference error")
# Mock sentence_transformers module
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
query_results = {"vector1": [Doc(id="1", fields={"content": "Document 1"})]}
with pytest.raises(RuntimeError, match="Failed to compute rerank scores"):
reranker.rerank(query_results)
def test_rerank_with_custom_batch_size(self):
"""Test rerank uses custom batch_size."""
mock_model = MagicMock()
import numpy as np
mock_scores = np.array([0.9, 0.8])
mock_model.predict.return_value = mock_scores
# Mock sentence_transformers module
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(
query="test", rerank_field="content", batch_size=64
)
query_results = {
"vector1": [
Doc(id="1", fields={"content": "Doc 1"}),
Doc(id="2", fields={"content": "Doc 2"}),
]
}
reranker.rerank(query_results)
# Verify batch_size is passed to predict
call_args = mock_model.predict.call_args
assert call_args[1]["batch_size"] == 64
@pytest.mark.skipif(
not RUN_INTEGRATION_TESTS,
reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
)
def test_real_sentence_transformer_rerank(self):
"""Integration test with real SentenceTransformer cross-encoder model.
To run this test, set environment variable:
export ZVEC_RUN_INTEGRATION_TESTS=1
Note: This test requires sentence-transformers package and will
download the MS MARCO MiniLM model (~80MB) on first run.
"""
# Create reranker with real model (using default lightweight model)
reranker = DefaultLocalReRanker(
query="What is machine learning?",
topn=3,
rerank_field="content",
)
# Prepare test documents
query_results = {
"vector1": [
Doc(
id="1",
score=0.8,
fields={
"content": "Machine learning is a subset of artificial intelligence that focuses on building systems that can learn from data."
},
),
Doc(
id="2",
score=0.7,
fields={
"content": "The weather is nice today with clear skies and sunshine."
},
),
Doc(
id="3",
score=0.75,
fields={
"content": "Deep learning is a specialized branch of machine learning using neural networks with multiple layers."
},
),
],
"vector2": [
Doc(
id="4",
score=0.6,
fields={
"content": "Python is a popular programming language for data science and machine learning applications."
},
),
Doc(
id="5",
score=0.65,
fields={
"content": "A recipe for chocolate cake includes flour, sugar, eggs, and cocoa powder."
},
),
],
}
# Call real model
results = reranker.rerank(query_results)
# Verify results
assert len(results) <= 3, "Should return at most topn documents"
assert len(results) > 0, "Should return at least one document"
# All results should have valid scores
for doc in results:
assert hasattr(doc, "score"), "Each document should have a score"
assert isinstance(doc.score, (int, float)), "Score should be numeric"
# Verify scores are in descending order
scores = [doc.score for doc in results]
assert scores == sorted(scores, reverse=True), (
"Results should be sorted by score in descending order"
)
# Verify relevant documents are ranked higher
# Documents 1, 3, and 4 are about machine learning, should rank higher
result_ids = [doc.id for doc in results]
# At least one of the ML-related documents should be in top results
ml_related_docs = {"1", "3", "4"}
assert any(doc_id in ml_related_docs for doc_id in result_ids[:2]), (
"ML-related documents should rank higher"
)
# Print results for manual verification (useful during development)
print("\nSentenceTransformer Reranking results:")
for i, doc in enumerate(results, 1):
print(f"{i}. ID={doc.id}, Score={doc.score:.4f}")
if doc.fields:
content = doc.field("content")
if content:
print(f" Content: {content[:80]}...")

5
python/tests/test_util.py
View File

@ -87,8 +87,3 @@ def test_require_module_calls_importlib(mock_import_module):
mock_import_module.assert_called_once_with("test_module")
assert result is mock_module
def test_require_module_with_openai():
with pytest.raises(ImportError) as exc_info:
require_module("openai")

37
python/zvec/__init__.py
View File

@ -27,8 +27,27 @@ if TYPE_CHECKING:
from . import model as model
# —— Extensions & typing ——
from .extension import DenseEmbeddingFunction, ReRanker, RrfReRanker, WeightedReRanker
# —— Extensions ——
from .extension import (
BM25EmbeddingFunction,
DefaultLocalDenseEmbedding,
DefaultLocalReRanker,
DefaultLocalSparseEmbedding,
DenseEmbeddingFunction,
OpenAIDenseEmbedding,
OpenAIFunctionBase,
QwenDenseEmbedding,
QwenFunctionBase,
QwenReRanker,
QwenSparseEmbedding,
ReRanker,
RrfReRanker,
SentenceTransformerFunctionBase,
SparseEmbeddingFunction,
WeightedReRanker,
)
# —— Typing ——
from .model import param as param
from .model import schema as schema
@ -100,10 +119,22 @@ __all__ = [
"HnswQueryParam",
"IVFQueryParam",
# Extensions
"ReRanker",
"DenseEmbeddingFunction",
"SparseEmbeddingFunction",
"QwenFunctionBase",
"OpenAIFunctionBase",
"SentenceTransformerFunctionBase",
"ReRanker",
"DefaultLocalDenseEmbedding",
"DefaultLocalSparseEmbedding",
"BM25EmbeddingFunction",
"OpenAIDenseEmbedding",
"QwenDenseEmbedding",
"QwenSparseEmbedding",
"RrfReRanker",
"WeightedReRanker",
"DefaultLocalReRanker",
"QwenReRanker",
# Typing
"DataType",
"MetricType",

12
python/zvec/common/constants.py
View File

@ -16,7 +16,19 @@ from __future__ import annotations
from typing import Optional, Union
import numpy as np
from typing_extensions import TypeVar
# VectorType: DenseVectorType | SparseVectorType
DenseVectorType = Union[list[float], list[int], np.ndarray]
SparseVectorType = dict[int, float]
VectorType = Optional[Union[DenseVectorType, SparseVectorType]]
# Embeddable: Text | Image | Audio
TEXT = str
IMAGE = Union[str, bytes, np.ndarray] # file path, raw bytes, or numpy array
AUDIO = Union[str, bytes, np.ndarray] # file path, raw bytes, or numpy array
Embeddable = Optional[Union[TEXT, IMAGE, AUDIO]]
# Multimodal Embeddable
MD = TypeVar("MD", bound=Embeddable, contravariant=True)

29
python/zvec/extension/__init__.py
View File

@ -13,14 +13,37 @@
# limitations under the License.
from __future__ import annotations
from .embedding import DenseEmbeddingFunction, QwenEmbeddingFunction
from .rerank import QwenReRanker, ReRanker, RrfReRanker, WeightedReRanker
from .bm25_embedding_function import BM25EmbeddingFunction
from .embedding_function import DenseEmbeddingFunction, SparseEmbeddingFunction
from .multi_vector_reranker import 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 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",
"DefaultLocalDenseEmbedding",
"DefaultLocalReRanker",
"DefaultLocalSparseEmbedding",
"DenseEmbeddingFunction",
"QwenEmbeddingFunction",
"OpenAIDenseEmbedding",
"OpenAIFunctionBase",
"QwenDenseEmbedding",
"QwenFunctionBase",
"QwenReRanker",
"QwenSparseEmbedding",
"ReRanker",
"RrfReRanker",
"SentenceTransformerFunctionBase",
"SparseEmbeddingFunction",
"WeightedReRanker",
]

375
python/zvec/extension/bm25_embedding_function.py Normal file
View File

@ -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

188
python/zvec/extension/embedding.py
View File

@ -1,188 +0,0 @@
# 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 abc import ABC, abstractmethod
from functools import lru_cache
from http import HTTPStatus
from typing import Optional, Union
from ..tool import require_module
from ..typing import DataType
class DenseEmbeddingFunction(ABC):
"""Abstract base class for dense vector embedding functions.
Dense embedding functions map text to fixed-length real-valued vectors.
Subclasses must implement the ``embed()`` method.
Args:
dimension (int): Dimensionality of the output embedding vector.
data_type (DataType, optional): Numeric type of the embedding.
Defaults to ``DataType.VECTOR_FP32``.
Note:
This class is callable: ``embedding_func("text")`` is equivalent to
``embedding_func.embed("text")``.
"""
def __init__(self, dimension: int, data_type: DataType = DataType.VECTOR_FP32):
self._dimension = dimension
self._data_type = data_type
@property
def dimension(self) -> int:
"""int: The expected dimensionality of the embedding vector."""
return self._dimension
@property
def data_type(self) -> DataType:
"""DataType: The numeric data type of the embedding (e.g., VECTOR_FP32)."""
return self._data_type
@abstractmethod
def embed(self, text: str) -> list[Union[int, float]]:
"""Generate a dense embedding vector for the input text.
Args:
text (str): Input text to embed.
Returns:
list[Union[int, float]]: A list of numbers representing the embedding.
Length must equal ``self.dimension``.
"""
raise NotImplementedError
def __call__(self, text: str) -> list[Union[int, float]]:
return self.embed(text)
class SparseEmbeddingFunction(ABC):
"""Abstract base class for sparse vector embedding functions.
Sparse embedding functions map text to a dictionary of {index: weight},
where only non-zero dimensions are stored.
Note:
Subclasses must implement the ``embed()`` method.
"""
@abstractmethod
def embed(self, text: str) -> dict[int, float]:
"""Generate a sparse embedding for the input text.
Args:
text (str): Input text to embed.
Returns:
dict[int, float]: Mapping from dimension index to non-zero weight.
"""
raise NotImplementedError
class QwenEmbeddingFunction(DenseEmbeddingFunction):
"""Dense embedding function using Qwen (DashScope) Text Embedding API.
This implementation uses the DashScope service to generate embeddings
via Qwen's text embedding models (e.g., ``text-embedding-v4``).
Args:
dimension (int): Desired embedding dimension (e.g., 1024).
model (str, optional): DashScope embedding model name.
Defaults to ``"text-embedding-v4"``.
api_key (Optional[str], optional): DashScope API key. If not provided,
reads from ``DASHSCOPE_API_KEY`` environment variable.
Raises:
ValueError: If API key is missing or input text is invalid.
Note:
Requires the ``dashscope`` Python package.
Embedding results are cached using ``functools.lru_cache`` (maxsize=10).
"""
def __init__(
self,
dimension: int,
model: str = "text-embedding-v4",
api_key: Optional[str] = None,
):
super().__init__(dimension, DataType.VECTOR_FP32)
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")
@property
def model(self) -> str:
"""str: The DashScope embedding model name in use."""
return self._model
def _connection(self):
dashscope = require_module("dashscope")
dashscope.api_key = self._api_key
return dashscope
@lru_cache(maxsize=10)
def embed(self, text: str) -> list[Union[int, float]]:
"""
Generate embedding for a given text using Qwen (via DashScope).
Args:
text (str): Input text to embed. Must be non-empty and valid string.
Returns:
list[Union[int, float]]: The dense embedding vector.
Raises:
ValueError: If input is invalid or API response is malformed.
RuntimeError: If network or internal error occurs during API call.
"""
if not isinstance(text, str):
raise TypeError(f"Expected 'text' to be str, got {type(text).__name__}")
text = text.strip()
if not text:
raise ValueError("Input text cannot be empty or whitespace only")
resp = self._connection().TextEmbedding.call(
model=self.model, input=text, dimension=self.dimension, output_type="dense"
)
if resp.status_code != HTTPStatus.OK:
error_msg = getattr(resp, "message", "Unknown error")
error_detail = f"Status={resp.status_code}, Message={error_msg}"
raise ValueError(f"QwenEmbedding failed: {error_detail}")
output = getattr(resp, "output", None)
if not isinstance(output, dict):
raise ValueError("Invalid response: missing or malformed 'output' field")
embeddings = output.get("embeddings")
if not isinstance(embeddings, list):
raise ValueError(
"Invalid response: 'embeddings' field is missing or not a list"
)
if len(embeddings) != 1:
raise ValueError(
f"Expected 1 embedding, got {len(embeddings)}. Response: {resp}"
)
first_emb = embeddings[0]
if not isinstance(first_emb, dict):
raise ValueError("Invalid response: embedding item is not a dictionary")
return list(first_emb.get("embedding"))

148
python/zvec/extension/embedding_function.py Normal file
View File

@ -0,0 +1,148 @@
# 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_extensions 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.
"""
...

174
python/zvec/extension/multi_vector_reranker.py Normal file
View File

@ -0,0 +1,174 @@
# 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 heapq
import math
from collections import defaultdict
from typing import Optional
from ..model.doc import Doc
from ..typing import MetricType
from .rerank_function import RerankFunction
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.
It assigns higher weight to documents that appear early in multiple result lists.
The RRF score for a document at rank ``r`` is: ``1 / (k + r + 1)``,
where ``k`` is the rank constant.
Note:
This re-ranker is specifically designed for multi-vector scenarios where
query results from multiple vector fields need to be combined.
Args:
topn (int, optional): Number of top documents to return. Defaults to 10.
rerank_field (Optional[str], optional): Ignored by RRF. Defaults to None.
rank_constant (int, optional): Smoothing constant ``k`` in RRF formula.
Larger values reduce the impact of early ranks. Defaults to 60.
"""
def __init__(
self,
topn: int = 10,
rerank_field: Optional[str] = None,
rank_constant: int = 60,
):
super().__init__(topn=topn, rerank_field=rerank_field)
self._rank_constant = rank_constant
@property
def rank_constant(self) -> int:
return self._rank_constant
def _rrf_score(self, rank: int) -> float:
return 1.0 / (self._rank_constant + rank + 1)
def rerank(self, query_results: dict[str, list[Doc]]) -> list[Doc]:
"""Apply Reciprocal Rank Fusion to combine multiple query results.
Args:
query_results (dict[str, list[Doc]]): Results from one or more vector queries.
Returns:
list[Doc]: Re-ranked documents with RRF scores in the ``score`` field.
"""
rrf_scores: dict[str, float] = defaultdict(float)
id_to_doc: dict[str, Doc] = {}
for _, query_result in query_results.items():
for rank, doc in enumerate(query_result):
doc_id = doc.id
rrf_score = self._rrf_score(rank)
rrf_scores[doc_id] += rrf_score
if doc_id not in id_to_doc:
id_to_doc[doc_id] = doc
top_docs = heapq.nlargest(self.topn, rrf_scores.items(), key=lambda x: x[1])
results: list[Doc] = []
for doc_id, rrf_score in top_docs:
doc = id_to_doc[doc_id]
new_doc = doc._replace(score=rrf_score)
results.append(new_doc)
return results
class WeightedReRanker(RerankFunction):
"""Re-ranker that combines scores from multiple vector fields using weights.
Each vector field's relevance score is normalized based on its metric type,
then scaled by a user-provided weight. Final scores are summed across fields.
Note:
This re-ranker is specifically designed for multi-vector scenarios where
query results from multiple vector fields need to be combined with
configurable weights.
Args:
topn (int, optional): Number of top documents to return. Defaults to 10.
rerank_field (Optional[str], optional): Ignored. Defaults to None.
metric (MetricType, optional): Distance metric used for score normalization.
Defaults to ``MetricType.L2``.
weights (Optional[dict[str, float]], optional): Weight per vector field.
Fields not listed use weight 1.0. Defaults to None.
Note:
Supported metrics: L2, IP, COSINE. Scores are normalized to [0, 1].
"""
def __init__(
self,
topn: int = 10,
rerank_field: Optional[str] = None,
metric: MetricType = MetricType.L2,
weights: Optional[dict[str, float]] = None,
):
super().__init__(topn=topn, rerank_field=rerank_field)
self._weights = weights or {}
self._metric = metric
@property
def weights(self) -> dict[str, float]:
"""dict[str, float]: Weight mapping for vector fields."""
return self._weights
@property
def metric(self) -> MetricType:
"""MetricType: Distance metric used for score normalization."""
return self._metric
def rerank(self, query_results: dict[str, list[Doc]]) -> list[Doc]:
"""Combine scores from multiple vector fields using weighted sum.
Args:
query_results (dict[str, list[Doc]]): Results per vector field.
Returns:
list[Doc]: Re-ranked documents with combined scores in ``score`` field.
"""
weighted_scores: dict[str, float] = defaultdict(float)
id_to_doc: dict[str, Doc] = {}
for vector_name, query_result in query_results.items():
for _, doc in enumerate(query_result):
doc_id = doc.id
weighted_score = self._normalize_score(
doc.score, self.metric
) * self.weights.get(vector_name, 1.0)
weighted_scores[doc_id] += weighted_score
if doc_id not in id_to_doc:
id_to_doc[doc_id] = doc
top_docs = heapq.nlargest(
self.topn, weighted_scores.items(), key=lambda x: x[1]
)
results: list[Doc] = []
for doc_id, weighted_score in top_docs:
doc = id_to_doc[doc_id]
new_doc = doc._replace(score=weighted_score)
results.append(new_doc)
return results
def _normalize_score(self, score: float, metric: MetricType) -> float:
if metric == MetricType.L2:
return 1.0 - 2 * math.atan(score) / math.pi
if metric == MetricType.IP:
return 0.5 + math.atan(score) / math.pi
if metric == MetricType.COSINE:
return 1.0 - score / 2.0
raise ValueError("Unsupported metric type")

238
python/zvec/extension/openai_embedding_function.py Normal file
View File

@ -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

149
python/zvec/extension/openai_function.py Normal file
View File

@ -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

537
python/zvec/extension/qwen_embedding_function.py Normal file
View File

@ -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()))

186
python/zvec/extension/qwen_function.py Normal file
View File

@ -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

162
python/zvec/extension/qwen_rerank_function.py Normal file
View File

@ -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
from typing import Optional
from ..model.doc import Doc
from .qwen_function import QwenFunctionBase
from .rerank_function import RerankFunction
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**.
topn (int, optional): Maximum number of documents to return after re-ranking.
Defaults to 10.
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",
... topn=5,
... 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,
topn: int = 10,
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.
topn (int): Number of top results to return.
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)
RerankFunction.__init__(self, topn=topn, rerank_field=rerank_field)
if not query:
raise ValueError("Query is required for QwenReRanker")
self._query = query
@property
def query(self) -> str:
"""str: Query text used for semantic re-ranking."""
return self._query
def rerank(self, query_results: dict[str, list[Doc]]) -> list[Doc]:
"""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 (dict[str, list[Doc]]): Mapping from vector field names
to lists of retrieved documents. Documents from all fields are
deduplicated and re-ranked together.
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 fields 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 []
# Collect and deduplicate documents
id_to_doc: dict[str, Doc] = {}
doc_ids: list[str] = []
contents: list[str] = []
for _, query_result in query_results.items():
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=self.topn,
)
# Build result list with updated scores
results: list[Doc] = []
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

343
python/zvec/extension/rerank.py
View File

@ -1,343 +0,0 @@
# 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 heapq
import math
import os
from abc import ABC, abstractmethod
from collections import defaultdict
from http import HTTPStatus
from typing import Optional
from ..model.doc import Doc
from ..tool import require_module
from ..typing import MetricType
class ReRanker(ABC):
"""Abstract base class for re-ranking search results.
Re-rankers refine the output of one or more vector queries by applying
a secondary scoring strategy. They are used in the ``query()`` method of
``Collection`` via the ``reranker`` parameter.
Args:
query (Optional[str], optional): Query text used for re-ranking.
Required for LLM-based re-rankers. Defaults to None.
topn (int, optional): Number of top documents to return after re-ranking.
Defaults to 10.
rerank_field (Optional[str], optional): Field name used as input for
re-ranking (e.g., document title or body). Defaults to None.
Note:
Subclasses must implement the ``rerank()`` method.
"""
def __init__(
self,
query: Optional[str] = None,
topn: int = 10,
rerank_field: Optional[str] = None,
):
self._query = query
self._topn = topn
self._rerank_field = rerank_field
@property
def topn(self) -> int:
"""int: Number of top documents to return after re-ranking."""
return self._topn
@property
def query(self) -> str:
"""str: Query text used for re-ranking."""
return self._query
@property
def rerank_field(self) -> Optional[str]:
"""Optional[str]: Field name used as re-ranking input."""
return self._rerank_field
@abstractmethod
def rerank(self, query_results: dict[str, list[Doc]]) -> list[Doc]:
"""Re-rank documents from one or more vector queries.
Args:
query_results (dict[str, list[Doc]]): Mapping from vector field name
to list of retrieved documents (sorted by relevance).
Returns:
list[Doc]: Re-ranked list of documents (length ≤ ``topn``),
with updated ``score`` fields.
"""
raise NotImplementedError
class RrfReRanker(ReRanker):
"""Re-ranker using Reciprocal Rank Fusion (RRF).
RRF combines results from multiple queries without requiring relevance scores.
It assigns higher weight to documents that appear early in multiple result lists.
The RRF score for a document at rank ``r`` is: ``1 / (k + r + 1)``,
where ``k`` is the rank constant.
Args:
query (Optional[str], optional): Ignored by RRF. Defaults to None.
topn (int, optional): Number of top documents to return. Defaults to 10.
rerank_field (Optional[str], optional): Ignored by RRF. Defaults to None.
rank_constant (int, optional): Smoothing constant ``k`` in RRF formula.
Larger values reduce the impact of early ranks. Defaults to 60.
"""
def __init__(
self,
query: Optional[str] = None,
topn: int = 10,
rerank_field: Optional[str] = None,
rank_constant: int = 60,
):
super().__init__(query, topn, rerank_field)
self._rank_constant = rank_constant
@property
def rank_constant(self) -> int:
return self._rank_constant
def _rrf_score(self, rank: int):
return 1.0 / (self._rank_constant + rank + 1)
def rerank(self, query_results: dict[str, list[Doc]]) -> list[Doc]:
"""Apply Reciprocal Rank Fusion to combine multiple query results.
Args:
query_results (dict[str, list[Doc]]): Results from one or more vector queries.
Returns:
list[Doc]: Re-ranked documents with RRF scores in the ``score`` field.
"""
rrf_scores: dict[str, float] = defaultdict(float)
id_to_doc: dict[str, Doc] = {}
for _, query_result in query_results.items():
for rank, doc in enumerate(query_result):
doc_id = doc.id
rrf_score = self._rrf_score(rank)
rrf_scores[doc_id] += rrf_score
if doc_id not in id_to_doc:
id_to_doc[doc_id] = doc
top_docs = heapq.nlargest(self.topn, rrf_scores.items(), key=lambda x: x[1])
results = []
for doc_id, rrf_score in top_docs:
doc = id_to_doc[doc_id]
new_doc = doc._replace(score=rrf_score)
results.append(new_doc)
return results
class WeightedReRanker(ReRanker):
"""Re-ranker that combines scores from multiple vector fields using weights.
Each vector field's relevance score is normalized based on its metric type,
then scaled by a user-provided weight. Final scores are summed across fields.
Args:
query (Optional[str], optional): Ignored. Defaults to None.
topn (int, optional): Number of top documents to return. Defaults to 10.
rerank_field (Optional[str], optional): Ignored. Defaults to None.
metric (MetricType, optional): Distance metric used for score normalization.
Defaults to ``MetricType.L2``.
weights (Optional[dict[str, float]], optional): Weight per vector field.
Fields not listed use weight 1.0. Defaults to None.
Note:
Supported metrics: L2, IP, COSINE. Scores are normalized to [0, 1].
"""
def __init__(
self,
query: Optional[str] = None,
topn: int = 10,
rerank_field: Optional[str] = None,
metric: MetricType = MetricType.L2,
weights: Optional[dict[str, float]] = None,
):
super().__init__(query, topn, rerank_field)
self._weights = weights
self._metric = metric
@property
def weights(self) -> dict[str, float]:
"""dict[str, float]: Weight mapping for vector fields."""
return self._weights
@property
def metric(self) -> MetricType:
"""MetricType: Distance metric used for score normalization."""
return self._metric
def rerank(self, query_results: dict[str, list[Doc]]) -> list[Doc]:
"""Combine scores from multiple vector fields using weighted sum.
Args:
query_results (dict[str, list[Doc]]): Results per vector field.
Returns:
list[Doc]: Re-ranked documents with combined scores in ``score`` field.
"""
weighted_scores: dict[str, float] = defaultdict(float)
id_to_doc: dict[str, Doc] = {}
for vector_name, query_result in query_results.items():
for _, doc in enumerate(query_result):
doc_id = doc.id
weighted_score = self._normalize_score(
doc.score, self.metric
) * self.weights.get(vector_name, 1.0)
weighted_scores[doc_id] += weighted_score
if doc_id not in id_to_doc:
id_to_doc[doc_id] = doc
top_docs = heapq.nlargest(
self.topn, weighted_scores.items(), key=lambda x: x[1]
)
results = []
for doc_id, weighted_score in top_docs:
doc = id_to_doc[doc_id]
new_doc = doc._replace(score=weighted_score)
results.append(new_doc)
return results
def _normalize_score(self, score: float, metric: MetricType) -> float:
if metric == MetricType.L2:
return 1.0 - 2 * math.atan(score) / math.pi
if metric == MetricType.IP:
return 0.5 + math.atan(score) / math.pi
if metric == MetricType.COSINE:
return 1.0 - score / 2.0
raise ValueError("Unsupported metric type")
class QwenReRanker(ReRanker):
"""Re-ranker using Qwen (DashScope) LLM-based re-ranking API.
This re-ranker sends documents to the DashScope TextReRank service for
cross-encoder style re-ranking based on semantic relevance to the query.
Args:
query (str): Query text for semantic re-ranking. **Required**.
topn (int, optional): Number of top documents to return. Defaults to 10.
rerank_field (str): Field name containing document text for re-ranking.
**Required**.
model (str, optional): DashScope re-ranking model name.
Defaults to ``"gte-rerank-v2"``.
api_key (Optional[str], optional): DashScope API key. If not provided,
reads from ``DASHSCOPE_API_KEY`` environment variable.
Raises:
ValueError: If ``query`` is missing, ``rerank_field`` is missing,
or API key is not provided.
Note:
Requires the ``dashscope`` Python package.
Documents without content in ``rerank_field`` are skipped.
"""
def __init__(
self,
query: Optional[str] = None,
topn: int = 10,
rerank_field: Optional[str] = None,
model: str = "gte-rerank-v2",
api_key: Optional[str] = None,
):
super().__init__(query, topn, rerank_field)
if not query:
raise ValueError("Query is required for reranking")
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")
@property
def model(self) -> str:
"""str: DashScope re-ranking model name."""
return self._model
def _connection(self):
dashscope = require_module("dashscope")
dashscope.api_key = self._api_key
return dashscope
def rerank(self, query_results: dict[str, list[Doc]]) -> list[Doc]:
"""Re-rank documents using Qwen's TextReRank API.
Args:
query_results (dict[str, list[Doc]]): Results from vector search.
Returns:
list[Doc]: Re-ranked documents with relevance scores from Qwen.
Raises:
ValueError: If API call fails or no valid documents are found.
"""
if not query_results:
return []
id_to_doc: dict[str, Doc] = {}
doc_ids = []
contents = []
for _, query_result in query_results.items():
for doc in query_result:
doc_id = doc.id
if doc_id in id_to_doc:
continue
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")
resp = self._connection().TextReRank.call(
model=self.model,
query=self.query,
documents=list(contents),
top_n=self.topn,
return_documents=False,
)
if resp.status_code != HTTPStatus.OK:
raise ValueError(
f"QwenReranker failed with status {resp.status_code}: {resp.message}"
)
results = []
for item in resp.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

69
python/zvec/extension/rerank_function.py Normal file
View File

@ -0,0 +1,69 @@
# 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 Optional
from ..model.doc import Doc
class RerankFunction(ABC):
"""Abstract base class for re-ranking search results.
Re-rankers refine the output of one or more vector queries by applying
a secondary scoring strategy. They are used in the ``query()`` method of
``Collection`` via the ``reranker`` parameter.
Args:
topn (int, optional): Number of top documents to return after re-ranking.
Defaults to 10.
rerank_field (Optional[str], optional): Field name used as input for
re-ranking (e.g., document title or body). Defaults to None.
Note:
Subclasses must implement the ``rerank()`` method.
"""
def __init__(
self,
topn: int = 10,
rerank_field: Optional[str] = None,
):
self._topn = topn
self._rerank_field = rerank_field
@property
def topn(self) -> int:
"""int: Number of top documents to return after re-ranking."""
return self._topn
@property
def rerank_field(self) -> Optional[str]:
"""Optional[str]: Field name used as re-ranking input."""
return self._rerank_field
@abstractmethod
def rerank(self, query_results: dict[str, list[Doc]]) -> list[Doc]:
"""Re-rank documents from one or more vector queries.
Args:
query_results (dict[str, list[Doc]]): Mapping from vector field name
to list of retrieved documents (sorted by relevance).
Returns:
list[Doc]: Re-ranked list of documents (length ≤ ``topn``),
with updated ``score`` fields.
"""
...

839
python/zvec/extension/sentence_transformer_embedding_function.py Normal file
View File

@ -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

150
python/zvec/extension/sentence_transformer_function.py Normal file
View File

@ -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")

384
python/zvec/extension/sentence_transformer_rerank_function.py Normal file
View File

@ -0,0 +1,384 @@
# 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 ..model.doc import Doc
from ..tool import require_module
from .rerank_function import RerankFunction
from .sentence_transformer_function import SentenceTransformerFunctionBase
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**.
topn (int, optional): Maximum number of documents to return after re-ranking.
Defaults to 10.
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.
topn (int): Maximum number of documents to return.
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",
... topn=5,
... 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="深度学习",
... topn=10,
... rerank_field="content",
... model_source="modelscope"
... )
>>> # Using larger model for better quality
>>> reranker = SentenceTransformerReRanker(
... query="neural networks",
... topn=5,
... 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,
topn: int = 10,
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.
topn (int): Number of top results to return.
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 function
RerankFunction.__init__(self, topn=topn, 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 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: dict[str, list[Doc]]) -> list[Doc]:
"""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 (dict[str, list[Doc]]): Mapping from vector field names
to lists of retrieved documents. Documents from all fields are
deduplicated and re-ranked together.
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 []
# Collect and deduplicate documents
id_to_doc: dict[str, Doc] = {}
doc_ids: list[str] = []
contents: list[str] = []
for _, query_result in query_results.items():
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[: self.topn]
# Build result list with updated scores
results: list[Doc] = []
for _, doc, score in top_scored_docs:
new_doc = doc._replace(score=score)
results.append(new_doc)
return results

2
python/zvec/tool/util.py
View File

@ -59,5 +59,5 @@ def require_module(module: str, mitigation: Optional[str] = None) -> Any:
else:
msg += f"please pip install '{top_level}'."
else:
msg += f"Please pip install '{package}."
msg += f"Please pip install '{package}'."
raise ImportError(msg) from e