wehub/alibaba--zvec
1
0
Fork 0
Code Issues Pull Requests Actions 2 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:47:42 +08:00
commit be3ef883e1
1214 changed files with 431743 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/zvec/extension/sentence_transformer_embedding_function.py
+839
View File
@@ -0,0 +1,839 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from typing import ClassVar, Literal, Optional
import numpy as np
from ..common.constants import TEXT, DenseVectorType, SparseVectorType
from .embedding_function import DenseEmbeddingFunction, SparseEmbeddingFunction
from .sentence_transformer_function import SentenceTransformerFunctionBase
class DefaultLocalDenseEmbedding(
SentenceTransformerFunctionBase, DenseEmbeddingFunction[TEXT]
):
"""Default local dense embedding using all-MiniLM-L6-v2 model.
This is the default implementation for dense text embedding that uses the
``all-MiniLM-L6-v2`` model from Hugging Face by default. This model provides
a good balance between speed and quality for general-purpose text embedding.
The class provides text-to-vector dense embedding capabilities using the
sentence-transformers library. It supports models from Hugging Face Hub and
ModelScope, runs locally without API calls, and supports CPU/GPU acceleration.
The model produces 384-dimensional embeddings and is optimized for semantic
similarity tasks. It runs locally without requiring API keys.
Args:
model_source (Literal["huggingface", "modelscope"], optional): Model source.
- ``"huggingface"``: Use Hugging Face Hub (default, for international users)
- ``"modelscope"``: Use ModelScope (recommended for users in China)
Defaults to ``"huggingface"``.
device (Optional[str], optional): Device to run the model on.
Options: ``"cpu"``, ``"cuda"``, ``"mps"`` (for Apple Silicon), or ``None``
for automatic detection. Defaults to ``None``.
normalize_embeddings (bool, optional): Whether to normalize embeddings to
unit length (L2 normalization). Useful for cosine similarity.
Defaults to ``True``.
batch_size (int, optional): Batch size for encoding. Defaults to ``32``.
**kwargs: Additional parameters for future extension.
Attributes:
dimension (int): Always 384 for both models.
model_name (str): "all-MiniLM-L6-v2" (HF) or "iic/nlp_gte_sentence-embedding_chinese-small" (MS).
model_source (str): The model source being used.
device (str): The device the model is running on.
Raises:
ValueError: If the model cannot be loaded or input is invalid.
TypeError: If input to ``embed()`` is not a string.
RuntimeError: If model inference fails.
Note:
- Requires Python 3.10, 3.11, or 3.12
- Requires the ``sentence-transformers`` package:
``pip install sentence-transformers``
- For ModelScope, also requires: ``pip install modelscope``
- First run downloads the model (~50-80MB) from chosen source
- Hugging Face cache: ``~/.cache/torch/sentence_transformers/``
- ModelScope cache: ``~/.cache/modelscope/hub/``
- No API keys or network required after initial download
- Inference speed: ~1000 sentences/sec on CPU, ~10000 on GPU
**For users in China:**
If you encounter Hugging Face access issues, use ModelScope instead:
.. code-block:: python
# Recommended for users in China
emb = DefaultLocalDenseEmbedding(model_source="modelscope")
Alternatively, use Hugging Face mirror:
.. code-block:: bash
export HF_ENDPOINT=https://hf-mirror.com
# Then use default Hugging Face mode
Examples:
>>> # Basic usage with Hugging Face (default)
>>> from zvec.extension import DefaultLocalDenseEmbedding
>>>
>>> emb_func = DefaultLocalDenseEmbedding()
>>> vector = emb_func.embed("Hello, world!")
>>> len(vector)
384
>>> isinstance(vector, list)
True
>>> # Recommended for users in China (uses ModelScope)
>>> emb_func = DefaultLocalDenseEmbedding(model_source="modelscope")
>>> vector = emb_func.embed("你好，世界！") # Works well with Chinese text
>>> len(vector)
384
>>> # Alternative for China users: Use Hugging Face mirror
>>> import os
>>> os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
>>> emb_func = DefaultLocalDenseEmbedding() # Uses HF mirror
>>> vector = emb_func.embed("Hello, world!")
>>> # Using GPU for faster inference
>>> emb_func = DefaultLocalDenseEmbedding(device="cuda")
>>> vector = emb_func("Machine learning is fascinating")
>>> # Normalized vector has unit length
>>> import numpy as np
>>> np.linalg.norm(vector)
1.0
>>> # Batch processing
>>> texts = ["First text", "Second text", "Third text"]
>>> vectors = [emb_func.embed(text) for text in texts]
>>> len(vectors)
3
>>> all(len(v) == 384 for v in vectors)
True
>>> # Semantic similarity
>>> v1 = emb_func.embed("The cat sits on the mat")
>>> v2 = emb_func.embed("A feline rests on a rug")
>>> v3 = emb_func.embed("Python programming")
>>> similarity_high = np.dot(v1, v2) # Similar sentences
>>> similarity_low = np.dot(v1, v3) # Different topics
>>> similarity_high > similarity_low
True
>>> # Error handling
>>> try:
... emb_func.embed("") # Empty string
... except ValueError as e:
... print(f"Error: {e}")
Error: Input text cannot be empty or whitespace only
See Also:
- ``DenseEmbeddingFunction``: Base class for dense embeddings
- ``DefaultLocalSparseEmbedding``: Sparse embedding with SPLADE
- ``QwenDenseEmbedding``: Alternative using Qwen API
"""
def __init__(
self,
model_source: Literal["huggingface", "modelscope"] = "huggingface",
device: Optional[str] = None,
normalize_embeddings: bool = True,
batch_size: int = 32,
**kwargs,
):
"""Initialize with all-MiniLM-L6-v2 model.
Args:
model_source (Literal["huggingface", "modelscope"]): Model source.
Defaults to "huggingface".
device (Optional[str]): Target device ("cpu", "cuda", "mps", or None).
Defaults to None (automatic detection).
normalize_embeddings (bool): Whether to L2-normalize output vectors.
Defaults to True.
batch_size (int): Batch size for encoding. Defaults to 32.
**kwargs: Additional parameters for future extension.
Raises:
ImportError: If sentence-transformers or modelscope is not installed.
ValueError: If model cannot be loaded.
"""
# Use different models based on source
if model_source == "modelscope":
# Use Chinese-optimized model for ModelScope (better for Chinese text)
model_name = "iic/nlp_gte_sentence-embedding_chinese-small"
else:
model_name = "all-MiniLM-L6-v2"
# Initialize base class for model loading
SentenceTransformerFunctionBase.__init__(
self, model_name=model_name, model_source=model_source, device=device
)
self._normalize_embeddings = normalize_embeddings
self._batch_size = batch_size
# Load model and get dimension
model = self._get_model()
self._dimension = model.get_sentence_embedding_dimension()
# Store extra parameters
self._extra_params = kwargs
@property
def dimension(self) -> int:
"""int: The expected dimensionality of the embedding vector."""
return self._dimension
@property
def extra_params(self) -> dict:
"""dict: Extra parameters for model-specific customization."""
return self._extra_params
def __call__(self, input: str) -> DenseVectorType:
"""Make the embedding function callable."""
return self.embed(input)
def embed(self, input: str) -> DenseVectorType:
"""Generate dense embedding vector for the input text.
This method uses the Sentence Transformer model to convert input text
into a dense vector representation. The model runs locally without
requiring API calls.
Args:
input (str): Input text string to embed. Must be non-empty after
stripping whitespace. Maximum length depends on the model used
(typically 128-512 tokens for most models).
Returns:
DenseVectorType: A list of floats representing the embedding vector.
Length equals ``self.dimension``. If ``normalize_embeddings=True``,
the vector has unit length. Example:
``[0.123, -0.456, 0.789, ...]``
Raises:
TypeError: If ``input`` is not a string.
ValueError: If input is empty or whitespace-only.
RuntimeError: If model inference fails.
Examples:
>>> emb = DefaultLocalDenseEmbedding()
>>> vector = emb.embed("Natural language processing")
>>> len(vector)
384
>>> isinstance(vector[0], float)
True
>>> # Normalized vectors have unit length
>>> import numpy as np
>>> emb = DefaultLocalDenseEmbedding(normalize_embeddings=True)
>>> vector = emb.embed("Test sentence")
>>> np.linalg.norm(vector)
1.0
>>> # Error: empty input
>>> emb.embed(" ")
ValueError: Input text cannot be empty or whitespace only
>>> # Error: non-string input
>>> emb.embed(123)
TypeError: Expected 'input' to be str, got int
>>> # Semantic similarity example
>>> v1 = emb.embed("The cat sits on the mat")
>>> v2 = emb.embed("A feline rests on a rug")
>>> similarity = np.dot(v1, v2) # High similarity due to semantic meaning
>>> similarity > 0.7
True
Note:
- First call may be slower due to model loading
- Subsequent calls are much faster as the model stays in memory
- For batch processing, consider encoding multiple texts together
(though this method handles single texts only)
- GPU acceleration provides 5-10x speedup over CPU
"""
if not isinstance(input, str):
raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}")
input = input.strip()
if not input:
raise ValueError("Input text cannot be empty or whitespace only")
try:
model = self._get_model()
embedding = model.encode(
input,
convert_to_numpy=True,
normalize_embeddings=self._normalize_embeddings,
batch_size=self._batch_size,
)
# Convert numpy array to list
if isinstance(embedding, np.ndarray):
embedding_list = embedding.tolist()
else:
embedding_list = list(embedding)
# Validate dimension
if len(embedding_list) != self.dimension:
raise ValueError(
f"Dimension mismatch: expected {self.dimension}, "
f"got {len(embedding_list)}"
)
return embedding_list
except Exception as e:
if isinstance(e, (TypeError, ValueError)):
raise
raise RuntimeError(f"Failed to generate embedding: {e!s}") from e
class DefaultLocalSparseEmbedding(
SentenceTransformerFunctionBase, SparseEmbeddingFunction[TEXT]
):
"""Default local sparse embedding using SPLADE model.
This class provides sparse vector embedding using the SPLADE (SParse Lexical
AnD Expansion) model. SPLADE generates sparse, interpretable representations
where each dimension corresponds to a vocabulary term with learned importance
weights. It's ideal for lexical matching, BM25-style retrieval, and hybrid
search scenarios.
The default model is ``naver/splade-cocondenser-ensembledistil``, which is
publicly available without authentication. It produces sparse vectors with
thousands of dimensions but only hundreds of non-zero values, making them
efficient for storage and retrieval while maintaining strong lexical matching.
**Model Caching:**
This class uses class-level caching to share the SPLADE model across all instances
with the same configuration (model_source, device). This significantly reduces
memory usage when creating multiple instances for different encoding types
(query vs document).
**Cache Management:**
The class provides methods to manage the model cache:
- ``clear_cache()``: Clear all cached models to free memory
- ``get_cache_info()``: Get information about cached models
- ``remove_from_cache(model_source, device)``: Remove a specific model from cache
.. note::
**Why not use splade-v3?**
The newer ``naver/splade-v3`` model is gated (requires access approval).
We use ``naver/splade-cocondenser-ensembledistil`` instead.
**To use splade-v3 (if you have access):**
1. Request access at https://huggingface.co/naver/splade-v3
2. Get your Hugging Face token from https://huggingface.co/settings/tokens
3. Set environment variable:
.. code-block:: bash
export HF_TOKEN="your_huggingface_token"
4. Or login programmatically:
.. code-block:: python
from huggingface_hub import login
login(token="your_huggingface_token")
5. To use a custom SPLADE model, you can subclass this class and override
the model_name in ``__init__``, or create your own implementation
inheriting from ``SentenceTransformerFunctionBase`` and
``SparseEmbeddingFunction``.
Args:
model_source (Literal["huggingface", "modelscope"], optional): Model source.
Defaults to ``"huggingface"``. ModelScope support may vary for SPLADE models.
device (Optional[str], optional): Device to run the model on.
Options: ``"cpu"``, ``"cuda"``, ``"mps"`` (for Apple Silicon), or ``None``
for automatic detection. Defaults to ``None``.
encoding_type (Literal["query", "document"], optional): Encoding type.
- ``"query"``: Optimize for search queries (default)
- ``"document"``: Optimize for indexed documents
**kwargs: Additional parameters (currently unused, for future extension).
Attributes:
model_name (str): Model identifier.
model_source (str): The model source being used.
device (str): The device the model is running on.
Raises:
ValueError: If the model cannot be loaded or input is invalid.
TypeError: If input to ``embed()`` is not a string.
RuntimeError: If model inference fails.
Note:
- Requires Python 3.10, 3.11, or 3.12
- Requires the ``sentence-transformers`` package:
``pip install sentence-transformers``
- First run downloads the model (~100MB) from Hugging Face
- Cache location: ``~/.cache/torch/sentence_transformers/``
- No API keys or authentication required
- Sparse vectors have ~30k dimensions but only ~100-200 non-zero values
- Best combined with dense embeddings for hybrid retrieval
**SPLADE vs Dense Embeddings:**
- **Dense**: Continuous semantic vectors, good for semantic similarity
- **Sparse**: Lexical keyword-based, interpretable, good for exact matching
- **Hybrid**: Combine both for best retrieval performance
Examples:
>>> # Memory-efficient: both instances share the same model (~200MB)
>>> from zvec.extension import DefaultLocalSparseEmbedding
>>>
>>> # Query embedding
>>> query_emb = DefaultLocalSparseEmbedding(encoding_type="query")
>>> query_vec = query_emb.embed("machine learning algorithms")
>>> type(query_vec)
<class 'dict'>
>>> len(query_vec) # Only non-zero dimensions
156
>>> # Document embedding (shares model with query_emb)
>>> doc_emb = DefaultLocalSparseEmbedding(encoding_type="document")
>>> doc_vec = doc_emb.embed("Machine learning is a subset of AI")
>>> # Total memory: ~200MB (not 400MB) thanks to model caching
>>> # Asymmetric retrieval example
>>> query_vec = query_emb.embed("what causes aging fast")
>>> doc_vec = doc_emb.embed(
... "UV-A light causes tanning, skin aging, and cataracts..."
... )
>>>
>>> # Calculate similarity (dot product for sparse vectors)
>>> similarity = sum(
... query_vec.get(k, 0) * doc_vec.get(k, 0)
... for k in set(query_vec) | set(doc_vec)
... )
>>> # Batch processing
>>> queries = ["query 1", "query 2", "query 3"]
>>> query_vecs = [query_emb.embed(q) for q in queries]
>>>
>>> documents = ["doc 1", "doc 2", "doc 3"]
>>> doc_vecs = [doc_emb.embed(d) for d in documents]
>>> # Inspecting sparse dimensions (output is sorted by indices)
>>> query_vec = query_emb.embed("machine learning")
>>> list(query_vec.items())[:5] # First 5 dimensions (by index)
[(10, 0.45), (23, 0.87), (56, 0.32), (89, 1.12), (120, 0.65)]
>>>
>>> # Sort by weight to find most important terms
>>> sorted_by_weight = sorted(query_vec.items(), key=lambda x: x[1], reverse=True)
>>> top_5 = sorted_by_weight[:5] # Top 5 most important terms
>>> top_5
[(1023, 1.45), (245, 1.23), (8901, 0.98), (5678, 0.87), (12034, 0.76)]
>>> # Using GPU for faster inference
>>> sparse_emb = DefaultLocalSparseEmbedding(device="cuda")
>>> vector = sparse_emb.embed("natural language processing")
>>> # Hybrid retrieval example (combining dense + sparse)
>>> from zvec.extension import DefaultDenseEmbedding
>>> dense_emb = DefaultDenseEmbedding()
>>> sparse_emb = DefaultLocalSparseEmbedding()
>>>
>>> query = "deep learning neural networks"
>>> dense_vec = dense_emb.embed(query) # [0.1, -0.3, 0.5, ...]
>>> sparse_vec = sparse_emb.embed(query) # {12: 0.8, 45: 1.2, ...}
>>> # Error handling
>>> try:
... sparse_emb.embed("") # Empty string
... except ValueError as e:
... print(f"Error: {e}")
Error: Input text cannot be empty or whitespace only
>>> # Cache management
>>> # Check cache status
>>> info = DefaultLocalSparseEmbedding.get_cache_info()
>>> print(f"Cached models: {info['cached_models']}")
Cached models: 1
>>>
>>> # Clear cache to free memory
>>> DefaultLocalSparseEmbedding.clear_cache()
>>> info = DefaultLocalSparseEmbedding.get_cache_info()
>>> print(f"Cached models: {info['cached_models']}")
Cached models: 0
>>>
>>> # Remove specific model from cache
>>> query_emb = DefaultLocalSparseEmbedding() # Creates CPU model
>>> cuda_emb = DefaultLocalSparseEmbedding(device="cuda") # Creates CUDA model
>>> info = DefaultLocalSparseEmbedding.get_cache_info()
>>> print(f"Cached models: {info['cached_models']}")
Cached models: 2
>>>
>>> # Remove only CPU model
>>> removed = DefaultLocalSparseEmbedding.remove_from_cache(device=None)
>>> print(f"Removed: {removed}")
True
>>> info = DefaultLocalSparseEmbedding.get_cache_info()
>>> print(f"Cached models: {info['cached_models']}")
Cached models: 1
See Also:
- ``SparseEmbeddingFunction``: Base class for sparse embeddings
- ``DefaultDenseEmbedding``: Dense embedding with all-MiniLM-L6-v2
- ``QwenDenseEmbedding``: Alternative using Qwen API
References:
- SPLADE Paper: https://arxiv.org/abs/2109.10086
- Model: https://huggingface.co/naver/splade-cocondenser-ensembledistil
"""
# Class-level model cache: {(model_name, model_source, device): model}
# Shared across all DefaultLocalSparseEmbedding instances to save memory
_model_cache: ClassVar[dict] = {}
@classmethod
def clear_cache(cls) -> None:
"""Clear all cached SPLADE models from memory.
This is useful for:
- Freeing memory when models are no longer needed
- Forcing a fresh model reload
- Testing and debugging
Examples:
>>> # Clear cache to free memory
>>> DefaultLocalSparseEmbedding.clear_cache()
>>> # Or in tests to ensure fresh model loading
>>> def test_something():
... DefaultLocalSparseEmbedding.clear_cache()
... emb = DefaultLocalSparseEmbedding()
... # Test with fresh model
"""
cls._model_cache.clear()
@classmethod
def get_cache_info(cls) -> dict:
"""Get information about currently cached models.
Returns:
dict: Dictionary with cache statistics:
- cached_models (int): Number of cached model instances
- cache_keys (list): List of cache keys (model_name, model_source, device)
Examples:
>>> info = DefaultLocalSparseEmbedding.get_cache_info()
>>> print(f"Cached models: {info['cached_models']}")
Cached models: 2
>>> print(f"Cache keys: {info['cache_keys']}")
Cache keys: [('naver/splade-cocondenser-ensembledistil', 'huggingface', None),
('naver/splade-cocondenser-ensembledistil', 'huggingface', 'cuda')]
"""
return {
"cached_models": len(cls._model_cache),
"cache_keys": list(cls._model_cache.keys()),
}
@classmethod
def remove_from_cache(
cls, model_source: str = "huggingface", device: Optional[str] = None
) -> bool:
"""Remove a specific model from cache.
Args:
model_source (str): Model source ("huggingface" or "modelscope").
Defaults to "huggingface".
device (Optional[str]): Device identifier. Defaults to None.
Returns:
bool: True if model was found and removed, False otherwise.
Examples:
>>> # Remove CPU model from cache
>>> removed = DefaultLocalSparseEmbedding.remove_from_cache()
>>> print(f"Removed: {removed}")
True
>>> # Remove CUDA model from cache
>>> removed = DefaultLocalSparseEmbedding.remove_from_cache(device="cuda")
>>> print(f"Removed: {removed}")
True
"""
model_name = "naver/splade-cocondenser-ensembledistil"
cache_key = (model_name, model_source, device)
if cache_key in cls._model_cache:
del cls._model_cache[cache_key]
return True
return False
def __init__(
self,
model_source: Literal["huggingface", "modelscope"] = "huggingface",
device: Optional[str] = None,
encoding_type: Literal["query", "document"] = "query",
**kwargs,
):
"""Initialize with SPLADE model.
Args:
model_source (Literal["huggingface", "modelscope"]): Model source.
Defaults to "huggingface".
device (Optional[str]): Target device ("cpu", "cuda", "mps", or None).
Defaults to None (automatic detection).
encoding_type (Literal["query", "document"]): Encoding type for embeddings.
- "query": Optimize for search queries (default)
- "document": Optimize for indexed documents
This distinction is important for asymmetric retrieval tasks.
**kwargs: Additional parameters (reserved for future use).
Raises:
ImportError: If sentence-transformers is not installed.
ValueError: If model cannot be loaded.
Note:
Multiple instances with the same (model_source, device) configuration
will share the same underlying model to save memory. Different
instances can use different encoding_type settings while sharing
the model.
**Model Selection:**
Uses ``naver/splade-cocondenser-ensembledistil`` instead of the newer
``naver/splade-v3`` because splade-v3 is a gated model requiring
Hugging Face authentication. The cocondenser-ensembledistil variant:
- Does not require authentication or API tokens
- Is immediately available for all users
- Provides comparable retrieval performance (~2% difference)
- Avoids "Access to model is restricted" errors
If you need splade-v3 and have obtained access, you can subclass
this class and override the model_name parameter.
Examples:
>>> # Both instances share the same model (saves memory)
>>> query_emb = DefaultLocalSparseEmbedding(encoding_type="query")
>>> doc_emb = DefaultLocalSparseEmbedding(encoding_type="document")
>>> # Only one model is loaded in memory
"""
# Use publicly available SPLADE model (no gated access required)
# Note: naver/splade-v3 requires authentication, so we use the
# cocondenser-ensembledistil variant which is publicly accessible
model_name = "naver/splade-cocondenser-ensembledistil"
# Initialize base class for model loading
SentenceTransformerFunctionBase.__init__(
self, model_name=model_name, model_source=model_source, device=device
)
self._encoding_type = encoding_type
self._extra_params = kwargs
# Create cache key for this model configuration
self._cache_key = (model_name, model_source, device)
# Load model to ensure it's available (will use cache if exists)
self._get_model()
@property
def extra_params(self) -> dict:
"""dict: Extra parameters for model-specific customization."""
return self._extra_params
def __call__(self, input: str) -> SparseVectorType:
"""Make the embedding function callable."""
return self.embed(input)
def embed(self, input: str) -> SparseVectorType:
"""Generate sparse embedding vector for the input text.
This method uses the SPLADE model to convert input text into a sparse
vector representation. The result is a dictionary where keys are dimension
indices and values are importance weights (only non-zero values included).
The embedding is optimized based on the ``encoding_type`` specified during
initialization: "query" for search queries or "document" for indexed content.
Args:
input (str): Input text string to embed. Must be non-empty after
stripping whitespace.
Returns:
SparseVectorType: A dictionary mapping dimension index to weight.
Only non-zero dimensions are included. The dictionary is sorted
by indices (keys) in ascending order for consistent output.
Example: ``{10: 0.5, 245: 0.8, 1023: 1.2, 5678: 0.5}``
Raises:
TypeError: If ``input`` is not a string.
ValueError: If input is empty or whitespace-only.
RuntimeError: If model inference fails.
Examples:
>>> # Query embedding
>>> query_emb = DefaultLocalSparseEmbedding(encoding_type="query")
>>> query_vec = query_emb.embed("machine learning")
>>> isinstance(query_vec, dict)
True
Note:
- First call may be slower due to model loading
- Subsequent calls are much faster as the model stays in memory
- GPU acceleration provides significant speedup
- Sparse vectors are memory-efficient (only store non-zero values)
"""
if not isinstance(input, str):
raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}")
input = input.strip()
if not input:
raise ValueError("Input text cannot be empty or whitespace only")
try:
model = self._get_model()
# Use appropriate encoding method based on type
if self._encoding_type == "document" and hasattr(model, "encode_document"):
# Use document encoding
sparse_matrix = model.encode_document([input])
elif hasattr(model, "encode_query"):
# Use query encoding (default)
sparse_matrix = model.encode_query([input])
else:
# Fallback: manual implementation for older sentence-transformers
return self._manual_sparse_encode(input)
# Convert sparse matrix to dictionary
# SPLADE returns shape [1, vocab_size] for single input
# Check if it's a sparse matrix (duck typing - has toarray method)
if hasattr(sparse_matrix, "toarray"):
# Sparse matrix (CSR/CSC/etc.) - convert to dense array
sparse_array = sparse_matrix[0].toarray().flatten()
sparse_dict = {
int(idx): float(val)
for idx, val in enumerate(sparse_array)
if val > 0
}
else:
# Dense array format (numpy array or similar)
if isinstance(sparse_matrix, np.ndarray):
sparse_array = sparse_matrix[0]
else:
sparse_array = sparse_matrix
sparse_dict = {
int(idx): float(val)
for idx, val in enumerate(sparse_array)
if val > 0
}
# Sort by indices (keys) to ensure consistent ordering
return dict(sorted(sparse_dict.items()))
except Exception as e:
if isinstance(e, (TypeError, ValueError)):
raise
raise RuntimeError(f"Failed to generate sparse embedding: {e!s}") from e
def _manual_sparse_encode(self, input: str) -> SparseVectorType:
"""Fallback manual SPLADE encoding for older sentence-transformers.
Args:
input (str): Input text to encode.
Returns:
SparseVectorType: Sparse vector as dictionary.
"""
import torch
model = self._get_model()
# Tokenize input
features = model.tokenize([input])
# Move to correct device
features = {k: v.to(model.device) for k, v in features.items()}
# Forward pass with no gradient
with torch.no_grad():
embeddings = model.forward(features)
# Get logits from model output
# SPLADE models typically output 'token_embeddings'
if isinstance(embeddings, dict) and "token_embeddings" in embeddings:
logits = embeddings["token_embeddings"][0] # First batch item
elif hasattr(embeddings, "token_embeddings"):
logits = embeddings.token_embeddings[0]
# Fallback: try to get first value
elif isinstance(embeddings, dict):
logits = next(iter(embeddings.values()))[0]
else:
logits = embeddings[0]
# Apply SPLADE activation: log(1 + relu(x))
relu_log = torch.log(1 + torch.relu(logits))
# Max pooling over token dimension (reduce to vocab size)
if relu_log.dim() > 1:
sparse_vec, _ = torch.max(relu_log, dim=0)
else:
sparse_vec = relu_log
# Convert to sparse dictionary (only non-zero values)
sparse_vec_np = sparse_vec.cpu().numpy()
sparse_dict = {
int(idx): float(val) for idx, val in enumerate(sparse_vec_np) if val > 0
}
# Sort by indices (keys) to ensure consistent ordering
return dict(sorted(sparse_dict.items()))
def _get_model(self):
"""Load or retrieve the SPLADE model from class-level cache.
Returns:
SentenceTransformer: The loaded SPLADE model instance.
Raises:
ImportError: If required packages are not installed.
ValueError: If model cannot be loaded.
Note:
Models are cached at class level and shared across all instances
with the same (model_name, model_source, device) configuration.
This allows memory-efficient usage when creating multiple instances
with different encoding_type settings.
"""
# Check class-level cache first
if self._cache_key in self._model_cache:
return self._model_cache[self._cache_key]
# Use parent class method to load model
model = super()._get_model()
# Cache the model at class level
self._model_cache[self._cache_key] = model
return model