project-nomad/admin/app/utils/eval/retrieval_metrics.ts

324 lines
12 KiB
TypeScript

/**
* Information-retrieval metrics for the RAG eval harness.
*
* Pure functions over ranked results — no I/O, no models, no clock. Everything
* here is deterministic, which is the point: a change in these numbers is
* unambiguously a change in the code, never in the weather.
*
* ## Two levels of measurement, and why both
*
* Retrieval returns *chunks*, but a golden answer lives in a *document*. The
* two questions we care about are different, so they are measured differently:
*
* - **Document level** (recall, MRR, nDCG): "did the answer's document make it
* into the context at all, and how near the top?" Chunks are collapsed to
* their document, keeping each document's best rank.
* - **Chunk level** (precision): "how much of what we injected is noise?" This
* is deliberately *not* deduped — five chunks from one irrelevant document
* cost five slots of a small model's context and should be counted five
* times.
*
* Reporting only one of these hides a real failure mode in the other.
*/
/** One retrieved chunk, reduced to what scoring needs. */
export type ScoredChunk = {
/** Corpus document this chunk came from; null if it could not be resolved. */
docId: string | null
/** The pipeline's final (post-rerank) score. */
score: number
/** The raw pre-rerank cosine score, when available. */
semanticScore?: number
}
export type RetrievalCase = {
id: string
tags: string[]
/** Ranked chunks, best first, exactly as the pipeline returned them. */
retrieved: ScoredChunk[]
/** Documents that genuinely answer the question. */
relevantDocIds: string[]
/** True when the right behaviour is to retrieve nothing useful. */
expectRefusal: boolean
}
/**
* Collapse a ranked chunk list to a ranked document list, keeping each
* document's best (earliest) position. Chunks with no resolvable document are
* dropped — they still occupy a rank in the chunk-level metrics, but they
* cannot be credited to any document.
*/
export function toDocumentRanking(retrieved: ScoredChunk[]): string[] {
const seen = new Set<string>()
const ranking: string[] = []
for (const chunk of retrieved) {
if (!chunk.docId || seen.has(chunk.docId)) continue
seen.add(chunk.docId)
ranking.push(chunk.docId)
}
return ranking
}
/**
* Fraction of the relevant documents that appear within the top `k` chunks.
*
* Returns null when there are no relevant documents (an out-of-corpus case) —
* recall is undefined there, and returning 0 would drag the mean down for
* questions that are *supposed* to retrieve nothing.
*/
export function recallAtK(retrieved: ScoredChunk[], relevantDocIds: string[], k: number): number | null {
if (relevantDocIds.length === 0) return null
const relevant = new Set(relevantDocIds)
const found = new Set<string>()
for (const chunk of retrieved.slice(0, k)) {
if (chunk.docId && relevant.has(chunk.docId)) found.add(chunk.docId)
}
return found.size / relevant.size
}
/**
* 1 if *any* relevant document appears in the top `k`, else 0.
*
* Distinct from recall on multi-hop questions: finding one of two required
* documents is a hit but only 0.5 recall. Reporting both is what separates
* "found something" from "found enough to answer".
*/
export function hitRateAtK(retrieved: ScoredChunk[], relevantDocIds: string[], k: number): number | null {
if (relevantDocIds.length === 0) return null
const relevant = new Set(relevantDocIds)
return retrieved.slice(0, k).some((c) => c.docId && relevant.has(c.docId)) ? 1 : 0
}
/**
* Fraction of the top `k` *chunks* that come from a relevant document.
*
* Not deduped, on purpose — this measures context pollution, and a small model
* drowning in four irrelevant chunks does not care that they share a source.
* Denominator is min(k, retrieved.length) so a pipeline returning 2 good chunks
* is not punished for the 3 it correctly declined to return.
*/
export function precisionAtK(retrieved: ScoredChunk[], relevantDocIds: string[], k: number): number | null {
if (relevantDocIds.length === 0) return null
const window = retrieved.slice(0, k)
if (window.length === 0) return 0
const relevant = new Set(relevantDocIds)
const hits = window.filter((c) => c.docId && relevant.has(c.docId)).length
return hits / window.length
}
/**
* Reciprocal of the rank of the first relevant chunk (1-indexed); 0 if none.
* Averaged over cases this is MRR.
*/
export function reciprocalRank(retrieved: ScoredChunk[], relevantDocIds: string[]): number | null {
if (relevantDocIds.length === 0) return null
const relevant = new Set(relevantDocIds)
const idx = retrieved.findIndex((c) => c.docId && relevant.has(c.docId))
return idx === -1 ? 0 : 1 / (idx + 1)
}
/**
* Normalized discounted cumulative gain at `k`, over the **document** ranking,
* with binary gain.
*
* DCG@k = Σ_{i=1..k} rel_i / log2(i + 1)
* IDCG@k = Σ_{i=1..min(R, k)} 1 / log2(i + 1) where R = |relevant|
* nDCG@k = DCG@k / IDCG@k
*
* The ideal ranking is defined against the *known* number of relevant documents
* rather than against however many happened to be retrieved. That distinction
* matters: normalizing against the retrieved set would score a run that found
* one of three required documents, and ranked it first, as a perfect 1.0.
*
* This is the metric that catches "right documents, wrong order" — a reranking
* regression that leaves recall untouched while pushing the answer to position
* five, where a 1B model with a 2-result budget will never see it.
*/
export function ndcgAtK(retrieved: ScoredChunk[], relevantDocIds: string[], k: number): number | null {
if (relevantDocIds.length === 0) return null
const relevant = new Set(relevantDocIds)
const docRanking = toDocumentRanking(retrieved).slice(0, k)
let dcg = 0
docRanking.forEach((docId, i) => {
if (relevant.has(docId)) dcg += 1 / Math.log2(i + 2) // i is 0-indexed, so rank = i + 1
})
let idcg = 0
for (let i = 0; i < Math.min(relevant.size, k); i++) {
idcg += 1 / Math.log2(i + 2)
}
return idcg === 0 ? 0 : dcg / idcg
}
/**
* Score separation: the gap between what relevant and irrelevant chunks score.
*
* This is how the similarity threshold gets calibrated with evidence instead of
* intuition. If the relevant p10 sits below the irrelevant p90, no threshold
* can cleanly separate them and the honest conclusion is that the *retriever*
* needs work, not the cutoff.
*/
export type ScoreDistribution = {
count: number
min: number
p10: number
median: number
p90: number
max: number
mean: number
}
export function describeScores(scores: number[]): ScoreDistribution | null {
if (scores.length === 0) return null
const sorted = [...scores].sort((a, b) => a - b)
const pct = (p: number) => {
// Nearest-rank percentile: no interpolation, so a reported value is always
// a value that actually occurred.
const idx = Math.min(sorted.length - 1, Math.max(0, Math.ceil((p / 100) * sorted.length) - 1))
return sorted[idx]
}
return {
count: sorted.length,
min: sorted[0],
p10: pct(10),
median: pct(50),
p90: pct(90),
max: sorted[sorted.length - 1],
mean: sorted.reduce((a, b) => a + b, 0) / sorted.length,
}
}
/** Mean over the entries that are actually defined; null if none are. */
export function meanOf(values: Array<number | null>): number | null {
const defined = values.filter((v): v is number => v !== null)
if (defined.length === 0) return null
return defined.reduce((a, b) => a + b, 0) / defined.length
}
export type RetrievalCaseResult = {
id: string
tags: string[]
retrievedCount: number
relevantDocIds: string[]
retrievedDocIds: string[]
expectRefusal: boolean
recall: Record<number, number | null>
hitRate: Record<number, number | null>
precision: Record<number, number | null>
ndcg: Record<number, number | null>
reciprocalRank: number | null
/** True when this case retrieved nothing at all. */
empty: boolean
}
export type RetrievalAggregate = {
cases: number
/** Cases with at least one relevant document (i.e. excluding refusal cases). */
answerable: number
recall: Record<number, number | null>
hitRate: Record<number, number | null>
precision: Record<number, number | null>
ndcg: Record<number, number | null>
mrr: number | null
/**
* Fraction of *answerable* questions that retrieved nothing. A non-zero value
* means the score threshold is filtering out real answers.
*/
emptyRateOnAnswerable: number | null
/**
* Fraction of *out-of-corpus* questions that retrieved something anyway. This
* is the other half of the threshold trade-off: context handed to the model
* for a question the corpus cannot answer is exactly what produces a
* confident, wrong reply.
*/
nonEmptyRateOnRefusal: number | null
relevantScores: ScoreDistribution | null
irrelevantScores: ScoreDistribution | null
}
export const DEFAULT_K_VALUES = [1, 3, 5, 10]
export function scoreCase(c: RetrievalCase, kValues: number[] = DEFAULT_K_VALUES): RetrievalCaseResult {
const byK = <T>(fn: (k: number) => T): Record<number, T> =>
Object.fromEntries(kValues.map((k) => [k, fn(k)]))
return {
id: c.id,
tags: c.tags,
retrievedCount: c.retrieved.length,
relevantDocIds: c.relevantDocIds,
retrievedDocIds: toDocumentRanking(c.retrieved),
expectRefusal: c.expectRefusal,
recall: byK((k) => recallAtK(c.retrieved, c.relevantDocIds, k)),
hitRate: byK((k) => hitRateAtK(c.retrieved, c.relevantDocIds, k)),
precision: byK((k) => precisionAtK(c.retrieved, c.relevantDocIds, k)),
ndcg: byK((k) => ndcgAtK(c.retrieved, c.relevantDocIds, k)),
reciprocalRank: reciprocalRank(c.retrieved, c.relevantDocIds),
empty: c.retrieved.length === 0,
}
}
export function aggregate(
cases: RetrievalCase[],
results: RetrievalCaseResult[],
kValues: number[] = DEFAULT_K_VALUES
): RetrievalAggregate {
const answerable = results.filter((r) => !r.expectRefusal)
const refusals = results.filter((r) => r.expectRefusal)
const byK = (pick: (r: RetrievalCaseResult, k: number) => number | null): Record<number, number | null> =>
Object.fromEntries(kValues.map((k) => [k, meanOf(results.map((r) => pick(r, k)))]))
// Split every retrieved chunk's score by whether its document was relevant.
// Refusal cases contribute only to the irrelevant side — by definition
// nothing they retrieve is relevant, and that is precisely the population a
// threshold needs to exclude.
const relevantScores: number[] = []
const irrelevantScores: number[] = []
for (const c of cases) {
const relevant = new Set(c.relevantDocIds)
for (const chunk of c.retrieved) {
const bucket = chunk.docId && relevant.has(chunk.docId) ? relevantScores : irrelevantScores
bucket.push(chunk.semanticScore ?? chunk.score)
}
}
return {
cases: results.length,
answerable: answerable.length,
recall: byK((r, k) => r.recall[k]),
hitRate: byK((r, k) => r.hitRate[k]),
precision: byK((r, k) => r.precision[k]),
ndcg: byK((r, k) => r.ndcg[k]),
mrr: meanOf(results.map((r) => r.reciprocalRank)),
emptyRateOnAnswerable:
answerable.length === 0 ? null : answerable.filter((r) => r.empty).length / answerable.length,
nonEmptyRateOnRefusal:
refusals.length === 0 ? null : refusals.filter((r) => !r.empty).length / refusals.length,
relevantScores: describeScores(relevantScores),
irrelevantScores: describeScores(irrelevantScores),
}
}
/** Aggregate restricted to cases carrying a given tag, for per-slice reporting. */
export function aggregateByTag(
cases: RetrievalCase[],
results: RetrievalCaseResult[],
kValues: number[] = DEFAULT_K_VALUES
): Record<string, RetrievalAggregate> {
const tags = new Set(results.flatMap((r) => r.tags))
const byId = new Map(cases.map((c) => [c.id, c]))
const out: Record<string, RetrievalAggregate> = {}
for (const tag of [...tags].sort()) {
const tagged = results.filter((r) => r.tags.includes(tag))
out[tag] = aggregate(
tagged.map((r) => byId.get(r.id)!).filter(Boolean),
tagged,
kValues
)
}
return out
}