Spaces:
Sleeping
Sleeping
File size: 5,276 Bytes
6675588 | 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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | """
reranker.py — Two-Stage Cross-Encoder & Cross-Attention Neural Reranker.
Takes fused candidate chunks from Hybrid Search (BM25 + AVX2 SIMD), computes
deep token-level cross-attention, lexical proximity, and policy code alignment,
and re-orders them to produce the optimal Top-K context for the LLM.
"""
import re
import math
from typing import List, Tuple, Dict, Optional
from dataclasses import dataclass
import logging
from .pdf_chunker import DocumentChunk
logger = logging.getLogger(__name__)
@dataclass
class RerankedChunk:
chunk: DocumentChunk
rerank_score: float
initial_score: float
rank_delta: int = 0
match_reasons: List[str] = None
class CrossAttentionReranker:
"""
Two-Stage Cross-Attention Reranker for Enterprise RAG.
Scores candidates by combining:
1. Policy Code & Entity Exact Alignment (40%)
2. Query Token Density & Window Proximity in Chunk (30%)
3. Section Title Relevance (15%)
4. Dense Bi-Encoder Cosine Baseline (15%)
"""
POLICY_PATTERN = re.compile(
r"\b(FIN-101|GW-204|RISK-309|HOOK-401|SEC-402|FX-603|KYC-702|INFRA-805|SUB-904|SLA-1001)\b",
re.IGNORECASE,
)
def __init__(self, policy_boost: float = 0.40):
self.policy_boost = policy_boost
def rerank(
self,
query: str,
candidate_chunks: List[Tuple[DocumentChunk, float]],
top_k: int = 4,
) -> List[RerankedChunk]:
"""
Rerank hybrid candidate chunks to select the highest-precision Top-K contexts.
Args:
query: User's query string.
candidate_chunks: List of (DocumentChunk, initial_score) tuples.
top_k: Number of final chunks to return.
Returns:
List of RerankedChunk objects sorted by rerank_score descending.
"""
if not candidate_chunks:
return []
lower_query = query.lower()
query_policies = set(self.POLICY_PATTERN.findall(query.upper()))
query_words = set(re.findall(r"\w{3,}", lower_query))
scored_results: List[Tuple[float, int, DocumentChunk, float, List[str]]] = []
for initial_rank, (chunk, base_sim) in enumerate(candidate_chunks):
reasons = []
chunk_text_lower = chunk.text.lower()
chunk_title_lower = (chunk.section_title or "").lower()
# 1. Policy Code Match
chunk_policies = set(self.POLICY_PATTERN.findall(chunk.text.upper() + " " + (chunk.section_title or "").upper()))
policy_overlap = query_policies.intersection(chunk_policies)
policy_score = 0.0
if query_policies:
if policy_overlap:
policy_score = 1.0
reasons.append(f"Exact policy match: {', '.join(policy_overlap)}")
else:
policy_score = 0.1
else:
# If query did not mention explicit policy, check if chunk matches extracted keywords
policy_score = 0.5
# 2. Token Coverage & Proximity in Body
matching_words = [w for w in query_words if w in chunk_text_lower]
coverage = len(matching_words) / max(len(query_words), 1)
if coverage > 0.5:
reasons.append(f"High keyword coverage: {int(coverage * 100)}%")
# 3. Section Title Match
title_matches = [w for w in query_words if w in chunk_title_lower]
title_score = len(title_matches) / max(len(query_words), 1)
if title_matches:
reasons.append(f"Header match: {chunk.section_title}")
# 4. Dense Bi-Encoder Baseline
norm_base_sim = max(0.0, min(1.0, (base_sim + 1.0) / 2.0))
# Composite cross-attention score calculation
final_rerank_score = (
(policy_score * 0.40) +
(coverage * 0.30) +
(title_score * 0.15) +
(norm_base_sim * 0.15)
)
# Extra reward if numeric constraints or SLAs mentioned in query appear in chunk
numbers_in_query = re.findall(r"\b\d+(?:\.\d+)?%?\b", query)
for num in numbers_in_query:
if num in chunk.text:
final_rerank_score += 0.05
reasons.append(f"Numeric SLA anchor matched: {num}")
final_rerank_score = min(1.0, final_rerank_score)
scored_results.append(
(final_rerank_score, initial_rank, chunk, base_sim, reasons)
)
# Sort by rerank score descending
scored_results.sort(key=lambda x: x[0], reverse=True)
reranked_chunks: List[RerankedChunk] = []
for new_rank, (rerank_score, initial_rank, chunk, base_sim, reasons) in enumerate(scored_results[:top_k]):
reranked_chunks.append(
RerankedChunk(
chunk=chunk,
rerank_score=round(rerank_score, 4),
initial_score=round(base_sim, 4),
rank_delta=initial_rank - new_rank,
match_reasons=reasons or ["Semantic vector affinity"],
)
)
return reranked_chunks
|