d63199d96d
mxbai-rerank-base-v2 ships a Sigmoid activation and evaluates it in bf16, so every strongly-relevant candidate rounds to exactly 1.0. Ties then leave the order to the stable sort, which preserves the incoming hybrid ranking: on 100 t2_finqa retrieval cases the reranker scored MAP 0.661 against 0.659 with no reranker at all, and 0.742 once the scores separate. Ask the model for logits and apply the sigmoid here, where it runs in float64. Scores stay 0-1, matching the cohere, vllm and zeroentropy rerankers. Also drop the remaining pyright references; the project type-checks with ty.
36 lines
1.2 KiB
Python
36 lines
1.2 KiB
Python
import asyncio
|
|
|
|
try:
|
|
from transformers import AutoModel
|
|
except ImportError as e: # pragma: no cover
|
|
raise ImportError(
|
|
"transformers is not installed. Please install it with `pip install transformers torch` "
|
|
"or use the jina optional dependency."
|
|
) from e
|
|
|
|
from haiku.rag.reranking.base import RerankerBase
|
|
from haiku.rag.store.models.chunk import Chunk
|
|
|
|
|
|
class JinaLocalReranker(RerankerBase): # pragma: no cover
|
|
"""Jina reranker using local model inference via transformers.
|
|
|
|
Note: The Jina Reranker v3 model is licensed under CC BY-NC 4.0,
|
|
which restricts commercial use.
|
|
"""
|
|
|
|
def __init__(self, model: str = "jinaai/jina-reranker-v3"):
|
|
self._model = model
|
|
self._reranker = AutoModel.from_pretrained(model, trust_remote_code=True)
|
|
self._reranker.eval()
|
|
|
|
async def _rerank(
|
|
self, query: str, chunks: list[Chunk], top_n: int = 10
|
|
) -> list[tuple[Chunk, float]]:
|
|
documents = [chunk.content for chunk in chunks]
|
|
|
|
results = await asyncio.to_thread(
|
|
lambda: self._reranker.rerank(query, documents, top_n=top_n)
|
|
)
|
|
|
|
return [(chunks[r["index"]], float(r["relevance_score"])) for r in results]
|