Spaces:
Running
Running
File size: 4,756 Bytes
6c3af4e 1a6d0d2 6c3af4e | 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 | // @vitest-environment node
/**
* Offline retrieval eval: runs the real reranker through rankSearchResults on
* the fixed golden set and scores the ranking with nDCG and recall against the
* human-labeled relevant results.
*
* It calls rankSearchResults with preserveTopResults=true, exactly the way the
* app does for text search (searchEndpointServerHook.ts), so the pinned-first
* result branch is exercised. The golden set therefore includes entries whose
* first candidate is irrelevant: that is the case where the pin costs the
* ranking, and without it the eval could not tell a good pin from a bad one.
*
* This is the regression signal for changes to the reranker or to
* rankSearchResults (the score filter, the top-result preservation, the sort).
* Golden-set structure is validated in goldenSet.test.ts (default suite).
*
* npx vitest run --config vitest.eval.config.ts retrieval
*
* Loads the real ONNX model, so it is excluded from the default suite and from
* CI (which has no model on disk); it is a local signal, matching the server
* integration-test house style.
*/
import { afterAll, beforeAll, describe, expect, it } from "vitest";
import { rankSearchResults } from "../server/rankSearchResults.ts";
import {
startRerankerService,
stopRerankerService,
} from "../server/rerankerService.ts";
import { goldenQueries } from "./goldenSet.ts";
import { ndcgAtK, recallAtK } from "./metrics.ts";
import { K, MIN_MEAN_NDCG, MIN_MEAN_RECALL } from "./thresholds.ts";
// The floors (MIN_MEAN_NDCG / MIN_MEAN_RECALL) and the no-op ceilings that
// keep them falsifiable live in thresholds.ts; see that file for the band
// they're set in and why the headroom is intentional.
interface PerQueryScore {
id: string;
ndcg: number;
recall: number;
}
describe("retrieval eval (real reranker)", () => {
const scores: PerQueryScore[] = [];
beforeAll(async () => {
await startRerankerService();
}, 600_000);
afterAll(async () => {
await stopRerankerService();
});
for (const golden of goldenQueries) {
it(`scores ${golden.id}`, async () => {
const candidates: [title: string, content: string, url: string][] =
golden.results.map(({ title, snippet, url }) => [title, snippet, url]);
// preserveTopResults=true matches the app's text-search call.
const ranked = await rankSearchResults(
golden.query,
"text",
candidates,
true,
);
expect(ranked.length, `${golden.id}: empty ranking`).toBeGreaterThan(0);
// The pin is the point of the [1,3] entries: the first candidate must
// stay first regardless of the model's scores. A floor-only assert can't
// catch a broken pin (it would raise the mean), so pin it directly.
expect(
ranked[0][2],
`${golden.id}: preserveTopResults did not pin the first candidate`,
).toBe(golden.results[0].url);
const rankedUrls = ranked.map(([, , url]) => url);
const relevantUrls = golden.relevant.map((i) => golden.results[i].url);
const ndcg = ndcgAtK(rankedUrls, relevantUrls, K);
const recall = recallAtK(rankedUrls, relevantUrls, K);
// Bounds that also catch a metric regression (e.g. nDCG > 1 from a
// duplicate url).
expect(ndcg, `${golden.id}: nDCG out of range`).toBeGreaterThanOrEqual(0);
expect(ndcg, `${golden.id}: nDCG out of range`).toBeLessThanOrEqual(1);
expect(
recall,
`${golden.id}: recall out of range`,
).toBeGreaterThanOrEqual(0);
expect(recall, `${golden.id}: recall out of range`).toBeLessThanOrEqual(
1,
);
scores.push({ id: golden.id, ndcg, recall });
}, 120_000);
}
it("meets the aggregate regression thresholds", () => {
// Print the breakdown before asserting, so a per-query failure still shows
// which queries scored what (the diagnostic you need when a regression fires).
const meanNdcg =
scores.length > 0
? scores.reduce((sum, s) => sum + s.ndcg, 0) / scores.length
: 0;
const meanRecall =
scores.length > 0
? scores.reduce((sum, s) => sum + s.recall, 0) / scores.length
: 0;
console.table(
scores.map((s) => ({
id: s.id,
[`nDCG@${K}`]: s.ndcg.toFixed(3),
[`recall@${K}`]: s.recall.toFixed(3),
})),
);
console.log(
`mean nDCG@${K}=${meanNdcg.toFixed(3)} mean recall@${K}=${meanRecall.toFixed(3)}`,
);
expect(scores).toHaveLength(goldenQueries.length);
expect(meanNdcg, `mean nDCG@${K} regressed`).toBeGreaterThanOrEqual(
MIN_MEAN_NDCG,
);
expect(meanRecall, `mean recall@${K} regressed`).toBeGreaterThanOrEqual(
MIN_MEAN_RECALL,
);
});
});
|