Spaces:
Sleeping
Sleeping
File size: 1,832 Bytes
363ca2b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | from sentence_transformers import CrossEncoder
from typing import List, Dict, Any
from src.config import config
class Reranker:
def __init__(self):
self.model = None # Lazy loading
def _get_model(self) -> CrossEncoder:
if self.model is None:
# Load CrossEncoder model
self.model = CrossEncoder(config.cross_encoder_model_name)
return self.model
def rerank(self, query: str, candidates: List[Dict[str, Any]], top_n: int = None) -> List[Dict[str, Any]]:
"""
Reranks a list of candidate chunks against the query using a Cross-Encoder.
Args:
query: The user's search query.
candidates: A list of dicts, each containing a Chunk object under the key "chunk".
top_n: Number of final chunks to return. Defaults to config.top_n_final.
Returns:
Sorted list of candidates with a 'cross_score' key, descending.
"""
if not candidates:
return []
top_n = top_n or config.top_n_final
model = self._get_model()
# Build pairs: (query, text)
pairs = [[query, item["chunk"].text] for item in candidates]
# Predict relevance scores
scores = model.predict(pairs, show_progress_bar=False)
# Add score to each candidate
reranked = []
for i, item in enumerate(candidates):
item_copy = item.copy()
# Map score to float (predict returns numpy array of floats)
item_copy["cross_score"] = float(scores[i])
reranked.append(item_copy)
# Sort descending by cross-encoder score
reranked.sort(key=lambda x: x["cross_score"], reverse=True)
# Cap at top_n
return reranked[:top_n]
|