chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,396 @@
|
||||
# Copyright 2025-present the zvec project
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Literal, Optional
|
||||
|
||||
from ..model.doc import Doc, DocList
|
||||
from ..tool import require_module
|
||||
from .rerank_function import RerankFunction
|
||||
from .sentence_transformer_function import SentenceTransformerFunctionBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..model.schema import FieldSchema, VectorSchema
|
||||
|
||||
|
||||
class DefaultLocalReRanker(SentenceTransformerFunctionBase, RerankFunction):
|
||||
"""Re-ranker using Sentence Transformer cross-encoder models for semantic re-ranking.
|
||||
|
||||
This re-ranker leverages pre-trained cross-encoder models to perform deep semantic
|
||||
re-ranking of search results. It runs locally without API calls, supports GPU
|
||||
acceleration, and works with models from Hugging Face or ModelScope.
|
||||
|
||||
Cross-encoder models evaluate query-document pairs jointly, providing more
|
||||
accurate relevance scores than bi-encoder (embedding-based) similarity.
|
||||
|
||||
Args:
|
||||
query (str): Query text for semantic re-ranking. **Required**.
|
||||
rerank_field (Optional[str], optional): Document field name to use as
|
||||
re-ranking input text. **Required** (e.g., "content", "title", "body").
|
||||
model_name (str, optional): Cross-encoder model identifier or local path.
|
||||
Defaults to ``"cross-encoder/ms-marco-MiniLM-L6-v2"`` (MS MARCO MiniLM).
|
||||
Common options:
|
||||
- ``"cross-encoder/ms-marco-MiniLM-L6-v2"``: Lightweight, fast (~80MB, recommended)
|
||||
- ``"cross-encoder/ms-marco-MiniLM-L12-v2"``: Better accuracy (~120MB)
|
||||
- ``"BAAI/bge-reranker-base"``: BGE Reranker Base (~280MB)
|
||||
- ``"BAAI/bge-reranker-large"``: BGE Reranker Large (highest quality, ~560MB)
|
||||
model_source (Literal["huggingface", "modelscope"], optional): Model source.
|
||||
Defaults to ``"huggingface"``.
|
||||
- ``"huggingface"``: Load from Hugging Face Hub
|
||||
- ``"modelscope"``: Load from ModelScope (recommended for users in China)
|
||||
device (Optional[str], optional): Device to run the model on.
|
||||
Options: ``"cpu"``, ``"cuda"``, ``"mps"`` (for Apple Silicon), or ``None``
|
||||
for automatic detection. Defaults to ``None``.
|
||||
batch_size (int, optional): Batch size for processing query-document pairs.
|
||||
Larger values speed up processing but use more memory. Defaults to ``32``.
|
||||
|
||||
Attributes:
|
||||
query (str): The query text used for re-ranking.
|
||||
rerank_field (Optional[str]): Field name used for re-ranking input.
|
||||
model_name (str): The cross-encoder model being used.
|
||||
model_source (str): The model source ("huggingface" or "modelscope").
|
||||
device (str): The device the model is running on.
|
||||
|
||||
Raises:
|
||||
ValueError: If ``query`` is empty/None, ``rerank_field`` is None,
|
||||
or model cannot be loaded.
|
||||
TypeError: If input types are invalid.
|
||||
RuntimeError: If model inference fails.
|
||||
|
||||
Note:
|
||||
- Requires Python 3.10, 3.11, or 3.12
|
||||
- Requires ``sentence-transformers`` package: ``pip install sentence-transformers``
|
||||
- For ModelScope support, also requires: ``pip install modelscope``
|
||||
- First run downloads the model (~80-560MB depending on model) from chosen source
|
||||
- No API keys or network required after initial download
|
||||
- Cross-encoders are slower than bi-encoders but more accurate
|
||||
- GPU acceleration provides significant speedup (5-10x)
|
||||
|
||||
**MS MARCO MiniLM-L6-v2 Model (Default):**
|
||||
|
||||
The default model ``cross-encoder/ms-marco-MiniLM-L6-v2`` is a lightweight and
|
||||
efficient cross-encoder trained on MS MARCO dataset. It provides:
|
||||
|
||||
- Fast inference speed (suitable for real-time applications)
|
||||
- Small model size (~80MB, quick to download)
|
||||
- Good balance between speed and accuracy
|
||||
- Trained on 500K+ query-document pairs
|
||||
- Public availability without authentication
|
||||
|
||||
**For users in China:**
|
||||
|
||||
If you encounter Hugging Face access issues, use ModelScope instead:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Recommended for users in China
|
||||
reranker = SentenceTransformerReRanker(
|
||||
query="机器学习算法",
|
||||
rerank_field="content",
|
||||
model_source="modelscope"
|
||||
)
|
||||
|
||||
Alternatively, use Hugging Face mirror:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
export HF_ENDPOINT=https://hf-mirror.com
|
||||
|
||||
Examples:
|
||||
>>> # Basic usage with default MS MARCO MiniLM model
|
||||
>>> from zvec.extension import SentenceTransformerReRanker
|
||||
>>>
|
||||
>>> reranker = SentenceTransformerReRanker(
|
||||
... query="machine learning algorithms",
|
||||
... rerank_field="content"
|
||||
... )
|
||||
>>>
|
||||
>>> # Use in collection.query()
|
||||
>>> results = collection.query(
|
||||
... data={"vector_field": query_vector},
|
||||
... reranker=reranker,
|
||||
... topk=20
|
||||
... )
|
||||
|
||||
>>> # Using ModelScope for users in China
|
||||
>>> reranker = SentenceTransformerReRanker(
|
||||
... query="深度学习",
|
||||
... rerank_field="content",
|
||||
... model_source="modelscope"
|
||||
... )
|
||||
|
||||
>>> # Using larger model for better quality
|
||||
>>> reranker = SentenceTransformerReRanker(
|
||||
... query="neural networks",
|
||||
... rerank_field="content",
|
||||
... model_name="BAAI/bge-reranker-large",
|
||||
... device="cuda",
|
||||
... batch_size=64
|
||||
... )
|
||||
|
||||
>>> # Direct rerank call (for testing)
|
||||
>>> query_results = {
|
||||
... "vector1": [
|
||||
... Doc(id="1", score=0.9, fields={"content": "Machine learning is..."}),
|
||||
... Doc(id="2", score=0.8, fields={"content": "Deep learning is..."}),
|
||||
... ]
|
||||
... }
|
||||
>>> reranked = reranker.rerank(query_results)
|
||||
>>> for doc in reranked:
|
||||
... print(f"ID: {doc.id}, Score: {doc.score:.4f}")
|
||||
ID: 2, Score: 0.9234
|
||||
ID: 1, Score: 0.8567
|
||||
|
||||
See Also:
|
||||
- ``RerankFunction``: Abstract base class for re-rankers
|
||||
- ``QwenReRanker``: Re-ranker using Qwen API
|
||||
- ``RrfReRanker``: Multi-vector re-ranker using RRF
|
||||
- ``WeightedReRanker``: Multi-vector re-ranker using weighted scores
|
||||
|
||||
References:
|
||||
- MS MARCO Cross-Encoder: https://huggingface.co/cross-encoder/ms-marco-MiniLM-L6-v2
|
||||
- BGE Reranker: https://huggingface.co/BAAI/bge-reranker-base
|
||||
- Cross-Encoder vs Bi-Encoder: https://www.sbert.net/examples/applications/cross-encoder/README.html
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
query: Optional[str] = None,
|
||||
rerank_field: Optional[str] = None,
|
||||
model_name: str = "cross-encoder/ms-marco-MiniLM-L6-v2",
|
||||
model_source: Literal["huggingface", "modelscope"] = "huggingface",
|
||||
device: Optional[str] = None,
|
||||
batch_size: int = 32,
|
||||
):
|
||||
"""Initialize SentenceTransformerReRanker with query and configuration.
|
||||
|
||||
Args:
|
||||
query (Optional[str]): Query text for semantic matching. Required.
|
||||
rerank_field (Optional[str]): Document field for re-ranking input.
|
||||
model_name (str): Cross-encoder model identifier.
|
||||
model_source (Literal["huggingface", "modelscope"]): Model source.
|
||||
device (Optional[str]): Target device ("cpu", "cuda", "mps", or None).
|
||||
batch_size (int): Batch size for processing query-document pairs.
|
||||
|
||||
Raises:
|
||||
ValueError: If query is empty or model cannot be loaded.
|
||||
"""
|
||||
# Initialize base class for model loading
|
||||
SentenceTransformerFunctionBase.__init__(
|
||||
self, model_name=model_name, model_source=model_source, device=device
|
||||
)
|
||||
|
||||
# Initialize rerank parameters
|
||||
self._rerank_field = rerank_field
|
||||
|
||||
# Validate query
|
||||
if not query:
|
||||
raise ValueError("Query is required for DefaultLocalReRanker")
|
||||
self._query = query
|
||||
self._batch_size = batch_size
|
||||
|
||||
# Load and validate cross-encoder model
|
||||
model = self._get_model()
|
||||
if not hasattr(model, "predict"):
|
||||
raise ValueError(
|
||||
f"Model '{model_name}' does not appear to be a cross-encoder model. "
|
||||
"Cross-encoder models should have a 'predict' method."
|
||||
)
|
||||
self._model = model
|
||||
|
||||
def _get_model(self):
|
||||
"""Load or retrieve the CrossEncoder model.
|
||||
|
||||
This overrides the base class method to load CrossEncoder instead of
|
||||
SentenceTransformer, as reranking requires cross-encoder models.
|
||||
|
||||
Returns:
|
||||
CrossEncoder: The loaded cross-encoder model instance.
|
||||
|
||||
Raises:
|
||||
ImportError: If required packages are not installed.
|
||||
ValueError: If model cannot be loaded.
|
||||
"""
|
||||
# Return cached model if exists
|
||||
if self._model is not None:
|
||||
return self._model
|
||||
|
||||
# Load cross-encoder model
|
||||
try:
|
||||
sentence_transformers = require_module("sentence_transformers")
|
||||
|
||||
if self._model_source == "modelscope":
|
||||
# Load from ModelScope
|
||||
require_module("modelscope")
|
||||
from modelscope.hub.snapshot_download import snapshot_download
|
||||
|
||||
# Download model to cache
|
||||
model_dir = snapshot_download(self._model_name)
|
||||
|
||||
# Load CrossEncoder from local path
|
||||
model = sentence_transformers.CrossEncoder(
|
||||
model_dir, device=self._device
|
||||
)
|
||||
else:
|
||||
# Load CrossEncoder from Hugging Face (default)
|
||||
model = sentence_transformers.CrossEncoder(
|
||||
self._model_name, device=self._device
|
||||
)
|
||||
|
||||
return model
|
||||
|
||||
except ImportError as e:
|
||||
if "modelscope" in str(e) and self._model_source == "modelscope":
|
||||
raise ImportError(
|
||||
"ModelScope support requires the 'modelscope' package. "
|
||||
"Please install it with: pip install modelscope"
|
||||
) from e
|
||||
raise
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Failed to load CrossEncoder model '{self._model_name}' "
|
||||
f"from {self._model_source}: {e!s}"
|
||||
) from e
|
||||
|
||||
@property
|
||||
def rerank_field(self) -> Optional[str]:
|
||||
"""Optional[str]: Field name used as re-ranking input."""
|
||||
return self._rerank_field
|
||||
|
||||
@property
|
||||
def query(self) -> str:
|
||||
"""str: Query text used for semantic re-ranking."""
|
||||
return self._query
|
||||
|
||||
@property
|
||||
def batch_size(self) -> int:
|
||||
"""int: Batch size for processing query-document pairs."""
|
||||
return self._batch_size
|
||||
|
||||
def rerank(
|
||||
self,
|
||||
query_results: list[list[Doc]],
|
||||
topn: int = 10,
|
||||
*,
|
||||
fields: list[FieldSchema | VectorSchema] | None = None, # noqa: ARG002
|
||||
) -> DocList:
|
||||
"""Re-rank documents using Sentence Transformer cross-encoder model.
|
||||
|
||||
Evaluates each query-document pair using the cross-encoder model to compute
|
||||
relevance scores. Documents are then sorted by these scores and the top-k
|
||||
results are returned.
|
||||
|
||||
Args:
|
||||
query_results (list[list[Doc]]): Per-sub-query lists of retrieved
|
||||
documents. Documents from all lists are deduplicated and
|
||||
re-ranked together.
|
||||
topn (int): Maximum number of documents to return.
|
||||
fields: Unused; present for interface compatibility.
|
||||
|
||||
Returns:
|
||||
list[Doc]: Re-ranked documents (up to ``topn``) with updated ``score``
|
||||
fields containing relevance scores from the cross-encoder model.
|
||||
|
||||
Raises:
|
||||
ValueError: If no valid documents are found or model inference fails.
|
||||
|
||||
Note:
|
||||
- Duplicate documents (same ID) across fields are processed once
|
||||
- Documents with empty/missing ``rerank_field`` content are skipped
|
||||
- Returned scores are logits from the cross-encoder model
|
||||
- Higher scores indicate higher relevance
|
||||
- Processing time is O(n) where n is the number of documents
|
||||
|
||||
Examples:
|
||||
>>> reranker = SentenceTransformerReRanker(
|
||||
... query="machine learning",
|
||||
... topn=3,
|
||||
... rerank_field="content"
|
||||
... )
|
||||
>>> query_results = {
|
||||
... "vector1": [
|
||||
... Doc(id="1", score=0.9, fields={"content": "ML basics"}),
|
||||
... Doc(id="2", score=0.8, fields={"content": "DL tutorial"}),
|
||||
... ]
|
||||
... }
|
||||
>>> reranked = reranker.rerank(query_results)
|
||||
>>> len(reranked) <= 3
|
||||
True
|
||||
"""
|
||||
if not query_results:
|
||||
return []
|
||||
|
||||
# Accept both dict (legacy) and list formats
|
||||
if isinstance(query_results, dict):
|
||||
query_results = list(query_results.values())
|
||||
|
||||
# Collect and deduplicate documents
|
||||
id_to_doc: dict[str, Doc] = {}
|
||||
doc_ids: list[str] = []
|
||||
contents: list[str] = []
|
||||
|
||||
for query_result in query_results:
|
||||
for doc in query_result:
|
||||
doc_id = doc.id
|
||||
if doc_id in id_to_doc:
|
||||
continue
|
||||
|
||||
# Extract text content from specified field
|
||||
field_value = doc.field(self.rerank_field)
|
||||
rank_content = str(field_value).strip() if field_value else ""
|
||||
if not rank_content:
|
||||
continue
|
||||
|
||||
id_to_doc[doc_id] = doc
|
||||
doc_ids.append(doc_id)
|
||||
contents.append(rank_content)
|
||||
|
||||
if not contents:
|
||||
raise ValueError("No documents to rerank")
|
||||
|
||||
try:
|
||||
# Use standard cross-encoder predict method
|
||||
pairs = [[self.query, content] for content in contents]
|
||||
scores = self._model.predict(
|
||||
pairs,
|
||||
batch_size=self.batch_size,
|
||||
show_progress_bar=False,
|
||||
convert_to_numpy=True,
|
||||
)
|
||||
|
||||
# Convert to float list if needed
|
||||
if hasattr(scores, "tolist"):
|
||||
scores = scores.tolist()
|
||||
else:
|
||||
scores = [float(s) for s in scores]
|
||||
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to compute rerank scores: {e!s}") from e
|
||||
|
||||
# Create scored documents
|
||||
scored_docs = [
|
||||
(doc_ids[i], id_to_doc[doc_ids[i]], scores[i]) for i in range(len(doc_ids))
|
||||
]
|
||||
|
||||
# Sort by score (descending) and take top-k
|
||||
scored_docs.sort(key=lambda x: x[2], reverse=True)
|
||||
top_scored_docs = scored_docs[:topn]
|
||||
|
||||
# Build result list with updated scores
|
||||
results: DocList = []
|
||||
for _, doc, score in top_scored_docs:
|
||||
new_doc = doc._replace(score=score)
|
||||
results.append(new_doc)
|
||||
|
||||
return results
|
||||
Reference in New Issue
Block a user