# Copyright 2026 Boring for Gemini Authors
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
"""
Cross-Encoder Reranker Module - V10.24
High-precision reranking using cross-encoder architecture.
Improves ranking quality by 10-15% over bi-encoder approaches.
This module provides:
1. Cross-encoder based reranking (when sentence-transformers available)
2. Lightweight heuristic reranker (fallback)
3. Ensemble reranking combining multiple signals
"""
import logging
from dataclasses import dataclass
from typing import Any
logger = logging.getLogger(__name__)
from boring.core.dependencies import DependencyManager
@dataclass
class RerankScore:
"""Score from cross-encoder reranking."""
original_score: float
rerank_score: float
combined_score: float
confidence: float
method: str # "cross_encoder", "heuristic", "ensemble"
class CrossEncoderReranker:
"""
Cross-encoder based reranker for high-precision code retrieval.
Cross-encoders process query-document pairs jointly, allowing for
deeper semantic understanding than bi-encoder approaches.
Models:
- ms-marco-MiniLM-L-6-v2: Fast, good for general text (default)
- cross-encoder/ms-marco-TinyBERT-L-2: Ultra-fast
- cross-encoder/stsb-roberta-base: Good for semantic similarity
"""
# Available models for different use cases
MODELS = {
"fast": "cross-encoder/ms-marco-TinyBERT-L-2",
"balanced": "cross-encoder/ms-marco-MiniLM-L-6-v2",
"accurate": "cross-encoder/stsb-roberta-base",
}
def __init__(self, model_name: str = "balanced", device: str = "cpu"):
"""
Initialize cross-encoder reranker.
Args:
model_name: Model preset ("fast", "balanced", "accurate") or HuggingFace model name
device: Device to use ("cpu", "cuda", "mps")
"""
self.model_name = self.MODELS.get(model_name, model_name)
self.device = device
self._model: Any | None = None
self._initialized = False
def _ensure_model(self) -> bool:
"""Lazily initialize the model."""
if self._initialized:
return self._model is not None
self._initialized = True
if not DependencyManager.check_chroma():
logger.debug(
"CrossEncoder not available (sentence-transformers missing), using heuristic reranking"
)
return False
try:
from sentence_transformers import CrossEncoder
self._model = CrossEncoder(self.model_name, device=self.device)
logger.info(f"CrossEncoder initialized: {self.model_name}")
return True
except Exception as e:
logger.warning(f"Failed to load CrossEncoder: {e}")
return False
def rerank(
self,
query: str,
documents: list[str],
original_scores: list[float],
top_k: int = 10,
weight_original: float = 0.3,
) -> list[tuple[int, RerankScore]]:
"""
Rerank documents using cross-encoder.
Args:
query: Search query
documents: List of document texts
original_scores: Original retrieval scores
top_k: Number of results to return
weight_original: Weight for original score in combination
Returns:
List of (original_index, RerankScore) sorted by combined score
"""
if not documents:
return []
if self._ensure_model():
return self._rerank_with_model(
query, documents, original_scores, top_k, weight_original
)
else:
return self._rerank_heuristic(query, documents, original_scores, top_k, weight_original)
def _rerank_with_model(
self,
query: str,
documents: list[str],
original_scores: list[float],
top_k: int,
weight_original: float,
) -> list[tuple[int, RerankScore]]:
"""Rerank using actual cross-encoder model."""
# Create query-document pairs
pairs = [(query, doc) for doc in documents]
# Get cross-encoder scores
try:
scores = self._model.predict(pairs, show_progress_bar=False)
except Exception as e:
logger.warning(f"Cross-encoder prediction failed: {e}")
return self._rerank_heuristic(query, documents, original_scores, top_k, weight_original)
# Normalize scores to [0, 1]
min_score = min(scores)
max_score = max(scores)
range_score = max_score - min_score if max_score > min_score else 1.0
normalized_scores = [(s - min_score) / range_score for s in scores]
# Combine with original scores
results = []
for i, (rerank_score, orig_score) in enumerate(
zip(normalized_scores, original_scores, strict=True)
):
combined = (1 - weight_original) * rerank_score + weight_original * orig_score
results.append(
(
i,
RerankScore(
original_score=orig_score,
rerank_score=rerank_score,
combined_score=combined,
confidence=0.9,
method="cross_encoder",
),
)
)
# Sort by combined score
results.sort(key=lambda x: x[1].combined_score, reverse=True)
return results[:top_k]
def _rerank_heuristic(
self,
query: str,
documents: list[str],
original_scores: list[float],
top_k: int,
weight_original: float,
) -> list[tuple[int, RerankScore]]:
"""Heuristic reranking when cross-encoder unavailable."""
query_lower = query.lower()
query_terms = set(query_lower.split())
results = []
for i, (doc, orig_score) in enumerate(zip(documents, original_scores, strict=True)):
doc_lower = doc.lower()
# Heuristic scoring factors
factors = []
# 1. Exact phrase match (strongest signal)
if query_lower in doc_lower:
factors.append(0.4)
# 2. Term coverage
matching_terms = sum(1 for term in query_terms if term in doc_lower)
term_coverage = matching_terms / len(query_terms) if query_terms else 0
factors.append(term_coverage * 0.3)
# 3. Position of first match (earlier is better)
first_match_pos = min(
(doc_lower.find(term) for term in query_terms if term in doc_lower),
default=len(doc_lower),
)
position_score = 1.0 - (first_match_pos / len(doc_lower)) if doc_lower else 0
factors.append(position_score * 0.2)
# 4. Length penalty (very long docs might be less relevant)
length_penalty = min(1.0, 500 / max(len(doc), 1))
factors.append(length_penalty * 0.1)
# Combine factors
rerank_score = sum(factors)
combined = (1 - weight_original) * rerank_score + weight_original * orig_score
results.append(
(
i,
RerankScore(
original_score=orig_score,
rerank_score=rerank_score,
combined_score=combined,
confidence=0.6,
method="heuristic",
),
)
)
results.sort(key=lambda x: x[1].combined_score, reverse=True)
return results[:top_k]
class EnsembleReranker:
"""
Ensemble reranker combining multiple signals.
Combines:
1. Cross-encoder/heuristic scores
2. Keyword matching
3. Code structure analysis
4. Usage patterns (from IntelligentRanker)
"""
def __init__(self, weights: dict[str, float] | None = None):
"""
Initialize ensemble reranker.
Args:
weights: Custom weights for each signal
"""
self.weights = weights or {
"semantic": 0.35,
"keyword": 0.25,
"structure": 0.20,
"usage": 0.20,
}
self.cross_encoder = CrossEncoderReranker()
def rerank(
self,
query: str,
chunks: list, # list of CodeChunk or similar
original_scores: list[float],
usage_scores: dict[str, float] | None = None,
top_k: int = 10,
) -> list[tuple[int, float]]:
"""
Rerank using ensemble of signals.
Args:
query: Search query
chunks: List of code chunks with content and metadata
original_scores: Original retrieval scores
usage_scores: Optional usage-based scores from IntelligentRanker
top_k: Number of results to return
Returns:
List of (original_index, final_score) sorted by score
"""
if not chunks:
return []
documents = [self._chunk_to_text(chunk) for chunk in chunks]
# 1. Get semantic/cross-encoder scores
ce_results = self.cross_encoder.rerank(query, documents, original_scores, top_k=len(chunks))
semantic_scores = {idx: score.rerank_score for idx, score in ce_results}
# 2. Keyword matching scores
keyword_scores = self._compute_keyword_scores(query, chunks)
# 3. Structure scores
structure_scores = self._compute_structure_scores(query, chunks)
# 4. Usage scores (if provided)
usage_scores = usage_scores or {}
# Combine all signals
final_scores = []
for i, chunk in enumerate(chunks):
chunk_id = getattr(chunk, "chunk_id", str(i))
combined = (
self.weights["semantic"] * semantic_scores.get(i, 0)
+ self.weights["keyword"] * keyword_scores.get(i, 0)
+ self.weights["structure"] * structure_scores.get(i, 0)
+ self.weights["usage"] * usage_scores.get(chunk_id, 0.5)
)
final_scores.append((i, combined))
final_scores.sort(key=lambda x: x[1], reverse=True)
return final_scores[:top_k]
def _chunk_to_text(self, chunk) -> str:
"""Convert chunk to searchable text."""
parts = []
if hasattr(chunk, "name") and chunk.name:
parts.append(f"Name: {chunk.name}")
if hasattr(chunk, "signature") and chunk.signature:
parts.append(f"Signature: {chunk.signature}")
if hasattr(chunk, "docstring") and chunk.docstring:
parts.append(f"Doc: {chunk.docstring}")
if hasattr(chunk, "content") and chunk.content:
# Limit content length
content = chunk.content[:500]
parts.append(f"Code: {content}")
return "\n".join(parts)
def _compute_keyword_scores(self, query: str, chunks: list) -> dict[int, float]:
"""Compute keyword matching scores."""
query_terms = set(query.lower().split())
scores = {}
for i, chunk in enumerate(chunks):
text = self._chunk_to_text(chunk).lower()
matching = sum(1 for term in query_terms if term in text)
scores[i] = matching / len(query_terms) if query_terms else 0
return scores
def _compute_structure_scores(self, query: str, chunks: list) -> dict[int, float]:
"""Compute code structure relevance scores."""
scores = {}
query_lower = query.lower()
# Determine expected chunk types from query
expected_types = []
if any(kw in query_lower for kw in ["class", "model", "schema"]):
expected_types.append("class")
if any(kw in query_lower for kw in ["function", "method", "def"]):
expected_types.extend(["function", "method"])
if any(kw in query_lower for kw in ["test", "spec"]):
expected_types.append("function") # Tests are usually functions
for i, chunk in enumerate(chunks):
score = 0.5 # Default
chunk_type = getattr(chunk, "chunk_type", "")
# Boost for matching chunk type
if expected_types and chunk_type in expected_types:
score += 0.3
# Boost for docstring presence
if hasattr(chunk, "docstring") and chunk.docstring:
score += 0.1
# Boost for reasonable size
if hasattr(chunk, "content"):
length = len(chunk.content)
if 50 < length < 500: # Sweet spot
score += 0.1
scores[i] = min(1.0, score)
return scores
# Singleton instances
_cross_encoder_reranker: CrossEncoderReranker | None = None
_ensemble_reranker: EnsembleReranker | None = None
def get_cross_encoder_reranker(model: str = "balanced") -> CrossEncoderReranker:
"""Get or create cross-encoder reranker singleton."""
global _cross_encoder_reranker
if _cross_encoder_reranker is None:
_cross_encoder_reranker = CrossEncoderReranker(model_name=model)
return _cross_encoder_reranker
def get_ensemble_reranker() -> EnsembleReranker:
"""Get or create ensemble reranker singleton."""
global _ensemble_reranker
if _ensemble_reranker is None:
_ensemble_reranker = EnsembleReranker()
return _ensemble_reranker
