60 lines
1.9 KiB
Python
60 lines
1.9 KiB
Python
from zeroentropy import AsyncZeroEntropy
|
|
|
|
from haiku.rag.reranking.base import RerankerBase
|
|
from haiku.rag.store.models.chunk import Chunk
|
|
|
|
|
|
class ZeroEntropyReranker(RerankerBase): # pragma: no cover
|
|
"""Zero Entropy reranker implementation using the zerank-1 model."""
|
|
|
|
def __init__(self, model: str = "zerank-1"):
|
|
"""Initialize the Zero Entropy reranker.
|
|
|
|
Args:
|
|
model: The Zero Entropy model to use (default: "zerank-1")
|
|
"""
|
|
self._model = model
|
|
# Zero Entropy SDK reads ZEROENTROPY_API_KEY from environment by default
|
|
self._client = AsyncZeroEntropy()
|
|
|
|
async def rerank(
|
|
self, query: str, chunks: list[Chunk], top_n: int = 10
|
|
) -> list[tuple[Chunk, float]]:
|
|
"""Rerank the given chunks based on relevance to the query.
|
|
|
|
Args:
|
|
query: The query to rank against
|
|
chunks: The chunks to rerank
|
|
top_n: The number of top results to return
|
|
|
|
Returns:
|
|
A list of (chunk, score) tuples, sorted by relevance
|
|
"""
|
|
if not chunks:
|
|
return []
|
|
|
|
# Prepare documents for Zero Entropy API
|
|
documents = [chunk.content for chunk in chunks]
|
|
|
|
# Call Zero Entropy reranking API
|
|
model_name = self._model or "zerank-1"
|
|
response = await self._client.models.rerank(
|
|
model=model_name,
|
|
query=query,
|
|
documents=documents,
|
|
)
|
|
|
|
# Extract results and map back to chunks
|
|
# Zero Entropy returns results sorted by relevance with scores
|
|
reranked_results = []
|
|
|
|
# Get top_n results
|
|
for i, result in enumerate(response.results[:top_n]):
|
|
# Zero Entropy returns index and score for each document
|
|
chunk_index = result.index
|
|
score = result.relevance_score
|
|
|
|
if chunk_index < len(chunks):
|
|
reranked_results.append((chunks[chunk_index], score))
|
|
|
|
return reranked_results
|