Selection
Score fusion and evidence-set selection algorithms.
"""Score fusion and evidence-set selection algorithms.
Retrieval is usually presented as a ranked-list problem, while generation
consumes a *set* of passages under a token budget. This module makes that
boundary explicit. It contains calibrated score fusion, weighted reciprocal
rank fusion, budgeted maximum-coverage selection, and diagnostics for evidence
lost between retrieval and generation.
The implementations are dependency-free and intentionally expose every
intermediate score. They are suitable for experiments and small corpora; a
production system should vectorize the same objectives and persist its fitted
calibration parameters.
"""
import math
from dataclasses import dataclass
from typing import Dict, Iterable, List, Mapping, Optional, Sequence, Set, Tuple
from .models import Chunk, SearchResult
@dataclass(frozen=True)
class SelectionCandidate:
"""A candidate evidence unit with claims it can support and a unit cost."""
result: SearchResult
supports: Tuple[str, ...]
cost: int
def __post_init__(self) -> None:
if self.cost <= 0:
raise ValueError("candidate cost must be positive")
@dataclass(frozen=True)
class CoverageSelection:
"""Result of greedy budgeted maximum coverage."""
selected: Tuple[SelectionCandidate, ...]
covered: Tuple[str, ...]
uncovered: Tuple[str, ...]
spent: int
utility: float
@dataclass(frozen=True)
class EvidenceFlow:
"""Document-level survival through candidate, rerank, and packing stages."""
relevant: Tuple[str, ...]
retrieved: Tuple[str, ...]
reranked: Tuple[str, ...]
packed: Tuple[str, ...]
retrieval_recall: float
rerank_survival: float
pack_survival: float
end_to_end_recall: float
lost_at_retrieval: Tuple[str, ...]
lost_at_rerank: Tuple[str, ...]
lost_at_pack: Tuple[str, ...]
def min_max(values: Sequence[float]) -> Tuple[float, ...]:
"""Scale values to ``[0, 1]`` with a stable all-equal convention."""
if not values:
return ()
low, high = min(values), max(values)
if high == low:
return tuple(1.0 for _ in values)
return tuple((value - low) / (high - low) for value in values)
def z_scores(values: Sequence[float]) -> Tuple[float, ...]:
"""Population z-scores; an all-equal list maps to zero."""
if not values:
return ()
mean = sum(values) / len(values)
variance = sum((value - mean) ** 2 for value in values) / len(values)
if variance == 0.0:
return tuple(0.0 for _ in values)
deviation = math.sqrt(variance)
return tuple((value - mean) / deviation for value in values)
def _candidate_table(
rankings: Mapping[str, Sequence[SearchResult]],
) -> Tuple[Dict[str, Chunk], Dict[str, Dict[str, SearchResult]]]:
chunks: Dict[str, Chunk] = {}
by_system: Dict[str, Dict[str, SearchResult]] = {}
for system, results in rankings.items():
system_results: Dict[str, SearchResult] = {}
for result in results:
chunks[result.chunk.id] = result.chunk
system_results[result.chunk.id] = result
by_system[system] = system_results
return chunks, by_system
def calibrated_comb_sum(
rankings: Mapping[str, Sequence[SearchResult]],
k: int,
weights: Optional[Mapping[str, float]] = None,
normalization: str = "minmax",
missing_score: float = 0.0,
) -> List[SearchResult]:
"""Fuse scores after per-system calibration.
``CombSUM`` is meaningful only when component scores share a scale. The
function therefore fits a deterministic query-local transform per input
ranking. Query-local min-max or z-score normalization is useful for a lab,
but a deployment should fit calibration on held-out judgments and monitor
it after corpus/model changes.
"""
if k <= 0 or not rankings:
return []
if normalization not in {"minmax", "zscore", "none"}:
raise ValueError("normalization must be minmax, zscore, or none")
chunks, by_system = _candidate_table(rankings)
normalized: Dict[str, Dict[str, float]] = {}
for system, results in rankings.items():
raw = [result.score for result in results]
scaled = (
min_max(raw)
if normalization == "minmax"
else z_scores(raw)
if normalization == "zscore"
else tuple(raw)
)
normalized[system] = {
result.chunk.id: score for result, score in zip(results, scaled)
}
fused: List[Tuple[float, str, Dict[str, float]]] = []
for chunk_id in chunks:
score = 0.0
components: Dict[str, float] = {}
for system in rankings:
weight = 1.0 if weights is None else weights.get(system, 1.0)
value = normalized[system].get(chunk_id, missing_score)
contribution = weight * value
components[system + "_calibrated"] = value
components[system + "_weighted"] = contribution
score += contribution
raw_result = by_system[system].get(chunk_id)
if raw_result is not None:
components[system + "_raw"] = raw_result.score
fused.append((score, chunk_id, components))
fused.sort(key=lambda item: (-item[0], item[1]))
return [
SearchResult(
chunk=chunks[chunk_id],
score=score,
rank=rank,
retriever="calibrated-combsum",
component_scores=components,
)
for rank, (score, chunk_id, components) in enumerate(fused[:k], start=1)
]
def reciprocal_rank_fusion(
rankings: Mapping[str, Sequence[SearchResult]],
k: int,
constant: int = 60,
weights: Optional[Mapping[str, float]] = None,
) -> List[SearchResult]:
"""Fuse ordinal ranks without assuming comparable component scores."""
if constant <= 0:
raise ValueError("constant must be positive")
if k <= 0 or not rankings:
return []
chunks, _ = _candidate_table(rankings)
totals: Dict[str, float] = {chunk_id: 0.0 for chunk_id in chunks}
components: Dict[str, Dict[str, float]] = {chunk_id: {} for chunk_id in chunks}
for system, results in rankings.items():
weight = 1.0 if weights is None else weights.get(system, 1.0)
for ordinal, result in enumerate(results, start=1):
contribution = weight / (constant + ordinal)
totals[result.chunk.id] += contribution
components[result.chunk.id][system + "_rrf"] = contribution
ranked = sorted(totals, key=lambda chunk_id: (-totals[chunk_id], chunk_id))[:k]
return [
SearchResult(
chunk=chunks[chunk_id],
score=totals[chunk_id],
rank=rank,
retriever="rrf",
component_scores=components[chunk_id],
)
for rank, chunk_id in enumerate(ranked, start=1)
]
def greedy_budgeted_coverage(
candidates: Sequence[SelectionCandidate],
required: Iterable[str],
budget: int,
claim_weights: Optional[Mapping[str, float]] = None,
relevance_weight: float = 0.05,
) -> CoverageSelection:
"""Select evidence by marginal supported-claim utility per unit cost.
Weighted maximum coverage under a knapsack budget is NP-hard. This greedy
approximation repeatedly chooses the candidate with the highest marginal
coverage-plus-relevance utility per cost. Stable ID tie-breaking keeps
experiments reproducible.
"""
if budget < 0:
raise ValueError("budget must be non-negative")
required_set = set(required)
weights = {claim: 1.0 for claim in required_set}
if claim_weights:
for claim, weight in claim_weights.items():
if weight < 0:
raise ValueError("claim weights must be non-negative")
if claim in required_set:
weights[claim] = weight
remaining = list(candidates)
selected: List[SelectionCandidate] = []
covered: Set[str] = set()
spent = 0
utility = 0.0
while remaining:
feasible = [candidate for candidate in remaining if spent + candidate.cost <= budget]
if not feasible:
break
scored = []
for candidate in feasible:
new_claims = (set(candidate.supports) & required_set) - covered
coverage_gain = sum(weights[claim] for claim in new_claims)
relevance_gain = relevance_weight * max(0.0, candidate.result.score)
gain = coverage_gain + relevance_gain
scored.append((gain / candidate.cost, gain, candidate.result.chunk.id, candidate))
scored.sort(key=lambda item: (-item[0], -item[1], item[2]))
_, gain, _, chosen = scored[0]
if gain <= 0.0:
break
selected.append(chosen)
covered.update(set(chosen.supports) & required_set)
spent += chosen.cost
utility += gain
remaining = [item for item in remaining if item.result.chunk.id != chosen.result.chunk.id]
ordered_covered = tuple(sorted(covered))
return CoverageSelection(
selected=tuple(selected),
covered=ordered_covered,
uncovered=tuple(sorted(required_set - covered)),
spent=spent,
utility=utility,
)
def _documents(results: Sequence[SearchResult]) -> Tuple[str, ...]:
seen: Set[str] = set()
ordered: List[str] = []
for result in results:
identifier = result.chunk.document_id
if identifier not in seen:
ordered.append(identifier)
seen.add(identifier)
return tuple(ordered)
def evidence_flow(
relevant_document_ids: Iterable[str],
retrieved: Sequence[SearchResult],
reranked: Sequence[SearchResult],
packed: Sequence[SearchResult],
) -> EvidenceFlow:
"""Attribute recall loss to retrieval, reranking, or context packing."""
relevant = tuple(sorted(set(relevant_document_ids)))
relevant_set = set(relevant)
retrieved_ids = _documents(retrieved)
reranked_ids = _documents(reranked)
packed_ids = _documents(packed)
retrieved_relevant = relevant_set & set(retrieved_ids)
reranked_relevant = relevant_set & set(reranked_ids)
packed_relevant = relevant_set & set(packed_ids)
def ratio(numerator: int, denominator: int) -> float:
return numerator / denominator if denominator else 1.0
return EvidenceFlow(
relevant=relevant,
retrieved=retrieved_ids,
reranked=reranked_ids,
packed=packed_ids,
retrieval_recall=ratio(len(retrieved_relevant), len(relevant_set)),
rerank_survival=ratio(len(reranked_relevant), len(retrieved_relevant)),
pack_survival=ratio(len(packed_relevant), len(reranked_relevant)),
end_to_end_recall=ratio(len(packed_relevant), len(relevant_set)),
lost_at_retrieval=tuple(sorted(relevant_set - retrieved_relevant)),
lost_at_rerank=tuple(sorted(retrieved_relevant - reranked_relevant)),
lost_at_pack=tuple(sorted(reranked_relevant - packed_relevant)),
)
def selection_regret(oracle_utility: float, observed_utility: float) -> float:
"""Normalized utility regret, clamped to ``[0, 1]``."""
if oracle_utility < 0 or observed_utility < 0:
raise ValueError("utilities must be non-negative")
if oracle_utility == 0:
return 0.0 if observed_utility == 0 else 0.0
return min(1.0, max(0.0, (oracle_utility - observed_utility) / oracle_utility))