39 lines
1.3 KiB
Python
39 lines
1.3 KiB
Python
import asyncio
|
|
|
|
try:
|
|
from sentence_transformers import (
|
|
CrossEncoder, # pyright: ignore[reportMissingImports]
|
|
)
|
|
except ImportError as e: # pragma: no cover
|
|
raise ImportError(
|
|
"sentence-transformers is not installed. Install it with "
|
|
"`pip install sentence-transformers` or use the cross-encoder optional dependency."
|
|
) from e
|
|
|
|
from haiku.rag.reranking.base import RerankerBase
|
|
from haiku.rag.store.models.chunk import Chunk
|
|
|
|
|
|
class CrossEncoderReranker(RerankerBase):
|
|
"""Reranker for any sentence-transformers CrossEncoder model.
|
|
|
|
Loads the model in-process. Pass any HuggingFace cross-encoder reranker
|
|
as ``model`` (e.g. ``BAAI/bge-reranker-v2-m3``, ``Qwen/Qwen3-Reranker-0.6B``,
|
|
``cross-encoder/ms-marco-MiniLM-L-6-v2``).
|
|
"""
|
|
|
|
def __init__(self, model: str):
|
|
self._model = model
|
|
self._reranker = CrossEncoder(model)
|
|
|
|
async def rerank(
|
|
self, query: str, chunks: list[Chunk], top_n: int = 10
|
|
) -> list[tuple[Chunk, float]]:
|
|
if not chunks:
|
|
return []
|
|
|
|
documents = [chunk.content for chunk in chunks]
|
|
rankings = await asyncio.to_thread(
|
|
lambda: self._reranker.rank(query, documents, top_k=top_n)
|
|
)
|
|
return [(chunks[r["corpus_id"]], float(r["score"])) for r in rankings]
|