legal-eye / tau_rag /rerank /chain.py
Legal-i's picture
Initial deploy: legal-eye Hebrew legal RAG (17K corpus, verbatim-from-precedent)
3be54c6 verified
"""ChainReranker — runs several rerankers in sequence."""
from __future__ import annotations
from typing import List
from ..core.types import Query, Retrieved
class ChainReranker:
"""Apply multiple rerankers one after another."""
name = "chain"
def __init__(self, stages):
self.stages = list(stages)
def rerank(self, query: Query, candidates: List[Retrieved], k: int) -> List[Retrieved]:
cur = candidates
for s in self.stages[:-1]:
cur = s.rerank(query, cur, max(k, len(cur)))
if self.stages:
cur = self.stages[-1].rerank(query, cur, k)
return cur[:k]