324 lines
12 KiB
TypeScript
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
|
|
}
|