295 lines
10 KiB
TypeScript
295 lines
10 KiB
TypeScript
/**
|
||
* ANSWER-MERGE self-ensembling pilot (BEAM 1M).
|
||
*
|
||
* Hypothesis: BEAM per-question scores have high re-run variance and compound
|
||
* rubric nuggets reward content UNION. Merging two INDEPENDENT answers to the
|
||
* same question (same v2 config) into one, then re-judging, may beat the single
|
||
* headline answer — for free at inference time when the two answers already
|
||
* exist.
|
||
*
|
||
* PAIR SOURCES (all in benchmarks/results/beam/):
|
||
* - HEADLINE (canonical): beam-1m-FULL700-gpt5-retv2.jsonl, first row per
|
||
* instance_id → 700 canonical answers+scores+rubric.
|
||
* - ALT (independent 2nd answer under the same config):
|
||
* (a) beam-1m-FULL700-gpt5-retv2.backup.jsonl — for ids with 2+ rows whose
|
||
* answer TEXTS DIFFER, the LATER row is an independent re-answer.
|
||
* (b) beam-1m-cal-gpt5-retrieval-v2.jsonl — 50 questions answered again
|
||
* under the same config.
|
||
* A pair is kept iff both answers are non-empty, materially differ (normalized
|
||
* inequality), and are not BOTH the exact IDK sentence.
|
||
*
|
||
* MERGE (gpt-5) → JUDGE merged answer with the official nugget judge (gpt-5),
|
||
* computeTau:false so we compare the merged plain nugget-mean against the
|
||
* headline plain nugget-mean `score` (apples-to-apples; tau is a non-headline
|
||
* diagnostic and adds many costly calls).
|
||
*
|
||
* Budget-guarded ($8 hard stop). Pair order is deterministically shuffled
|
||
* (seed 42) so a partial run is still a fair sample. Resumable: already-written
|
||
* instance_ids are skipped.
|
||
*
|
||
* Run: npx tsx benchmarks/harness/scripts/_merge-pilot.ts [--budget 8]
|
||
*/
|
||
|
||
import fs from 'node:fs';
|
||
import path from 'node:path';
|
||
import { createBeamOpenAiClient, OPENAI_PRICING } from '../src/beam-openai-client.js';
|
||
import { judgeQuestion } from '../src/beam-nugget-judge.js';
|
||
|
||
const RESULTS_DIR = path.resolve('benchmarks/results/beam');
|
||
const HEADLINE = path.join(RESULTS_DIR, 'beam-1m-FULL700-gpt5-retv2.jsonl');
|
||
const BACKUP = path.join(RESULTS_DIR, 'beam-1m-FULL700-gpt5-retv2.backup.jsonl');
|
||
const CAL = path.join(RESULTS_DIR, 'beam-1m-cal-gpt5-retrieval-v2.jsonl');
|
||
const OUT = path.join(RESULTS_DIR, 'beam-1m-merge-pilot.jsonl');
|
||
|
||
const MERGE_SYSTEM =
|
||
'You are combining two draft answers to the same question, both produced from ' +
|
||
'the same memory context. Produce ONE final answer that includes ALL specific, ' +
|
||
'non-contradictory content from both drafts (names, dates, numbers, versions, ' +
|
||
'events, causes, outcomes), organized clearly, with no meta-commentary about ' +
|
||
'drafts. If the drafts disagree on a fact, state the contradiction explicitly ' +
|
||
'and present both values. If BOTH drafts say there is not enough information, ' +
|
||
"output exactly: I don't have enough information to answer this question. " +
|
||
"If only one draft has substantive content, use that draft's content.";
|
||
|
||
interface BeamRow {
|
||
instance_id: string;
|
||
conv?: number;
|
||
memory_ability: string;
|
||
question: string;
|
||
answer: string;
|
||
score: number;
|
||
nugget_scores: Array<{ nugget: string; score: number; reason: string }>;
|
||
}
|
||
|
||
function readJsonl(p: string): BeamRow[] {
|
||
return fs
|
||
.readFileSync(p, 'utf-8')
|
||
.split('\n')
|
||
.filter(l => l.trim())
|
||
.map(l => {
|
||
try {
|
||
return JSON.parse(l) as BeamRow;
|
||
} catch {
|
||
return null;
|
||
}
|
||
})
|
||
.filter((r): r is BeamRow => r !== null);
|
||
}
|
||
|
||
const norm = (t: string): string =>
|
||
(t ?? '')
|
||
.replace(/[‘’]/g, "'")
|
||
.replace(/[“”]/g, '"')
|
||
.replace(/\s+/g, ' ')
|
||
.trim()
|
||
.toLowerCase();
|
||
|
||
const IDK = "i don't have enough information to answer this question";
|
||
const isExactIDK = (t: string): boolean => {
|
||
const n = norm(t);
|
||
return n === IDK || n === `${IDK}.`;
|
||
};
|
||
const looksIDK = (t: string): boolean => norm(t).startsWith(IDK);
|
||
const isEmpty = (t: string): boolean => !t || !t.trim();
|
||
|
||
/** Deterministic PRNG (mulberry32) + Fisher-Yates for a seed-42 shuffle. */
|
||
function mulberry32(seed: number): () => number {
|
||
let a = seed >>> 0;
|
||
return () => {
|
||
a |= 0;
|
||
a = (a + 0x6d2b79f5) | 0;
|
||
let t = Math.imul(a ^ (a >>> 15), 1 | a);
|
||
t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t;
|
||
return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
|
||
};
|
||
}
|
||
function shuffle<T>(arr: T[], seed: number): T[] {
|
||
const rnd = mulberry32(seed);
|
||
const a = [...arr];
|
||
for (let i = a.length - 1; i > 0; i--) {
|
||
const j = Math.floor(rnd() * (i + 1));
|
||
[a[i], a[j]] = [a[j], a[i]];
|
||
}
|
||
return a;
|
||
}
|
||
|
||
interface Pair {
|
||
instanceId: string;
|
||
memoryAbility: string;
|
||
question: string;
|
||
headline: BeamRow;
|
||
alt: BeamRow;
|
||
altSource: 'backup' | 'cal';
|
||
}
|
||
|
||
function buildPairs(): Pair[] {
|
||
const headline = readJsonl(HEADLINE);
|
||
const headMap = new Map<string, BeamRow>();
|
||
for (const r of headline) if (!headMap.has(r.instance_id)) headMap.set(r.instance_id, r);
|
||
|
||
const usable = (head: BeamRow, alt: BeamRow): boolean => {
|
||
if (isEmpty(alt.answer) || isEmpty(head.answer)) return false;
|
||
if (norm(alt.answer) === norm(head.answer)) return false; // not materially different
|
||
if (isExactIDK(alt.answer) && isExactIDK(head.answer)) return false;
|
||
return true;
|
||
};
|
||
|
||
// (a) backup: later row among 2+ distinct-answer rows.
|
||
const backup = readJsonl(BACKUP);
|
||
const byId = new Map<string, BeamRow[]>();
|
||
for (const r of backup) {
|
||
if (!byId.has(r.instance_id)) byId.set(r.instance_id, []);
|
||
byId.get(r.instance_id)!.push(r);
|
||
}
|
||
const pairs = new Map<string, Pair>();
|
||
for (const [id, rows] of byId) {
|
||
if (rows.length < 2) continue;
|
||
if (new Set(rows.map(r => norm(r.answer))).size < 2) continue; // all identical
|
||
const head = headMap.get(id);
|
||
if (!head) continue;
|
||
const alt = rows[rows.length - 1];
|
||
if (!usable(head, alt)) continue;
|
||
pairs.set(id, {
|
||
instanceId: id,
|
||
memoryAbility: head.memory_ability,
|
||
question: head.question,
|
||
headline: head,
|
||
alt,
|
||
altSource: 'backup',
|
||
});
|
||
}
|
||
|
||
// (b) cal: independent re-answers; add ids not already covered by backup.
|
||
const cal = readJsonl(CAL);
|
||
const calMap = new Map<string, BeamRow>();
|
||
for (const r of cal) if (!calMap.has(r.instance_id)) calMap.set(r.instance_id, r);
|
||
for (const [id, alt] of calMap) {
|
||
if (pairs.has(id)) continue;
|
||
const head = headMap.get(id);
|
||
if (!head) continue;
|
||
if (!usable(head, alt)) continue;
|
||
pairs.set(id, {
|
||
instanceId: id,
|
||
memoryAbility: head.memory_ability,
|
||
question: head.question,
|
||
headline: head,
|
||
alt,
|
||
altSource: 'cal',
|
||
});
|
||
}
|
||
|
||
// Canonical sort (by instance_id) then deterministic seed-42 shuffle.
|
||
const sorted = [...pairs.values()].sort((a, b) => a.instanceId.localeCompare(b.instanceId));
|
||
return shuffle(sorted, 42);
|
||
}
|
||
|
||
function mergeUser(question: string, draftA: string, draftB: string): string {
|
||
return `QUESTION:\n${question}\n\nDRAFT A:\n${draftA}\n\nDRAFT B:\n${draftB}`;
|
||
}
|
||
|
||
async function main(): Promise<void> {
|
||
const budgetArg = process.argv.indexOf('--budget');
|
||
const BUDGET = budgetArg >= 0 ? Number(process.argv[budgetArg + 1]) : 8;
|
||
|
||
const client = createBeamOpenAiClient({ model: 'gpt-5', pricing: OPENAI_PRICING['gpt-5'] });
|
||
|
||
const allPairs = buildPairs();
|
||
|
||
// Resume: skip instance_ids already written.
|
||
const done = new Set<string>();
|
||
if (fs.existsSync(OUT)) {
|
||
for (const line of fs.readFileSync(OUT, 'utf-8').split('\n')) {
|
||
if (!line.trim()) continue;
|
||
try {
|
||
done.add((JSON.parse(line) as { instance_id: string }).instance_id);
|
||
} catch {
|
||
/* skip */
|
||
}
|
||
}
|
||
}
|
||
|
||
const abilityDist: Record<string, number> = {};
|
||
for (const p of allPairs) abilityDist[p.memoryAbility] = (abilityDist[p.memoryAbility] ?? 0) + 1;
|
||
console.log(`[merge-pilot] ${allPairs.length} pairs total | already done: ${done.size} | budget $${BUDGET}`);
|
||
console.log(`[merge-pilot] ability dist: ${JSON.stringify(abilityDist)}`);
|
||
|
||
const outStream = fs.createWriteStream(OUT, { flags: 'a' });
|
||
let spend = 0;
|
||
let processed = 0;
|
||
let mergeFails = 0;
|
||
|
||
for (const pair of allPairs) {
|
||
if (spend >= BUDGET) {
|
||
console.warn(`[merge-pilot] budget $${BUDGET} reached — stopping (spend=$${spend.toFixed(3)})`);
|
||
break;
|
||
}
|
||
if (done.has(pair.instanceId)) continue;
|
||
|
||
const rubric = pair.headline.nugget_scores.map(n => n.nugget);
|
||
if (rubric.length === 0) continue;
|
||
|
||
// ── MERGE (gpt-5) ──
|
||
const mergeRes = await client.chat({
|
||
system: MERGE_SYSTEM,
|
||
user: mergeUser(pair.question, pair.headline.answer, pair.alt.answer),
|
||
maxTokens: 2000,
|
||
});
|
||
spend += mergeRes.costUsd;
|
||
if (mergeRes.failureMode || isEmpty(mergeRes.text)) {
|
||
mergeFails++;
|
||
console.warn(` [${pair.instanceId}] merge failed (${mergeRes.failureMode ?? 'empty'}) — skipping`);
|
||
continue;
|
||
}
|
||
const merged = mergeRes.text.trim();
|
||
|
||
// ── JUDGE merged answer (gpt-5, plain nugget-mean) ──
|
||
const { judgement, llmResults } = await judgeQuestion(
|
||
client,
|
||
{ question: pair.question, rubric, memoryAbility: pair.memoryAbility, answer: merged },
|
||
{ computeTau: false },
|
||
);
|
||
const judgeCost = llmResults.reduce((s, r) => s + r.costUsd, 0);
|
||
spend += judgeCost;
|
||
|
||
const headlineNug = pair.headline.nugget_scores.map(n => n.score);
|
||
const mergedNug = judgement.nuggetScores.map(n => n.score);
|
||
|
||
const row = {
|
||
instance_id: pair.instanceId,
|
||
memory_ability: pair.memoryAbility,
|
||
alt_source: pair.altSource,
|
||
question: pair.question,
|
||
score_headline: pair.headline.score,
|
||
score_alt: pair.alt.score,
|
||
score_merged: judgement.score,
|
||
nuggets: rubric,
|
||
headline_nugget_scores: headlineNug,
|
||
merged_nugget_scores: mergedNug,
|
||
headline_idk: looksIDK(pair.headline.answer),
|
||
alt_idk: looksIDK(pair.alt.answer),
|
||
merged_idk: looksIDK(merged),
|
||
answer_merged: merged,
|
||
merge_cost_usd: round4(mergeRes.costUsd),
|
||
judge_cost_usd: round4(judgeCost),
|
||
};
|
||
outStream.write(`${JSON.stringify(row)}\n`);
|
||
processed++;
|
||
const delta = pair.headline.score === judgement.score ? 'TIE' : judgement.score > pair.headline.score ? 'WIN' : 'LOSS';
|
||
process.stdout.write(
|
||
` [${pair.memoryAbility.padEnd(24)}] H=${pair.headline.score.toFixed(2)} A=${pair.alt.score.toFixed(2)} M=${judgement.score.toFixed(2)} ${delta} $${spend.toFixed(3)}\n`,
|
||
);
|
||
}
|
||
|
||
outStream.end();
|
||
console.log(`\n[merge-pilot] DONE. processed=${processed} mergeFails=${mergeFails} spend=$${spend.toFixed(3)} → ${OUT}`);
|
||
}
|
||
|
||
function round4(x: number): number {
|
||
return Math.round(x * 1e4) / 1e4;
|
||
}
|
||
|
||
main().catch(err => {
|
||
console.error(err);
|
||
process.exit(1);
|
||
});
|