Coverage for src/lilbee/retrieval/reranker.py: 100%

1"""Cross-encoder reranking for search results. 

2 

3Optional precision pass that scores each (query, chunk) pair through the 

4active provider's ``rerank`` method. Only active when 

5``cfg.reranker_model`` is set. 

6 

7Core technique: Nogueira & Cho 2019, "Passage Re-ranking with BERT" 

8(https://arxiv.org/abs/1901.04085). 

9 

10Position-aware blending: derived from learning-to-rank literature 

11(Burges et al. 2005). Top positions trust hybrid fusion more, lower 

12positions trust the reranker more. 

13""" 

14 

15from __future__ import annotations 

16 

17import logging 

18from typing import NamedTuple 

19 

20from lilbee.core.config import Config 

21from lilbee.data.store import SearchChunk 

22from lilbee.retrieval.query.dedup import fusion_norms as compute_fusion_norms 

23from lilbee.retrieval.query.dedup import normalize_scores 

24 

25log = logging.getLogger(__name__) 

26 

27 

28class ScoredChunk(NamedTuple): 

29 """A search chunk paired with its blended score.""" 

30 

31 score: float 

32 chunk: SearchChunk 

33 

34 

35_TOP_POSITION_CUTOFF = 3 

36_MID_POSITION_CUTOFF = 10 

37 

38_BLEND_SCHEDULE = { 

39 "top": (0.70, 0.30), 

40 "mid": (0.50, 0.50), 

41 "bottom": (0.30, 0.70), 

42} 

43 

44 

45def _blend_scores( 

46 to_rerank: list[SearchChunk], norm_scores: list[float], fusion_norms: list[float] 

47) -> list[ScoredChunk]: 

48 """Blend fusion scores with reranker scores using position-aware weights. 

49 

50 Both inputs are already min-max normalized to [0, 1] (``fusion_norms`` 

51 across the pool's canonical scores, ``norm_scores`` across the reranker 

52 scores), so the blend weights compare like with like. Each chunk is 

53 copied with ``rerank_score`` set to its blended score; the input chunks 

54 are left untouched. 

55 """ 

56 blended: list[ScoredChunk] = [] 

57 for i, (chunk, rerank_score, fusion_norm) in enumerate( 

58 zip(to_rerank, norm_scores, fusion_norms, strict=True) 

59 ): 

60 if i < _TOP_POSITION_CUTOFF: 

61 fw, rw = _BLEND_SCHEDULE["top"] 

62 elif i < _MID_POSITION_CUTOFF: 

63 fw, rw = _BLEND_SCHEDULE["mid"] 

64 else: 

65 fw, rw = _BLEND_SCHEDULE["bottom"] 

66 

67 final_score = fw * fusion_norm + rw * rerank_score 

68 scored = chunk.model_copy(update={"rerank_score": final_score}) 

69 blended.append(ScoredChunk(final_score, scored)) 

70 return blended 

71 

72 

73class Reranker: 

74 """Cross-encoder reranker with position-aware blending. 

75 

76 Delegates scoring to the active provider's ``rerank``; blends the result with 

77 the normalized retrieval fusion signal so a confident hybrid hit keeps its 

78 standing against a reranker that favours a weaker chunk (Nogueira & Cho 2019, 

79 https://arxiv.org/abs/1901.04085). 

80 """ 

81 

82 def __init__(self, config: Config) -> None: 

83 self._config = config 

84 

85 def rerank( 

86 self, 

87 query: str, 

88 results: list[SearchChunk], 

89 candidates: int | None = None, 

90 ) -> list[SearchChunk]: 

91 """Rerank search results through the provider's ``rerank`` method.""" 

92 if not self._config.reranker_model: 

93 return results 

94 if candidates is None: 

95 candidates = self._config.rerank_candidates 

96 to_rerank = results[:candidates] 

97 remainder = results[candidates:] 

98 

99 if not to_rerank: 

100 return results 

101 

102 scores = _score_candidates(query, to_rerank) 

103 if scores is None: 

104 return results 

105 

106 floor = self._config.rerank_min_score 

107 if floor is not None: 

108 # Absolute floor on the raw provider score: min-max normalization 

109 # erases calibration, so without this a uniformly irrelevant pool 

110 # still yields a 1.0-scored "best" candidate. 

111 kept = [(s, c) for s, c in zip(scores, to_rerank, strict=True) if s >= floor] 

112 if len(kept) < len(to_rerank): 

113 log.info( 

114 "Reranker dropped %d of %d candidates below rerank_min_score", 

115 len(to_rerank) - len(kept), 

116 len(to_rerank), 

117 ) 

118 if not kept: 

119 return remainder 

120 scores = [s for s, _ in kept] 

121 to_rerank = [c for _, c in kept] 

122 

123 norm_scores = normalize_scores(scores) 

124 if self._config.rerank_blend: 

125 fusion_norms = compute_fusion_norms(to_rerank) 

126 scored = _blend_scores(to_rerank, norm_scores, fusion_norms) 

127 else: 

128 # Pure cross-encoder ordering: no fusion blend, so the reranker's 

129 # effect is unattenuated (and measurable in isolation). 

130 scored = [ 

131 ScoredChunk(s, c.model_copy(update={"rerank_score": s})) 

132 for s, c in zip(norm_scores, to_rerank, strict=True) 

133 ] 

134 scored_sorted = sorted(scored, key=lambda x: x.score, reverse=True) 

135 

136 reranked = [chunk for _, chunk in scored_sorted] 

137 return reranked + remainder 

138 

139 

140def _score_candidates(query: str, to_rerank: list[SearchChunk]) -> list[float] | None: 

141 """Call the active provider's rerank; return None on error after logging. 

142 

143 A provider that returns the wrong number of scores is contained here 

144 like any other failure: the scores cannot be paired with the candidates, 

145 so the pass is skipped and retrieval order stands. 

146 """ 

147 # circular: services -> reranker via Searcher; deferred so test-time 

148 # monkeypatching of ``lilbee.app.services.get_services`` stays effective. 

149 from lilbee.app.services import get_services 

150 

151 try: 

152 provider = get_services().provider 

153 # Title-prefixed passages: a cross-encoder cannot judge a chunk whose 

154 # relevance depends on its parent document from the bare text. 

155 scores = provider.rerank( 

156 query, [f"{c.title}\n{c.chunk}" if c.title else c.chunk for c in to_rerank] 

157 ) 

158 except Exception as exc: 

159 log.warning("Reranker failed; skipping rerank pass: %s", exc, exc_info=True) 

160 return None 

161 if len(scores) != len(to_rerank): 

162 log.warning( 

163 "Reranker returned %d scores for %d candidates; skipping rerank pass", 

164 len(scores), 

165 len(to_rerank), 

166 ) 

167 return None 

168 return scores