41 lines
1.3 KiB
Python
41 lines
1.3 KiB
Python
import asyncio
|
|
|
|
try:
|
|
from transformers import (
|
|
AutoModel, # pyright: ignore[reportMissingImports]
|
|
)
|
|
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]]:
|
|
if not chunks:
|
|
return []
|
|
|
|
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]
|