File size: 641 Bytes
3be54c6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | """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]
|