1+ """Reranker utilities for RAG chunk reranking.
2+
3+ This module contains functionality for reranking RAG chunks using cross-encoder models
4+ to improve the relevance of retrieved documents in RAG applications.
5+ """
6+
7+ import asyncio
8+ from typing import Any
9+
10+ import constants
11+ from configuration import configuration
12+ from log import get_logger
13+ from models .common .turn_summary import RAGChunk
14+
15+ logger = get_logger (__name__ )
16+
17+ # Lazy-loaded cross-encoder models for reranking RAG chunks (CPU-bound, use in thread).
18+ # Cache models by name to avoid reloading the same model multiple times.
19+ # Not a constant; pylint invalid-name is disabled for this module-level singleton.
20+ _cross_encoder_models : dict [str , Any ] = {} # pylint: disable=invalid-name
21+ _cross_encoder_load_lock = asyncio .Lock ()
22+
23+
24+ async def _get_cross_encoder (model_name : str ) -> Any :
25+ """Return the lazy-loaded cross-encoder model for reranking.
26+
27+ Args:
28+ model_name: Name of the cross-encoder model to load.
29+
30+ Returns:
31+ Loaded CrossEncoder model instance, or None if loading fails.
32+ """
33+ # Check if reranking is enabled before attempting to load the model
34+ if not configuration .reranker .enabled :
35+ logger .debug ("Reranker is disabled, not loading cross-encoder model" )
36+ return None
37+
38+ if model_name in _cross_encoder_models :
39+ return _cross_encoder_models [model_name ]
40+ async with _cross_encoder_load_lock :
41+ if model_name in _cross_encoder_models :
42+ return _cross_encoder_models [model_name ]
43+ try :
44+ from sentence_transformers import (
45+ CrossEncoder ,
46+ ) # pylint: disable=import-outside-toplevel
47+
48+ model = await asyncio .to_thread (CrossEncoder , model_name )
49+ _cross_encoder_models [model_name ] = model
50+ logger .info ("Loaded cross-encoder for RAG reranking: %s" , model_name )
51+ except Exception as e : # pylint: disable=broad-exception-caught
52+ logger .warning (
53+ "Could not load cross-encoder for reranking (%s): %s" , model_name , e
54+ )
55+ _cross_encoder_models [model_name ] = None
56+ return _cross_encoder_models [model_name ]
57+
58+
59+ # pylint: disable=too-many-locals,too-many-branches
60+ async def rerank_chunks_with_cross_encoder (
61+ query : str ,
62+ chunks : list [RAGChunk ],
63+ top_k : int ,
64+ ) -> list [RAGChunk ]:
65+ """Rerank chunks using configurable cross-encoder model.
66+
67+ Args:
68+ query: The search query
69+ chunks: RAG chunks to rerank (should contain original weighted scores)
70+ top_k: Number of top chunks to return
71+
72+ Returns:
73+ Top top_k chunks sorted by combined cross-encoder and weighted score (descending)
74+ """
75+ if not chunks :
76+ return []
77+
78+ try :
79+ # Get the cached cross-encoder model
80+ model_name = constants .DEFAULT_CROSS_ENCODER_MODEL
81+ model = await _get_cross_encoder (model_name )
82+ if model is None :
83+ raise RuntimeError (f"Failed to load cross-encoder model: { model_name } )
84+
85+ logger .debug ("Using cross-encoder model: %s" , model_name )
86+
87+ # Create query-chunk pairs for scoring
88+ pairs = [(query , chunk .content ) for chunk in chunks ]
89+ scores = await asyncio .to_thread (model .predict , pairs )
90+
91+ if hasattr (scores , "tolist" ):
92+ scores = scores .tolist ()
93+
94+ # Normalize cross-encoder scores to [0,1] range using min-max normalization
95+ if len (scores ) > 1 :
96+ min_score = min (scores )
97+ max_score = max (scores )
98+ score_range = max_score - min_score
99+ if score_range > 0 :
100+ normalized_ce_scores = [
101+ (score - min_score ) / score_range for score in scores
102+ ]
103+ else :
104+ # All scores are identical, assign 0.5 to all
105+ normalized_ce_scores = [0.5 ] * len (scores )
106+ else :
107+ # Single score, assign 1.0
108+ normalized_ce_scores = [1.0 ] * len (scores )
109+
110+ # Extract original weighted scores and normalize them
111+ original_scores = [
112+ chunk .score if chunk .score is not None else 0.0 for chunk in chunks
113+ ]
114+
115+ if len (original_scores ) > 1 :
116+ min_orig = min (original_scores )
117+ max_orig = max (original_scores )
118+ orig_range = max_orig - min_orig
119+ if orig_range > 0 :
120+ normalized_orig_scores = [
121+ (score - min_orig ) / orig_range for score in original_scores
122+ ]
123+ else :
124+ # All original scores identical, assign 0.5 to all
125+ normalized_orig_scores = [0.5 ] * len (original_scores )
126+ else :
127+ # Single score, assign 1.0
128+ normalized_orig_scores = [1.0 ] * len (original_scores )
129+
130+ # Combine cross-encoder scores with original weighted scores
131+ # (favor original weighted scores)
132+ # This ensures score multipliers are still influential in the final ranking
133+ # Weight: 30% cross-encoder, 70% original weighted scores
134+ combined_scores = [
135+ (0.3 * ce_score + 0.7 * orig_score )
136+ for ce_score , orig_score in zip (
137+ normalized_ce_scores , normalized_orig_scores , strict = True
138+ )
139+ ]
140+
141+ # Combine scores with chunks and sort by combined score (descending)
142+ indexed = list (zip (combined_scores , chunks , strict = True ))
143+ indexed .sort (key = lambda x : x [0 ], reverse = True )
144+ top_indexed = indexed [:top_k ]
145+
146+ # Log the score combination results
147+ logger .info (
148+ "Cross-encoder scoring completed: combined %d cross-encoder + "
149+ "original scores (30%%/70%% mix), returning top %d chunks" ,
150+ len (chunks ),
151+ len (top_indexed ),
152+ )
153+ if logger .isEnabledFor (10 ): # DEBUG level
154+ for i , (score , chunk ) in enumerate (top_indexed [:3 ]): # Show top 3
155+ logger .debug (
156+ "Reranked chunk %d: source=%s, combined_score=%.3f, content_preview='%.50s...'" ,
157+ i + 1 ,
158+ chunk .source ,
159+ score ,
160+ chunk .content ,
161+ )
162+
163+ # Return RAGChunk list with combined scores
164+ return [
165+ RAGChunk (
166+ content = chunk .content ,
167+ source = chunk .source ,
168+ score = float (score ),
169+ attributes = chunk .attributes ,
170+ )
171+ for score , chunk in top_indexed
172+ ]
173+
174+ except Exception as e : # pylint: disable=broad-exception-caught
175+ logger .warning (
176+ "Cross-encoder reranking failed, falling back to original scoring: %s" , e
177+ )
178+ # Fallback: sort by original score and take top_k
179+ sorted_chunks = sorted (
180+ chunks ,
181+ key = lambda c : c .score if c .score is not None else float ("-inf" ),
182+ reverse = True ,
183+ )
184+ return sorted_chunks [:top_k ]
185+
186+
187+ def apply_byok_rerank_boost (
188+ chunks : list [RAGChunk ], boost : float = constants .BYOK_RAG_RERANK_BOOST
189+ ) -> list [RAGChunk ]:
190+ """Apply a score multiplier to BYOK chunks (source != OKP) and re-sort by score.
191+
192+ Args:
193+ chunks: RAG chunks after reranking (may be from BYOK or Solr).
194+ boost: Multiplier applied to BYOK chunk scores. Solr chunks unchanged.
195+
196+ Returns:
197+ Same chunks with BYOK scores boosted, sorted by score descending.
198+ """
199+ boosted = []
200+ for chunk in chunks :
201+ score = chunk .score if chunk .score is not None else float ("-inf" )
202+ if chunk .source != constants .OKP_RAG_ID :
203+ score = score * boost
204+ boosted .append (
205+ RAGChunk (
206+ content = chunk .content ,
207+ source = chunk .source ,
208+ score = score ,
209+ attributes = chunk .attributes ,
210+ )
211+ )
212+ boosted .sort (
213+ key = lambda c : c .score if c .score is not None else float ("-inf" ),
214+ reverse = True ,
215+ )
216+ return boosted
0 commit comments