diff --git a/.github/workflows/linux_arm64_docker_ci.yml b/.github/workflows/linux_arm64_docker_ci.yml index 96a0f32..5e02a95 100644 --- a/.github/workflows/linux_arm64_docker_ci.yml +++ b/.github/workflows/linux_arm64_docker_ci.yml @@ -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 diff --git a/.github/workflows/linux_x64_docker_ci.yml b/.github/workflows/linux_x64_docker_ci.yml index 2edd099..2c5bb2f 100644 --- a/.github/workflows/linux_x64_docker_ci.yml +++ b/.github/workflows/linux_x64_docker_ci.yml @@ -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 diff --git a/pyproject.toml b/pyproject.toml index dee6728..de14714 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/python/tests/test_embedding.py b/python/tests/test_embedding.py index 0eb5d6b..e0a57a1 100644 --- a/python/tests/test_embedding.py +++ b/python/tests/test_embedding.py @@ -15,20 +15,31 @@ from __future__ import annotations import os from http import HTTPStatus -from unittest.mock import MagicMock, patch - +from unittest.mock import MagicMock, patch, Mock +import numpy as np import pytest -from zvec.extension import QwenEmbeddingFunction +from zvec.extension import ( + BM25EmbeddingFunction, + DefaultLocalDenseEmbedding, + DefaultLocalSparseEmbedding, + OpenAIDenseEmbedding, + QwenDenseEmbedding, + QwenSparseEmbedding, +) + +# Environment variable to control integration tests +# Set ZVEC_RUN_INTEGRATION_TESTS=1 to run real API/model tests +RUN_INTEGRATION_TESTS = os.environ.get("ZVEC_RUN_INTEGRATION_TESTS", "0") == "1" # ---------------------------- -# QwenEmbeddingFunction Test Case +# QwenDenseEmbedding Test Case # ---------------------------- -class TestQwenEmbeddingFunction: +class TestQwenDenseEmbedding: def test_init_with_api_key(self): # Test initialization with explicit API key - embedding_func = QwenEmbeddingFunction(dimension=128, api_key="test_key") + embedding_func = QwenDenseEmbedding(dimension=128, api_key="test_key") assert embedding_func.dimension == 128 assert embedding_func.model == "text-embedding-v4" assert embedding_func._api_key == "test_key" @@ -36,33 +47,28 @@ class TestQwenEmbeddingFunction: @patch.dict(os.environ, {"DASHSCOPE_API_KEY": "env_key"}) def test_init_with_env_api_key(self): # Test initialization with API key from environment - embedding_func = QwenEmbeddingFunction(dimension=128) + embedding_func = QwenDenseEmbedding(dimension=128) assert embedding_func._api_key == "env_key" - def test_init_without_api_key(self): - # Test initialization without API key raises ValueError - with pytest.raises(ValueError, match="DashScope API key is required"): - QwenEmbeddingFunction(dimension=128) - @patch.dict(os.environ, {"DASHSCOPE_API_KEY": ""}) def test_init_with_empty_env_api_key(self): # Test initialization with empty API key from environment with pytest.raises(ValueError, match="DashScope API key is required"): - QwenEmbeddingFunction(dimension=128) + QwenDenseEmbedding(dimension=128) def test_model_property(self): - embedding_func = QwenEmbeddingFunction(dimension=128, api_key="test_key") + embedding_func = QwenDenseEmbedding(dimension=128, api_key="test_key") assert embedding_func.model == "text-embedding-v4" - embedding_func = QwenEmbeddingFunction( + embedding_func = QwenDenseEmbedding( dimension=128, model="custom-model", api_key="test_key" ) assert embedding_func.model == "custom-model" - @patch("zvec.extension.embedding.require_module") + @patch("zvec.extension.qwen_function.require_module") def test_embed_with_empty_text(self, mock_require_module): # Test embed method with empty text raises ValueError - embedding_func = QwenEmbeddingFunction(dimension=128, api_key="test_key") + embedding_func = QwenDenseEmbedding(dimension=128, api_key="test_key") with pytest.raises( ValueError, match="Input text cannot be empty or whitespace only" @@ -72,7 +78,7 @@ class TestQwenEmbeddingFunction: with pytest.raises(TypeError): embedding_func.embed(None) - @patch("zvec.extension.embedding.require_module") + @patch("zvec.extension.qwen_function.require_module") def test_embed_success(self, mock_require_module): # Test successful embedding mock_dashscope = MagicMock() @@ -82,18 +88,20 @@ class TestQwenEmbeddingFunction: mock_dashscope.TextEmbedding.call.return_value = mock_response mock_require_module.return_value = mock_dashscope - embedding_func = QwenEmbeddingFunction(dimension=128, api_key="test_key") + embedding_func = QwenDenseEmbedding(dimension=3, api_key="test_key") + # Clear cache to avoid interference + embedding_func.embed.cache_clear() result = embedding_func.embed("test text") assert result == [0.1, 0.2, 0.3] mock_dashscope.TextEmbedding.call.assert_called_once_with( model="text-embedding-v4", input="test text", - dimension=128, + dimension=3, output_type="dense", ) - @patch("zvec.extension.embedding.require_module") + @patch("zvec.extension.qwen_function.require_module") def test_embed_http_error(self, mock_require_module): # Test embedding with HTTP error mock_dashscope = MagicMock() @@ -103,29 +111,1989 @@ class TestQwenEmbeddingFunction: mock_dashscope.TextEmbedding.call.return_value = mock_response mock_require_module.return_value = mock_dashscope - embedding_func = QwenEmbeddingFunction(dimension=128, api_key="test_key") + embedding_func = QwenDenseEmbedding(dimension=128, api_key="test_key") + embedding_func.embed.cache_clear() with pytest.raises(ValueError): embedding_func.embed("test text") - @patch("zvec.extension.embedding.require_module") + @patch("zvec.extension.qwen_function.require_module") def test_embed_invalid_response(self, mock_require_module): # Test embedding with invalid response (wrong number of embeddings) mock_dashscope = MagicMock() mock_response = MagicMock() mock_response.status_code = HTTPStatus.OK - mock_response.output.embeddings = [] + mock_response.output = {"embeddings": []} mock_dashscope.TextEmbedding.call.return_value = mock_response mock_require_module.return_value = mock_dashscope - embedding_func = QwenEmbeddingFunction(dimension=128, api_key="test_key") + embedding_func = QwenDenseEmbedding(dimension=128, api_key="test_key") + embedding_func.embed.cache_clear() with pytest.raises(ValueError): embedding_func.embed("test text") - @pytest.mark.skip(reason="Qwen Embedding is not available in CI") - def test_embed(self): - # Test embedding with invalid dimension - embedding_func = QwenEmbeddingFunction(dimension=128, api_key="xxx") + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + def test_real_embed_success(self): + """Integration test with real DashScope API. + + To run this test, set environment variable: + export ZVEC_RUN_INTEGRATION_TESTS=1 + export DASHSCOPE_API_KEY=your-api-key + """ + embedding_func = QwenDenseEmbedding(dimension=128) dense = embedding_func("test text") assert len(dense) == 128 + + +# ---------------------------- +# QwenSparseEmbedding Test Case +# ---------------------------- +class TestQwenSparseEmbedding: + """Test suite for QwenSparseEmbedding (Qwen sparse embedding via DashScope API).""" + + def test_init_with_api_key(self): + """Test initialization with explicit API key.""" + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + assert embedding_func._dimension == 1024 + assert embedding_func.model == "text-embedding-v4" + assert embedding_func._api_key == "test_key" + # encoding_type defaults to "query" via extra_params + assert embedding_func.extra_params.get("encoding_type", "query") == "query" + + def test_init_with_custom_encoding_type(self): + """Test initialization with custom encoding type.""" + embedding_func = QwenSparseEmbedding( + dimension=1024, encoding_type="document", api_key="test_key" + ) + assert embedding_func.extra_params.get("encoding_type") == "document" + + @patch.dict(os.environ, {"DASHSCOPE_API_KEY": "env_key"}) + def test_init_with_env_api_key(self): + """Test initialization with API key from environment.""" + embedding_func = QwenSparseEmbedding(dimension=1024) + assert embedding_func._api_key == "env_key" + + @patch.dict(os.environ, {"DASHSCOPE_API_KEY": ""}) + def test_init_without_api_key(self): + """Test initialization fails without API key.""" + with pytest.raises(ValueError, match="DashScope API key is required"): + QwenSparseEmbedding(dimension=1024) + + def test_model_property(self): + """Test model property.""" + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + assert embedding_func.model == "text-embedding-v4" + + embedding_func = QwenSparseEmbedding( + dimension=1024, model="text-embedding-v3", api_key="test_key" + ) + assert embedding_func.model == "text-embedding-v3" + + def test_encoding_type_property(self): + """Test encoding_type via extra_params.""" + query_emb = QwenSparseEmbedding( + dimension=1024, encoding_type="query", api_key="test_key" + ) + assert query_emb.extra_params.get("encoding_type") == "query" + + doc_emb = QwenSparseEmbedding( + dimension=1024, encoding_type="document", api_key="test_key" + ) + assert doc_emb.extra_params.get("encoding_type") == "document" + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_with_empty_text(self, mock_require_module): + """Test embed method with empty text raises ValueError.""" + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + + with pytest.raises( + ValueError, match="Input text cannot be empty or whitespace only" + ): + embedding_func.embed("") + + with pytest.raises( + ValueError, match="Input text cannot be empty or whitespace only" + ): + embedding_func.embed(" ") + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_with_non_string_input(self, mock_require_module): + """Test embed method with non-string input raises TypeError.""" + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + + with pytest.raises(TypeError, match="Expected 'input' to be str"): + embedding_func.embed(123) + + with pytest.raises(TypeError, match="Expected 'input' to be str"): + embedding_func.embed(None) + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_success(self, mock_require_module): + """Test successful sparse embedding generation.""" + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.OK + # Sparse embedding returns array of {index, value, token} objects + mock_response.output = { + "embeddings": [ + { + "sparse_embedding": [ + {"index": 10, "value": 0.5, "token": "机器"}, + {"index": 245, "value": 0.8, "token": "学习"}, + {"index": 1023, "value": 1.2, "token": "算法"}, + ] + } + ] + } + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + # Clear cache to avoid interference + embedding_func.embed.cache_clear() + result = embedding_func.embed("test text") + + # Verify result is a dict + assert isinstance(result, dict) + # Verify keys are integers + assert all(isinstance(k, int) for k in result.keys()) + # Verify values are floats + assert all(isinstance(v, float) for v in result.values()) + # Verify all values are positive + assert all(v > 0 for v in result.values()) + # Verify sorted by indices + keys = list(result.keys()) + assert keys == sorted(keys) + # Verify specific keys + assert keys == [10, 245, 1023] + + mock_dashscope.TextEmbedding.call.assert_called_once_with( + model="text-embedding-v4", + input="test text", + dimension=1024, + output_type="sparse", + text_type="query", + ) + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_with_document_encoding_type(self, mock_require_module): + """Test embedding with document encoding type.""" + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.OK + mock_response.output = { + "embeddings": [ + { + "sparse_embedding": [ + {"index": 5, "value": 0.3, "token": "文档"}, + {"index": 100, "value": 0.7, "token": "内容"}, + {"index": 500, "value": 0.9, "token": "检索"}, + ] + } + ] + } + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenSparseEmbedding( + dimension=1024, encoding_type="document", api_key="test_key" + ) + embedding_func.embed.cache_clear() + result = embedding_func.embed("test document") + + assert isinstance(result, dict) + assert list(result.keys()) == [5, 100, 500] + + # Verify text_type parameter is "document" + call_args = mock_dashscope.TextEmbedding.call.call_args + assert call_args[1]["text_type"] == "document" + assert call_args[1]["output_type"] == "sparse" + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_output_sorted_by_indices(self, mock_require_module): + """Test that output is always sorted by indices in ascending order.""" + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.OK + # Return unsorted indices + mock_response.output = { + "embeddings": [ + { + "sparse_embedding": [ + {"index": 9999, "value": 1.5, "token": "A"}, + {"index": 5, "value": 2.0, "token": "B"}, + {"index": 1234, "value": 0.8, "token": "C"}, + {"index": 77, "value": 3.2, "token": "D"}, + {"index": 500, "value": 1.1, "token": "E"}, + ] + } + ] + } + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + embedding_func.embed.cache_clear() + result = embedding_func.embed("test sorting") + + # Verify keys are sorted + result_keys = list(result.keys()) + assert result_keys == sorted(result_keys) + # Verify expected sorted order + assert result_keys == [5, 77, 500, 1234, 9999] + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_filters_zero_values(self, mock_require_module): + """Test that zero and negative values are filtered out.""" + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.OK + # Include zero and negative values + mock_response.output = { + "embeddings": [ + { + "sparse_embedding": [ + {"index": 10, "value": 0.5, "token": "正"}, + { + "index": 20, + "value": 0.0, + "token": "零", + }, # Should be filtered + { + "index": 30, + "value": -0.3, + "token": "负", + }, # Should be filtered + {"index": 40, "value": 0.8, "token": "正"}, + { + "index": 50, + "value": 0.0, + "token": "零", + }, # Should be filtered + ] + } + ] + } + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + embedding_func.embed.cache_clear() + result = embedding_func.embed("test filtering") + + # Only positive values should remain + assert list(result.keys()) == [10, 40] + assert all(v > 0 for v in result.values()) + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_http_error(self, mock_require_module): + """Test embedding with HTTP error.""" + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.BAD_REQUEST + mock_response.message = "Bad Request" + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + embedding_func.embed.cache_clear() + + with pytest.raises(ValueError, match="DashScope API error"): + embedding_func.embed("test text") + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_invalid_response_no_embeddings(self, mock_require_module): + """Test embedding with invalid response (no embeddings).""" + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.OK + mock_response.output = {"embeddings": []} + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + embedding_func.embed.cache_clear() + + with pytest.raises(ValueError, match="Expected exactly 1 embedding"): + embedding_func.embed("test text") + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_invalid_response_not_dict(self, mock_require_module): + """Test embedding with invalid response (sparse_embedding not list).""" + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.OK + # sparse_embedding should be list, not dict + mock_response.output = { + "embeddings": [{"sparse_embedding": {"index": 10, "value": 0.5}}] + } + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + embedding_func.embed.cache_clear() + + with pytest.raises( + ValueError, match="'sparse_embedding' field is missing or not a list" + ): + embedding_func.embed("test text") + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_callable_interface(self, mock_require_module): + """Test that embedding function is callable.""" + mock_dashscope = MagicMock() + mock_response = MagicMock() + mock_response.status_code = HTTPStatus.OK + mock_response.output = { + "embeddings": [ + { + "sparse_embedding": [ + {"index": 100, "value": 1.0, "token": "测试"}, + {"index": 200, "value": 0.5, "token": "调用"}, + ] + } + ] + } + mock_dashscope.TextEmbedding.call.return_value = mock_response + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + embedding_func.embed.cache_clear() + + # Test calling the function directly + result = embedding_func("test text") + assert isinstance(result, dict) + assert list(result.keys()) == [100, 200] + + @patch("zvec.extension.qwen_function.require_module") + def test_embed_api_connection_error(self, mock_require_module): + """Test handling of API connection errors.""" + mock_dashscope = MagicMock() + mock_dashscope.TextEmbedding.call.side_effect = Exception("Connection timeout") + mock_require_module.return_value = mock_dashscope + + embedding_func = QwenSparseEmbedding(dimension=1024, api_key="test_key") + embedding_func.embed.cache_clear() + + with pytest.raises(RuntimeError, match="Failed to call DashScope API"): + embedding_func.embed("test text") + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + def test_real_embed_success(self): + """Integration test with real DashScope API. + + To run this test, set environment variable: + export ZVEC_RUN_INTEGRATION_TESTS=1 + export DASHSCOPE_API_KEY=your-api-key + """ + # Test query embedding + query_emb = QwenSparseEmbedding(dimension=1024, encoding_type="query") + query_vec = query_emb.embed("machine learning") + + assert isinstance(query_vec, dict) + assert len(query_vec) > 0 + assert all(isinstance(k, int) for k in query_vec.keys()) + assert all(isinstance(v, float) and v > 0 for v in query_vec.values()) + + # Verify sorted output + keys = list(query_vec.keys()) + assert keys == sorted(keys) + + # Test document embedding + doc_emb = QwenSparseEmbedding(dimension=1024, encoding_type="document") + doc_vec = doc_emb.embed("Machine learning is a subset of AI") + + assert isinstance(doc_vec, dict) + assert len(doc_vec) > 0 + + # Verify sorted output + doc_keys = list(doc_vec.keys()) + assert doc_keys == sorted(doc_keys) + + +# ---------------------------- +# OpenAIDenseEmbedding Test Case +# ---------------------------- +class TestOpenAIDenseEmbedding: + def test_init_with_api_key(self): + """Test initialization with explicit API key.""" + embedding_func = OpenAIDenseEmbedding(api_key="sk-test-key") + assert embedding_func.dimension == 1536 # Default for text-embedding-3-small + assert embedding_func.model == "text-embedding-3-small" + assert embedding_func._api_key == "sk-test-key" + + @patch.dict(os.environ, {"OPENAI_API_KEY": "sk-env-key"}) + def test_init_with_env_api_key(self): + """Test initialization with API key from environment.""" + embedding_func = OpenAIDenseEmbedding() + assert embedding_func._api_key == "sk-env-key" + + @patch.dict(os.environ, {"OPENAI_API_KEY": ""}) + def test_init_without_api_key(self): + """Test initialization fails without API key.""" + with pytest.raises(ValueError, match="OpenAI API key is required"): + OpenAIDenseEmbedding() + + def test_init_with_custom_dimension(self): + """Test initialization with custom dimension.""" + embedding_func = OpenAIDenseEmbedding( + model="text-embedding-3-large", dimension=1024, api_key="sk-test" + ) + assert embedding_func.dimension == 1024 + assert embedding_func.model == "text-embedding-3-large" + + def test_init_with_base_url(self): + """Test initialization with custom base URL.""" + embedding_func = OpenAIDenseEmbedding( + api_key="sk-test", base_url="https://custom.openai.com/" + ) + assert embedding_func._base_url == "https://custom.openai.com/" + + def test_model_property(self): + """Test model property.""" + embedding_func = OpenAIDenseEmbedding(api_key="sk-test") + assert embedding_func.model == "text-embedding-3-small" + + embedding_func = OpenAIDenseEmbedding( + model="text-embedding-ada-002", api_key="sk-test" + ) + assert embedding_func.model == "text-embedding-ada-002" + + def test_extra_params(self): + """Test extra_params property.""" + # Test without extra params + embedding_func = OpenAIDenseEmbedding(api_key="sk-test") + assert embedding_func.extra_params == {} + + # Test with extra params + embedding_func = OpenAIDenseEmbedding( + api_key="sk-test", + encoding_format="float", + user="test-user", + ) + assert embedding_func.extra_params == { + "encoding_format": "float", + "user": "test-user", + } + + @patch("zvec.extension.openai_function.require_module") + def test_embed_with_empty_text(self, mock_require_module): + """Test embed method with empty text raises ValueError.""" + embedding_func = OpenAIDenseEmbedding(api_key="sk-test") + + with pytest.raises( + ValueError, match="Input text cannot be empty or whitespace only" + ): + embedding_func.embed("") + + with pytest.raises( + ValueError, match="Input text cannot be empty or whitespace only" + ): + embedding_func.embed(" ") + + @patch("zvec.extension.openai_function.require_module") + def test_embed_with_non_string_input(self, mock_require_module): + """Test embed method with non-string input raises TypeError.""" + embedding_func = OpenAIDenseEmbedding(api_key="sk-test") + + with pytest.raises(TypeError, match="Expected 'input' to be str"): + embedding_func.embed(123) + + with pytest.raises(TypeError, match="Expected 'input' to be str"): + embedding_func.embed(None) + + @patch("zvec.extension.openai_function.require_module") + def test_embed_success(self, mock_require_module): + """Test successful embedding generation.""" + # Mock OpenAI client + mock_openai = Mock() + mock_client = Mock() + mock_response = Mock() + + # Create mock embedding data + fake_embedding = [0.1, 0.2, 0.3] + mock_embedding_obj = Mock() + mock_embedding_obj.embedding = fake_embedding + mock_response.data = [mock_embedding_obj] + + mock_client.embeddings.create.return_value = mock_response + mock_openai.OpenAI.return_value = mock_client + mock_require_module.return_value = mock_openai + + embedding_func = OpenAIDenseEmbedding(dimension=3, api_key="sk-test") + embedding_func.embed.cache_clear() + result = embedding_func.embed("test text") + + assert result == [0.1, 0.2, 0.3] + mock_client.embeddings.create.assert_called_once_with( + model="text-embedding-3-small", input="test text", dimensions=3 + ) + + @patch("zvec.extension.openai_function.require_module") + def test_embed_with_custom_model(self, mock_require_module): + """Test embedding with custom model.""" + mock_openai = Mock() + mock_client = Mock() + mock_response = Mock() + + fake_embedding = [0.1] * 1536 + mock_embedding_obj = Mock() + mock_embedding_obj.embedding = fake_embedding + mock_response.data = [mock_embedding_obj] + + mock_client.embeddings.create.return_value = mock_response + mock_openai.OpenAI.return_value = mock_client + mock_require_module.return_value = mock_openai + + embedding_func = OpenAIDenseEmbedding( + model="text-embedding-ada-002", api_key="sk-test" + ) + embedding_func.embed.cache_clear() + result = embedding_func.embed("test text") + + assert len(result) == 1536 + mock_client.embeddings.create.assert_called_once_with( + model="text-embedding-ada-002", input="test text" + ) + + @patch("zvec.extension.openai_function.require_module") + def test_embed_api_error(self, mock_require_module): + """Test handling of API errors.""" + mock_openai = Mock() + mock_client = Mock() + + # Simulate API error + api_error = Mock() + api_error.__class__.__name__ = "APIError" + mock_openai.APIError = type("APIError", (Exception,), {}) + mock_openai.APIConnectionError = type("APIConnectionError", (Exception,), {}) + + mock_client.embeddings.create.side_effect = mock_openai.APIError( + "Rate limit exceeded" + ) + mock_openai.OpenAI.return_value = mock_client + mock_require_module.return_value = mock_openai + + embedding_func = OpenAIDenseEmbedding(api_key="sk-test") + embedding_func.embed.cache_clear() + + with pytest.raises(RuntimeError, match="Failed to call OpenAI API"): + embedding_func.embed("test text") + + @patch("zvec.extension.openai_function.require_module") + def test_embed_invalid_response(self, mock_require_module): + """Test handling of invalid API response.""" + mock_openai = Mock() + mock_client = Mock() + mock_response = Mock() + + # Empty response data + mock_response.data = [] + + mock_client.embeddings.create.return_value = mock_response + mock_openai.OpenAI.return_value = mock_client + mock_openai.APIError = type("APIError", (Exception,), {}) + mock_openai.APIConnectionError = type("APIConnectionError", (Exception,), {}) + mock_require_module.return_value = mock_openai + + embedding_func = OpenAIDenseEmbedding(api_key="sk-test") + embedding_func.embed.cache_clear() + + with pytest.raises(ValueError, match="no embedding data returned"): + embedding_func.embed("test text") + + @patch("zvec.extension.openai_function.require_module") + def test_embed_dimension_mismatch(self, mock_require_module): + """Test handling of dimension mismatch.""" + mock_openai = Mock() + mock_client = Mock() + mock_response = Mock() + + # Return embedding with wrong dimension + fake_embedding = [0.1] * 512 + mock_embedding_obj = Mock() + mock_embedding_obj.embedding = fake_embedding + mock_response.data = [mock_embedding_obj] + + mock_client.embeddings.create.return_value = mock_response + mock_openai.OpenAI.return_value = mock_client + mock_openai.APIError = type("APIError", (Exception,), {}) + mock_openai.APIConnectionError = type("APIConnectionError", (Exception,), {}) + mock_require_module.return_value = mock_openai + + embedding_func = OpenAIDenseEmbedding(dimension=1536, api_key="sk-test") + embedding_func.embed.cache_clear() + + with pytest.raises(ValueError, match="Dimension mismatch"): + embedding_func.embed("test text") + + @patch("zvec.extension.openai_function.require_module") + def test_embed_callable(self, mock_require_module): + """Test that embedding function is callable.""" + mock_openai = Mock() + mock_client = Mock() + mock_response = Mock() + + fake_embedding = [0.1] * 1536 + mock_embedding_obj = Mock() + mock_embedding_obj.embedding = fake_embedding + mock_response.data = [mock_embedding_obj] + + mock_client.embeddings.create.return_value = mock_response + mock_openai.OpenAI.return_value = mock_client + mock_openai.APIError = type("APIError", (Exception,), {}) + mock_openai.APIConnectionError = type("APIConnectionError", (Exception,), {}) + mock_require_module.return_value = mock_openai + + embedding_func = OpenAIDenseEmbedding(api_key="sk-test") + embedding_func.embed.cache_clear() + + # Test calling the function directly + result = embedding_func("test text") + assert isinstance(result, list) + assert len(result) == 1536 + + @patch("zvec.extension.openai_function.require_module") + def test_embed_with_base_url(self, mock_require_module): + """Test embedding with custom base URL.""" + mock_openai = Mock() + mock_client = Mock() + mock_response = Mock() + + fake_embedding = [0.1] * 1536 + mock_embedding_obj = Mock() + mock_embedding_obj.embedding = fake_embedding + mock_response.data = [mock_embedding_obj] + + mock_client.embeddings.create.return_value = mock_response + mock_openai.OpenAI.return_value = mock_client + mock_openai.APIError = type("APIError", (Exception,), {}) + mock_openai.APIConnectionError = type("APIConnectionError", (Exception,), {}) + mock_require_module.return_value = mock_openai + + embedding_func = OpenAIDenseEmbedding( + api_key="sk-test", base_url="https://custom.openai.com/" + ) + embedding_func.embed.cache_clear() + result = embedding_func.embed("test text") + + # Verify client was created with custom base URL + mock_openai.OpenAI.assert_called_once_with( + api_key="sk-test", base_url="https://custom.openai.com/" + ) + assert len(result) == 1536 + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + def test_real_embed_success(self): + """Integration test with real OpenAI API. + + To run this test, set environment variable: + export ZVEC_RUN_INTEGRATION_TESTS=1 + export OPENAI_API_KEY=sk-... + """ + embedding_func = OpenAIDenseEmbedding( + model="text-embedding-v4", + dimension=256, + base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", + ) + vector = embedding_func.embed("Hello, world!") + assert len(vector) == 256 + assert isinstance(vector, list) + assert all(isinstance(x, float) for x in vector) + + +# ---------------------------- +# DefaultLocalDenseEmbedding Test Case +# ---------------------------- +class TestDefaultLocalDenseEmbedding: + """Test cases for DefaultLocalDenseEmbedding.""" + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_init_success(self, mock_require_module): + """Test successful initialization with mocked model.""" + # Mock sentence_transformers module + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + mock_model.device = "cpu" + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + # Initialize embedding function + emb_func = DefaultLocalDenseEmbedding() + + # Assertions + assert emb_func.dimension == 384 + assert emb_func.model_name == "all-MiniLM-L6-v2" + assert emb_func.model_source == "huggingface" + assert emb_func.device == "cpu" + mock_st.SentenceTransformer.assert_called_once_with( + "all-MiniLM-L6-v2", device=None, trust_remote_code=True + ) + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_init_with_custom_device(self, mock_require_module): + """Test initialization with custom device.""" + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + mock_model.device = "cuda" + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + emb_func = DefaultLocalDenseEmbedding(device="cuda") + + assert emb_func.device == "cuda" + mock_st.SentenceTransformer.assert_called_once_with( + "all-MiniLM-L6-v2", device="cuda", trust_remote_code=True + ) + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_init_with_modelscope(self, mock_require_module): + """Test initialization with ModelScope as model source.""" + mock_st = Mock() + mock_ms = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + mock_model.device = "cpu" + mock_st.SentenceTransformer.return_value = mock_model + + def require_module_side_effect(module_name): + if module_name == "sentence_transformers": + return mock_st + elif module_name == "modelscope": + return mock_ms + raise ImportError(f"No module named '{module_name}'") + + mock_require_module.side_effect = require_module_side_effect + + # Mock snapshot_download at the correct import location + with patch( + "modelscope.hub.snapshot_download.snapshot_download", + return_value="/path/to/cached/model", + ): + emb_func = DefaultLocalDenseEmbedding(model_source="modelscope") + + # Assertions + assert emb_func.dimension == 384 + assert emb_func.model_name == "iic/nlp_gte_sentence-embedding_chinese-small" + assert emb_func.model_source == "modelscope" + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_init_with_invalid_model_source(self, mock_require_module): + """Test initialization with invalid model_source raises ValueError.""" + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + with pytest.raises(ValueError, match="Invalid model_source"): + DefaultLocalDenseEmbedding(model_source="invalid_source") + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_embed_success(self, mock_require_module): + """Test successful embedding generation.""" + # Mock embedding output + fake_embedding = np.random.rand(384).astype(np.float32) + + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + + # Configure encode method + mock_model.encode = Mock(return_value=fake_embedding) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + emb_func = DefaultLocalDenseEmbedding() + result = emb_func.embed("Hello, world!") + + # Assertions + assert isinstance(result, list) + assert len(result) == 384 + assert all(isinstance(x, float) for x in result) + mock_model.encode.assert_called_once_with( + "Hello, world!", + convert_to_numpy=True, + normalize_embeddings=True, + batch_size=32, + ) + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_embed_with_normalization(self, mock_require_module): + """Test embedding with L2 normalization.""" + # Create a normalized vector + fake_embedding = np.random.rand(384).astype(np.float32) + fake_embedding = fake_embedding / np.linalg.norm(fake_embedding) + + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + + # Configure encode method + mock_model.encode = Mock(return_value=fake_embedding) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + emb_func = DefaultLocalDenseEmbedding(normalize_embeddings=True) + result = emb_func.embed("Test sentence") + + # Check if vector is normalized (L2 norm should be close to 1.0) + result_array = np.array(result) + norm = np.linalg.norm(result_array) + assert abs(norm - 1.0) < 1e-5 + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_embed_empty_string(self, mock_require_module): + """Test embedding with empty string raises ValueError.""" + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + emb_func = DefaultLocalDenseEmbedding() + + with pytest.raises(ValueError, match="Input text cannot be empty"): + emb_func.embed("") + + with pytest.raises(ValueError, match="Input text cannot be empty"): + emb_func.embed(" ") + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_embed_non_string_input(self, mock_require_module): + """Test embedding with non-string input raises TypeError.""" + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + emb_func = DefaultLocalDenseEmbedding() + + with pytest.raises(TypeError, match="Expected 'input' to be str"): + emb_func.embed(123) + + with pytest.raises(TypeError, match="Expected 'input' to be str"): + emb_func.embed(None) + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_embed_callable(self, mock_require_module): + """Test that embedding function is callable.""" + fake_embedding = np.random.rand(384).astype(np.float32) + + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + + # Configure encode method + mock_model.encode = Mock(return_value=fake_embedding) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + emb_func = DefaultLocalDenseEmbedding() + + # Test calling the function directly + result = emb_func("Test text") + assert isinstance(result, list) + assert len(result) == 384 + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_semantic_similarity(self, mock_require_module): + """Test semantic similarity between similar and different texts.""" + # Create mock embeddings for similar and different texts + similar_emb_1 = np.array([1.0, 0.0, 0.0] + [0.0] * 381, dtype=np.float32) + similar_emb_2 = np.array([0.9, 0.1, 0.0] + [0.0] * 381, dtype=np.float32) + different_emb = np.array([0.0, 0.0, 1.0] + [0.0] * 381, dtype=np.float32) + + # Normalize + similar_emb_1 = similar_emb_1 / np.linalg.norm(similar_emb_1) + similar_emb_2 = similar_emb_2 / np.linalg.norm(similar_emb_2) + different_emb = different_emb / np.linalg.norm(different_emb) + + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + + # Configure encode method with side_effect for multiple calls + mock_model.encode = Mock( + side_effect=[similar_emb_1, similar_emb_2, different_emb] + ) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + emb_func = DefaultLocalDenseEmbedding() + + 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") + + # Calculate similarities + similarity_high = np.dot(v1, v2) + similarity_low = np.dot(v1, v3) + + assert similarity_high > similarity_low + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_model_loading_error(self, mock_require_module): + """Test handling of model loading failure.""" + # Clear model cache + from zvec.extension.sentence_transformer_embedding_function import ( + DefaultLocalSparseEmbedding, + ) + + DefaultLocalSparseEmbedding.clear_cache() + mock_st = Mock() + mock_st.SentenceTransformer.side_effect = Exception("Model not found") + mock_require_module.return_value = mock_st + + with pytest.raises( + ValueError, match="Failed to load Sentence Transformer model" + ): + DefaultLocalDenseEmbedding() + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_modelscope_import_error(self, mock_require_module): + """Test handling of ModelScope import error.""" + mock_st = Mock() + + def require_module_side_effect(module_name): + if module_name == "sentence_transformers": + return mock_st + elif module_name == "modelscope": + raise ImportError("No module named 'modelscope'") + + mock_require_module.side_effect = require_module_side_effect + + with pytest.raises( + ImportError, match="ModelScope support requires the 'modelscope' package" + ): + DefaultLocalDenseEmbedding(model_source="modelscope") + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_embed_dimension_mismatch(self, mock_require_module): + """Test handling of dimension mismatch in embedding output.""" + # Return embedding with wrong dimension + fake_embedding = np.random.rand(256).astype(np.float32) + + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + + # Configure encode method + mock_model.encode = Mock(return_value=fake_embedding) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + emb_func = DefaultLocalDenseEmbedding() + + with pytest.raises(ValueError, match="Dimension mismatch"): + emb_func.embed("Test text") + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + def test_real_embedding_generation(self): + """Integration test with real model (requires sentence-transformers). + + To run this test, set environment variable: + export ZVEC_RUN_INTEGRATION_TESTS=1 + + Note: First run will download the model (~80MB). + """ + emb_func = DefaultLocalDenseEmbedding() + + # Test basic embedding + vector = emb_func.embed("Hello, world!") + assert len(vector) == 384 + assert isinstance(vector, list) + assert all(isinstance(x, float) for x in vector) + + # Test normalization + norm = np.linalg.norm(vector) + assert abs(norm - 1.0) < 1e-5 + + # Test 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 language") + + similarity_high = np.dot(v1, v2) + similarity_low = np.dot(v1, v3) + assert similarity_high > similarity_low + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_model_properties(self, mock_require_module): + """Test model_name and model_source properties.""" + mock_st = Mock() + mock_model = Mock() + mock_model.get_sentence_embedding_dimension.return_value = 384 + mock_model.device = "cpu" + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + # Test Hugging Face + emb_func_hf = DefaultLocalDenseEmbedding(model_source="huggingface") + assert emb_func_hf.model_name == "all-MiniLM-L6-v2" + assert emb_func_hf.model_source == "huggingface" + + # Test ModelScope + with patch( + "modelscope.hub.snapshot_download.snapshot_download", + return_value="/path/to/model", + ): + mock_ms = Mock() + mock_require_module.side_effect = ( + lambda m: mock_st if m == "sentence_transformers" else mock_ms + ) + emb_func_ms = DefaultLocalDenseEmbedding(model_source="modelscope") + assert ( + emb_func_ms.model_name == "iic/nlp_gte_sentence-embedding_chinese-small" + ) + assert emb_func_ms.model_source == "modelscope" + + +# ----------------------------------- +# DefaultLocalSparseEmbedding Test Case +# ----------------------------------- +class TestDefaultLocalSparseEmbedding: + """Test suite for DefaultLocalSparseEmbedding (SPLADE sparse embedding). + + Note: + DefaultLocalSparseEmbedding uses naver/splade-cocondenser-ensembledistil + instead of naver/splade-v3 because: + + - splade-v3 is a gated model requiring Hugging Face authentication + - cocondenser-ensembledistil is publicly accessible + - Performance difference is minimal (~2%) + - Avoids "Access to model is restricted" errors + + This allows all users to run tests without authentication setup. + """ + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_init_success(self, mock_require_module): + """Test successful initialization. + + Verifies that DefaultLocalSparseEmbedding initializes with the publicly + accessible naver/splade-cocondenser-ensembledistil model instead of + the gated naver/splade-v3 model. + """ + mock_st = Mock() + mock_model = Mock() + mock_model.device = "cpu" + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding() + + assert sparse_emb.model_name == "naver/splade-cocondenser-ensembledistil" + assert sparse_emb.model_source == "huggingface" + assert sparse_emb.device == "cpu" + mock_st.SentenceTransformer.assert_called_once_with( + "naver/splade-cocondenser-ensembledistil", + device=None, + trust_remote_code=True, + ) + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_init_with_custom_device(self, mock_require_module): + """Test initialization with custom device.""" + mock_st = Mock() + mock_model = Mock() + mock_model.device = "cuda" + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding(device="cuda") + + assert sparse_emb.device == "cuda" + mock_st.SentenceTransformer.assert_called_once_with( + "naver/splade-cocondenser-ensembledistil", + device="cuda", + trust_remote_code=True, + ) + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_embed_success(self, mock_require_module): + """Test successful sparse embedding generation with official API.""" + import numpy as np + + # Clear model cache to ensure fresh mock + from zvec.extension.sentence_transformer_embedding_function import ( + DefaultLocalSparseEmbedding, + ) + + DefaultLocalSparseEmbedding.clear_cache() + + # Create a mock sparse matrix that simulates scipy.sparse behavior + # The code will call: sparse_matrix[0].toarray().flatten() + mock_sparse_matrix = Mock() + + # Create a dense array representation with vocab_size=30522 + vocab_size = 30522 + dense_array = np.zeros(vocab_size) + # Set specific non-zero values at indices [10, 245, 1023, 5678] + dense_array[10] = 0.5 + dense_array[245] = 0.8 + dense_array[1023] = 1.2 + dense_array[5678] = 0.3 + + # Mock the method chain: sparse_matrix[0].toarray().flatten() + mock_row = Mock() + mock_dense = Mock() + mock_row.toarray.return_value = mock_dense + mock_dense.flatten.return_value = dense_array + mock_sparse_matrix.__getitem__ = Mock(return_value=mock_row) + + # Also mock hasattr check for 'toarray' + mock_sparse_matrix.toarray = Mock() + + mock_st = Mock() + mock_model = Mock() + mock_model.device = "cpu" + + # Configure mock methods to return sparse matrix + # Must set return_value BEFORE hasattr() check in the code + mock_model.encode_query = Mock(return_value=mock_sparse_matrix) + mock_model.encode_document = Mock(return_value=mock_sparse_matrix) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding() + result = sparse_emb.embed("machine learning") + + # Verify result is a dictionary + assert isinstance(result, dict) + # Verify keys are integers and values are floats + assert all(isinstance(k, int) for k in result.keys()) + assert all(isinstance(v, float) for v in result.values()) + # Verify all values are positive + assert all(v > 0 for v in result.values()) + # Sparse vectors should have specific dimensions + assert len(result) == 4 + + # Verify output is sorted by indices (keys) + keys = list(result.keys()) + assert keys == sorted(keys), ( + "Sparse vector keys must be sorted in ascending order" + ) + + # Verify expected keys + assert keys == [10, 245, 1023, 5678] + + # Verify encode_query was called with a list + mock_model.encode_query.assert_called_once() + call_args = mock_model.encode_query.call_args[0][0] + assert isinstance(call_args, list) + assert call_args == ["machine learning"] + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_embed_empty_input(self, mock_require_module): + """Test embedding with empty input.""" + mock_st = Mock() + mock_model = Mock() + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding() + + with pytest.raises(ValueError, match="Input text cannot be empty"): + sparse_emb.embed("") + + with pytest.raises(ValueError, match="Input text cannot be empty"): + sparse_emb.embed(" ") + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_embed_non_string_input(self, mock_require_module): + """Test embedding with non-string input.""" + mock_st = Mock() + mock_model = Mock() + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding() + + with pytest.raises(TypeError, match="Expected 'input' to be str"): + sparse_emb.embed(123) + + with pytest.raises(TypeError, match="Expected 'input' to be str"): + sparse_emb.embed(["text"]) + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_callable_interface(self, mock_require_module): + """Test that DefaultSparseEmbedding is callable.""" + import numpy as np + + # Clear model cache + from zvec.extension.sentence_transformer_embedding_function import ( + DefaultLocalSparseEmbedding, + ) + + DefaultLocalSparseEmbedding.clear_cache() + + # Create a mock sparse matrix + mock_sparse_matrix = Mock() + + # Create a dense array representation with vocab_size=30522 + vocab_size = 30522 + dense_array = np.zeros(vocab_size) + # Set specific non-zero values at indices [100, 200, 300] + dense_array[100] = 1.0 + dense_array[200] = 0.5 + dense_array[300] = 0.8 + + # Mock the method chain: sparse_matrix[0].toarray().flatten() + mock_row = Mock() + mock_dense = Mock() + mock_row.toarray.return_value = mock_dense + mock_dense.flatten.return_value = dense_array + mock_sparse_matrix.__getitem__ = Mock(return_value=mock_row) + + # Also mock hasattr check for 'toarray' + mock_sparse_matrix.toarray = Mock() + + mock_st = Mock() + mock_model = Mock() + mock_model.device = "cpu" + + # Configure mock methods + mock_model.encode_query = Mock(return_value=mock_sparse_matrix) + mock_model.encode_document = Mock(return_value=mock_sparse_matrix) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding() + + # Test callable interface + result = sparse_emb("test input") + assert isinstance(result, dict) + assert all(isinstance(k, int) for k in result.keys()) + + # Verify sorted output + keys = list(result.keys()) + assert keys == sorted(keys), "Callable interface must also return sorted keys" + assert keys == [100, 200, 300] + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_model_loading_failure(self, mock_require_module): + """Test handling of model loading failure.""" + # Clear model cache to ensure the test actually tries to load the model + from zvec.extension.sentence_transformer_embedding_function import ( + DefaultLocalSparseEmbedding, + ) + + DefaultLocalSparseEmbedding.clear_cache() + + mock_st = Mock() + mock_st.SentenceTransformer.side_effect = Exception("Model not found") + mock_require_module.return_value = mock_st + + with pytest.raises( + ValueError, match="Failed to load Sentence Transformer model" + ): + DefaultLocalSparseEmbedding() + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_inference_failure(self, mock_require_module): + """Test handling of inference failure.""" + # Clear model cache + from zvec.extension.sentence_transformer_embedding_function import ( + DefaultLocalSparseEmbedding, + ) + + DefaultLocalSparseEmbedding.clear_cache() + + mock_st = Mock() + mock_model = Mock() + mock_model.device = "cpu" + + # Configure mock methods to raise RuntimeError + mock_model.encode_query = Mock(side_effect=RuntimeError("CUDA out of memory")) + mock_model.encode_document = Mock( + side_effect=RuntimeError("CUDA out of memory") + ) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding() + + with pytest.raises(RuntimeError, match="Failed to generate sparse embedding"): + sparse_emb.embed("test input") + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_sparse_vector_properties(self, mock_require_module): + """Test properties of sparse vectors (sparsity, non-zero values, sorted order).""" + import numpy as np + + # Clear model cache + from zvec.extension.sentence_transformer_embedding_function import ( + DefaultLocalSparseEmbedding, + ) + + DefaultLocalSparseEmbedding.clear_cache() + + # Create a mock sparse matrix that simulates scipy.sparse behavior + # The code will call: sparse_matrix[0].toarray().flatten() + mock_sparse_matrix = Mock() + + # Create a dense array representation with vocab_size=30522 + vocab_size = 30522 + dense_array = np.zeros(vocab_size) + # Set specific non-zero values at indices [50, 100, 200, 400, 500] + dense_array[50] = 3.0 + dense_array[100] = 2.0 + dense_array[200] = 1.5 + dense_array[400] = 2.5 + dense_array[500] = 1.8 + + # Mock the method chain: sparse_matrix[0].toarray().flatten() + mock_row = Mock() + mock_dense = Mock() + mock_row.toarray.return_value = mock_dense + mock_dense.flatten.return_value = dense_array + mock_sparse_matrix.__getitem__ = Mock(return_value=mock_row) + + # Also mock hasattr check for 'toarray' + mock_sparse_matrix.toarray = Mock() + + mock_st = Mock() + mock_model = Mock() + mock_model.device = "cpu" + + # Configure mock methods + mock_model.encode_query = Mock(return_value=mock_sparse_matrix) + mock_model.encode_document = Mock(return_value=mock_sparse_matrix) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding() + result = sparse_emb.embed("test") + + # Verify sparsity: result should have much fewer dimensions than vocab_size + assert len(result) < vocab_size + # All values should be positive + assert all(v > 0 for v in result.values()) + + # Verify keys are sorted in ascending order + keys = list(result.keys()) + assert keys == sorted(keys), "Sparse vector keys must be sorted" + + # Verify the specific non-zero indices are present and sorted + # Expected order: [50, 100, 200, 400, 500] (sorted) + expected_keys = [50, 100, 200, 400, 500] + assert keys == expected_keys, f"Expected {expected_keys}, got {keys}" + + # First key should be smallest + if len(result) > 0: + first_key = next(iter(result.keys())) + assert first_key == min(result.keys()), "First key must be the smallest" + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_output_sorted_by_indices(self, mock_require_module): + """Test that output dictionary is always sorted by indices (keys) in ascending order.""" + import numpy as np + + # Clear model cache + from zvec.extension.sentence_transformer_embedding_function import ( + DefaultLocalSparseEmbedding, + ) + + DefaultLocalSparseEmbedding.clear_cache() + + # Create sparse output with deliberately out-of-order indices + # Non-sequential indices: 9999, 5, 1234, 77, 500 + mock_sparse_matrix = Mock() + + # Create a dense array representation with vocab_size=30522 + vocab_size = 30522 + dense_array = np.zeros(vocab_size) + # Set specific non-zero values at out-of-order indices + dense_array[9999] = 1.5 + dense_array[5] = 2.0 + dense_array[1234] = 0.8 + dense_array[77] = 3.2 + dense_array[500] = 1.1 + + # Mock the method chain: sparse_matrix[0].toarray().flatten() + mock_row = Mock() + mock_dense = Mock() + mock_row.toarray.return_value = mock_dense + mock_dense.flatten.return_value = dense_array + mock_sparse_matrix.__getitem__ = Mock(return_value=mock_row) + + # Also mock hasattr check for 'toarray' + mock_sparse_matrix.toarray = Mock() + + mock_st = Mock() + mock_model = Mock() + mock_model.device = "cpu" + + # Configure mock methods + mock_model.encode_query = Mock(return_value=mock_sparse_matrix) + mock_model.encode_document = Mock(return_value=mock_sparse_matrix) + + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding() + result = sparse_emb.embed("test sorting") + + # Extract keys from result + result_keys = list(result.keys()) + + # Verify keys are sorted + assert result_keys == sorted(result_keys), ( + f"Keys must be sorted in ascending order. " + f"Got: {result_keys}, Expected: {sorted(result_keys)}" + ) + + # Verify expected keys are present and in correct order + # Expected sorted order: [5, 77, 500, 1234, 9999] + expected_sorted_keys = [5, 77, 500, 1234, 9999] + assert result_keys == expected_sorted_keys, ( + f"All expected keys should be present in sorted order. " + f"Expected: {expected_sorted_keys}, Got: {result_keys}" + ) + + # Verify first and last keys + assert result_keys[0] == 5, "First key must be minimum" + assert result_keys[-1] == 9999, "Last key must be maximum" + + # Verify iteration order matches sorted order + for i, (key, value) in enumerate(result.items()): + if i > 0: + prev_key = list(result.keys())[i - 1] + assert key > prev_key, ( + f"Key at position {i} must be greater than previous key" + ) + + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_device_property(self, mock_require_module): + """Test device property returns correct device.""" + mock_st = Mock() + mock_model = Mock() + mock_model.device = "cuda" + mock_st.SentenceTransformer.return_value = mock_model + mock_require_module.return_value = mock_st + + sparse_emb = DefaultLocalSparseEmbedding(device="cuda") + assert sparse_emb.device == "cuda" + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test: requires ZVEC_RUN_INTEGRATION_TESTS=1 and model download", + ) + @patch("zvec.extension.sentence_transformer_function.require_module") + def test_modelscope_source(self, mock_require_module): + """Test initialization with ModelScope source.""" + mock_st = Mock() + mock_ms = Mock() + mock_model = Mock() + mock_model.device = "cpu" + mock_st.SentenceTransformer.return_value = mock_model + + # Mock ModelScope snapshot_download + with patch( + "modelscope.hub.snapshot_download.snapshot_download", + return_value="/cache/splade-cocondenser", + ): + mock_require_module.side_effect = ( + lambda m: mock_st if m == "sentence_transformers" else mock_ms + ) + + sparse_emb = DefaultLocalSparseEmbedding(model_source="modelscope") + + assert sparse_emb.model_name == "naver/splade-cocondenser-ensembledistil" + assert sparse_emb.model_source == "modelscope" + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test: requires ZVEC_RUN_INTEGRATION_TESTS=1 and model download", + ) + def test_integration_real_model(self): + """Integration test with real SPLADE model (requires model download). + + This test uses naver/splade-cocondenser-ensembledistil instead of + naver/splade-v3 because splade-v3 requires Hugging Face authentication. + The cocondenser-ensembledistil model is publicly accessible and provides + comparable performance. + + To run this test: + export ZVEC_RUN_INTEGRATION_TESTS=1 + pytest tests/test_embedding.py::TestDefaultSparseEmbedding::test_integration_real_model -v + + Note: First run will download ~100MB model from Hugging Face. + + Alternative models: + If you have access to splade-v3, you can create a custom embedding + class following the example in DefaultSparseEmbedding docstring. + """ + # Clear model cache to ensure fresh load + from zvec.extension.sentence_transformer_embedding_function import ( + DefaultLocalSparseEmbedding, + ) + + DefaultLocalSparseEmbedding.clear_cache() + + sparse_emb = DefaultLocalSparseEmbedding() + + # Test with real input + text = "machine learning and artificial intelligence" + result = sparse_emb.embed(text) + + # Verify result structure + assert isinstance(result, dict) + assert len(result) > 0 + assert all(isinstance(k, int) and k >= 0 for k in result.keys()) + assert all(isinstance(v, float) and v > 0 for v in result.values()) + + # SPLADE typically produces 100-300 non-zero dimensions + assert 50 < len(result) < 500 + + # Verify keys are sorted in ascending order + keys = list(result.keys()) + assert keys == sorted(keys), "Real model output must be sorted by indices" + + # Test callable interface + result2 = sparse_emb(text) + assert result == result2 + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test: requires ZVEC_RUN_INTEGRATION_TESTS=1", + ) + def test_integration_multiple_inputs(self): + """Integration test with multiple different inputs.""" + # Clear model cache + from zvec.extension.sentence_transformer_embedding_function import ( + DefaultLocalSparseEmbedding, + ) + + DefaultLocalSparseEmbedding.clear_cache() + + sparse_emb = DefaultLocalSparseEmbedding() + + texts = [ + "Hello, world!", + "Machine learning is fascinating", + "Python programming language", + ] + + results = [sparse_emb.embed(text) for text in texts] + + # All results should be different + assert len(results) == 3 + assert all(isinstance(r, dict) for r in results) + + # Different inputs should produce different sparse vectors + assert results[0] != results[1] + assert results[1] != results[2] + + # All results must be sorted by indices + for i, result in enumerate(results): + keys = list(result.keys()) + assert keys == sorted(keys), f"Result {i} must have sorted keys" + + +# ---------------------------- +# BM25EmbeddingFunction Test Case +# ---------------------------- +class TestBM25EmbeddingFunction: + """Test suite for BM25EmbeddingFunction (BM25-based sparse embedding using DashText SDK).""" + + def test_init_with_built_in_encoder(self): + """Test successful initialization with built-in encoder (no corpus).""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + # Test with default language (Chinese) + bm25 = BM25EmbeddingFunction() + + assert bm25.corpus_size == 0 + assert bm25.encoding_type == "query" + assert bm25.language == "zh" + mock_dashtext.SparseVectorEncoder.default.assert_called_once_with(name="zh") + + def test_init_with_custom_encoder(self): + """Test successful initialization with custom encoder (with corpus).""" + corpus = [ + "a cat is a feline and likes to purr", + "a dog is the human's best friend", + "a bird is a beautiful animal that can fly", + ] + + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + mock_dashtext.SparseVectorEncoder.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + bm25 = BM25EmbeddingFunction(corpus=corpus, b=0.75, k1=1.2) + + assert bm25.corpus_size == 3 + assert bm25.encoding_type == "query" + mock_dashtext.SparseVectorEncoder.assert_called_once_with(b=0.75, k1=1.2) + mock_encoder.train.assert_called_once_with(corpus) + + def test_init_with_empty_corpus(self): + """Test initialization with empty corpus raises ValueError.""" + with pytest.raises(ValueError, match="Corpus must be a non-empty list"): + BM25EmbeddingFunction(corpus=[]) + + def test_init_with_invalid_corpus(self): + """Test initialization with invalid corpus elements.""" + with pytest.raises(ValueError, match="All corpus documents must be strings"): + BM25EmbeddingFunction(corpus=["text", 123, "another"]) + + with pytest.raises(ValueError, match="All corpus documents must be strings"): + BM25EmbeddingFunction(corpus=[None, "text"]) + + def test_init_with_language_parameter(self): + """Test initialization with different language settings.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + # Test English language + bm25_en = BM25EmbeddingFunction(language="en") + assert bm25_en.language == "en" + mock_dashtext.SparseVectorEncoder.default.assert_called_with(name="en") + + def test_init_with_encoding_type(self): + """Test initialization with different encoding types.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + # Test document encoding type + bm25_doc = BM25EmbeddingFunction(encoding_type="document") + assert bm25_doc.encoding_type == "document" + + def test_init_with_missing_dashtext_library(self): + """Test initialization fails when dashtext library is not installed.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_require.side_effect = ImportError("dashtext package is required") + + with pytest.raises(ImportError, match="dashtext package is required"): + BM25EmbeddingFunction() + + def test_embed_with_query_encoding(self): + """Test successful sparse embedding generation with query encoding.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + + # Mock encode_queries to return sparse vector + mock_encoder.encode_queries.return_value = { + 5: 0.89, + 12: 1.45, + 23: 0.67, + 45: 1.12, + } + + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + bm25 = BM25EmbeddingFunction(encoding_type="query") + # Clear LRU cache to ensure fresh call + bm25.embed.cache_clear() + result = bm25.embed("cat purr loud") + + # Verify result structure + assert isinstance(result, dict) + assert all(isinstance(k, int) for k in result.keys()) + assert all(isinstance(v, float) for v in result.values()) + + # Verify all values are positive + assert all(v > 0 for v in result.values()) + + # Verify output is sorted by indices + keys = list(result.keys()) + assert keys == sorted(keys), "Output must be sorted by indices" + + # Verify expected keys from mock response + assert result == {5: 0.89, 12: 1.45, 23: 0.67, 45: 1.12} + + # Verify encode_queries was called + mock_encoder.encode_queries.assert_called_once_with("cat purr loud") + + def test_embed_with_document_encoding(self): + """Test successful sparse embedding generation with document encoding.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + + # Mock encode_documents to return sparse vector + mock_encoder.encode_documents.return_value = {10: 1.5, 20: 2.3} + + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + bm25 = BM25EmbeddingFunction(encoding_type="document") + bm25.embed.cache_clear() + result = bm25.embed("document text") + + assert result == {10: 1.5, 20: 2.3} + mock_encoder.encode_documents.assert_called_once_with("document text") + + def test_embed_with_empty_input(self): + """Test embedding with empty input raises ValueError.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + bm25 = BM25EmbeddingFunction() + + with pytest.raises(ValueError, match="Input text cannot be empty"): + bm25.embed("") + + with pytest.raises(ValueError, match="Input text cannot be empty"): + bm25.embed(" ") + + def test_embed_with_non_string_input(self): + """Test embedding with non-string input raises TypeError.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + bm25 = BM25EmbeddingFunction() + + # Test with hashable non-string types - should get our custom error message + with pytest.raises(TypeError, match="Expected 'input' to be str"): + bm25.embed(123) + + with pytest.raises(TypeError, match="Expected 'input' to be str"): + bm25.embed(None) + + # Test with unhashable type (list) + # Note: lru_cache raises TypeError("unhashable type: 'list'") before our type check + # This is still a valid type error, just caught at a different layer + with pytest.raises(TypeError, match="unhashable type"): + bm25.embed(["text"]) + + def test_embed_callable_interface(self): + """Test that BM25EmbeddingFunction is callable.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + mock_encoder.encode_queries.return_value = {10: 1.5} + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + bm25 = BM25EmbeddingFunction() + bm25.embed.cache_clear() + + # Test callable interface + result = bm25("test query") + assert isinstance(result, dict) + assert 10 in result + + def test_embed_output_sorted_by_indices(self): + """Test that output is always sorted by indices in ascending order.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + + # Mock encode_queries with unsorted indices + mock_encoder.encode_queries.return_value = { + 9999: 1.5, + 5: 2.0, + 1234: 0.8, + 77: 3.2, + 500: 1.1, + } + + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + bm25 = BM25EmbeddingFunction() + bm25.embed.cache_clear() + result = bm25.embed("test query") + + # Verify keys are sorted + result_keys = list(result.keys()) + assert result_keys == sorted(result_keys), ( + f"Keys must be sorted. Got: {result_keys}, Expected: {sorted(result_keys)}" + ) + + # Verify expected sorted order: [5, 77, 500, 1234, 9999] + expected_keys = [5, 77, 500, 1234, 9999] + assert result_keys == expected_keys + + def test_embed_filters_zero_values(self): + """Test that zero and negative values are filtered out.""" + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + + # Mock encode_queries with zero and negative values + mock_encoder.encode_queries.return_value = { + 0: 1.5, # Positive - should be included + 1: 0.0, # Zero - should be filtered + 2: -0.5, # Negative - should be filtered + } + + mock_dashtext.SparseVectorEncoder.default.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + bm25 = BM25EmbeddingFunction() + bm25.embed.cache_clear() + result = bm25.embed("test") + + # Only positive token should be in result + assert 0 in result + assert 1 not in result # Zero value filtered + assert 2 not in result # Negative value filtered + assert all(v > 0 for v in result.values()) + + def test_properties(self): + """Test property accessors.""" + corpus = ["doc1", "doc2", "doc3"] + + with patch( + "zvec.extension.bm25_embedding_function.require_module" + ) as mock_require: + mock_dashtext = Mock() + mock_encoder = Mock() + mock_dashtext.SparseVectorEncoder.return_value = mock_encoder + mock_require.return_value = mock_dashtext + + bm25 = BM25EmbeddingFunction( + corpus=corpus, + encoding_type="document", + language="en", + b=0.8, + k1=1.5, + custom_param="test", + ) + + assert bm25.corpus_size == 3 + assert bm25.encoding_type == "document" + assert bm25.language == "en" + assert bm25.extra_params == {"custom_param": "test"} + + @pytest.mark.skipif( + not RUN_INTEGRATION_TESTS, + reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.", + ) + def test_real_dashtext_bm25_embedding(self): + """Integration test with real DashText library. + + To run this test: + export ZVEC_RUN_INTEGRATION_TESTS=1 + pip install dashtext + + Note: This test requires the dashtext package to be installed. + """ + # Test built-in encoder (Chinese) + bm25_zh = BM25EmbeddingFunction(language="zh", encoding_type="query") + + query_zh = "什么是向量检索服务" + result_zh = bm25_zh.embed(query_zh) + + assert isinstance(result_zh, dict) + assert len(result_zh) > 0 + assert all(isinstance(k, int) for k in result_zh.keys()) + assert all(isinstance(v, float) and v > 0 for v in result_zh.values()) + + # Verify sorted output + keys = list(result_zh.keys()) + assert keys == sorted(keys), "Real DashText BM25 output must be sorted" + + # Test custom corpus + corpus = [ + "The cat sits on the mat", + "The dog plays in the garden", + "Birds fly in the sky", + "Fish swim in the water", + ] + + bm25_custom = BM25EmbeddingFunction(corpus=corpus, encoding_type="query") + + query_en = "cat on mat" + result_en = bm25_custom.embed(query_en) + + assert isinstance(result_en, dict) + assert len(result_en) > 0 + assert all(isinstance(k, int) for k in result_en.keys()) + assert all(isinstance(v, float) and v > 0 for v in result_en.values()) + + # Test callable interface + result2 = bm25_custom(query_en) + assert result_en == result2 + + # Verify properties + assert bm25_custom.corpus_size == 4 diff --git a/python/tests/test_reranker.py b/python/tests/test_reranker.py index 5b2c177..dced1dd 100644 --- a/python/tests/test_reranker.py +++ b/python/tests/test_reranker.py @@ -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]}...") diff --git a/python/tests/test_util.py b/python/tests/test_util.py index bac8926..c5a56c1 100644 --- a/python/tests/test_util.py +++ b/python/tests/test_util.py @@ -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") diff --git a/python/zvec/__init__.py b/python/zvec/__init__.py index ec35829..1c8fdfc 100644 --- a/python/zvec/__init__.py +++ b/python/zvec/__init__.py @@ -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", diff --git a/python/zvec/common/constants.py b/python/zvec/common/constants.py index 56b82fd..c8da216 100644 --- a/python/zvec/common/constants.py +++ b/python/zvec/common/constants.py @@ -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) diff --git a/python/zvec/extension/__init__.py b/python/zvec/extension/__init__.py index 83421b5..597f91b 100644 --- a/python/zvec/extension/__init__.py +++ b/python/zvec/extension/__init__.py @@ -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", ] diff --git a/python/zvec/extension/bm25_embedding_function.py b/python/zvec/extension/bm25_embedding_function.py new file mode 100644 index 0000000..51ab5ac --- /dev/null +++ b/python/zvec/extension/bm25_embedding_function.py @@ -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 diff --git a/python/zvec/extension/embedding.py b/python/zvec/extension/embedding.py deleted file mode 100644 index 1bbb096..0000000 --- a/python/zvec/extension/embedding.py +++ /dev/null @@ -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")) diff --git a/python/zvec/extension/embedding_function.py b/python/zvec/extension/embedding_function.py new file mode 100644 index 0000000..a58ba23 --- /dev/null +++ b/python/zvec/extension/embedding_function.py @@ -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. + """ + ... diff --git a/python/zvec/extension/multi_vector_reranker.py b/python/zvec/extension/multi_vector_reranker.py new file mode 100644 index 0000000..ba3a236 --- /dev/null +++ b/python/zvec/extension/multi_vector_reranker.py @@ -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") diff --git a/python/zvec/extension/openai_embedding_function.py b/python/zvec/extension/openai_embedding_function.py new file mode 100644 index 0000000..03a34ed --- /dev/null +++ b/python/zvec/extension/openai_embedding_function.py @@ -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 diff --git a/python/zvec/extension/openai_function.py b/python/zvec/extension/openai_function.py new file mode 100644 index 0000000..d3f4de2 --- /dev/null +++ b/python/zvec/extension/openai_function.py @@ -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 diff --git a/python/zvec/extension/qwen_embedding_function.py b/python/zvec/extension/qwen_embedding_function.py new file mode 100644 index 0000000..7bdb69b --- /dev/null +++ b/python/zvec/extension/qwen_embedding_function.py @@ -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) + + >>> 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())) diff --git a/python/zvec/extension/qwen_function.py b/python/zvec/extension/qwen_function.py new file mode 100644 index 0000000..b15ee4b --- /dev/null +++ b/python/zvec/extension/qwen_function.py @@ -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 diff --git a/python/zvec/extension/qwen_rerank_function.py b/python/zvec/extension/qwen_rerank_function.py new file mode 100644 index 0000000..9b4a66b --- /dev/null +++ b/python/zvec/extension/qwen_rerank_function.py @@ -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 diff --git a/python/zvec/extension/rerank.py b/python/zvec/extension/rerank.py deleted file mode 100644 index 021f6ed..0000000 --- a/python/zvec/extension/rerank.py +++ /dev/null @@ -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 diff --git a/python/zvec/extension/rerank_function.py b/python/zvec/extension/rerank_function.py new file mode 100644 index 0000000..c558a2b --- /dev/null +++ b/python/zvec/extension/rerank_function.py @@ -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. + """ + ... diff --git a/python/zvec/extension/sentence_transformer_embedding_function.py b/python/zvec/extension/sentence_transformer_embedding_function.py new file mode 100644 index 0000000..032f02e --- /dev/null +++ b/python/zvec/extension/sentence_transformer_embedding_function.py @@ -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) + + >>> 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 diff --git a/python/zvec/extension/sentence_transformer_function.py b/python/zvec/extension/sentence_transformer_function.py new file mode 100644 index 0000000..1ba1662 --- /dev/null +++ b/python/zvec/extension/sentence_transformer_function.py @@ -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") diff --git a/python/zvec/extension/sentence_transformer_rerank_function.py b/python/zvec/extension/sentence_transformer_rerank_function.py new file mode 100644 index 0000000..58c5838 --- /dev/null +++ b/python/zvec/extension/sentence_transformer_rerank_function.py @@ -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 diff --git a/python/zvec/tool/util.py b/python/zvec/tool/util.py index a836876..409a4d5 100644 --- a/python/zvec/tool/util.py +++ b/python/zvec/tool/util.py @@ -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