208 lines
10 KiB
TypeScript
208 lines
10 KiB
TypeScript
#!/usr/bin/env tsx
|
||
/**
|
||
* E2 analysis — assemble the 2×2 store×prompt table, decompose main effects
|
||
* and interaction, attach the confound classification, and write
|
||
* beam-e2-FINAL.json.
|
||
*
|
||
* Metrics per cell: mean nugget score (0..1) and pass rate (judgment==PASS).
|
||
* Uncertainty: conversation-cluster bootstrap. The 70 questions are 2 per
|
||
* conversation × 35 conversations; resampling INDEPENDENT questions would
|
||
* understate variance because the two questions from one conversation share a
|
||
* store. So we resample the 35 conversation clusters with replacement (B=10000)
|
||
* and recompute every statistic on each resample.
|
||
*
|
||
* Cells:
|
||
* A raw v2 (retain dated turns + conflict-aware prompt)
|
||
* B raw incumbent (retain dated turns + prefer-most-recent prompt)
|
||
* C reconciled v2 (collapsed current-state store + conflict-aware)
|
||
* D reconciled incumbent (collapsed current-state store + prefer-recent)
|
||
*
|
||
* STORE main effect = mean(raw {A,B}) − mean(reconciled {C,D})
|
||
* PROMPT main effect = mean(v2 {A,C}) − mean(incumbent {B,D})
|
||
* INTERACTION = (A−B) − (C−D) [ = does the prompt gap depend on store ]
|
||
*/
|
||
|
||
import fs from 'node:fs';
|
||
import path from 'node:path';
|
||
import url from 'node:url';
|
||
|
||
const here = url.fileURLToPath(import.meta.url);
|
||
const outDir = path.resolve(path.dirname(here), '..', '..', 'results', 'beam');
|
||
|
||
interface Row {
|
||
instance_id: string;
|
||
conv: number;
|
||
score: number;
|
||
judgment: string;
|
||
}
|
||
|
||
function load(file: string): Row[] {
|
||
const p = path.join(outDir, file);
|
||
if (!fs.existsSync(p)) throw new Error(`missing ${file}`);
|
||
return fs.readFileSync(p, 'utf-8').split('\n').filter(l => l.trim()).map(l => {
|
||
const r = JSON.parse(l);
|
||
return { instance_id: r.instance_id, conv: r.conv, score: r.score, judgment: r.judgment };
|
||
});
|
||
}
|
||
|
||
const CELL_FILES: Record<string, string> = {
|
||
A: 'beam-e2-cellA.jsonl', B: 'beam-e2-cellB.jsonl',
|
||
C: 'beam-e2-cellC.jsonl', D: 'beam-e2-cellD.jsonl',
|
||
};
|
||
|
||
function mean(xs: number[]): number { return xs.reduce((s, x) => s + x, 0) / xs.length; }
|
||
function passRate(rows: Row[]): number { return rows.filter(r => r.judgment === 'PASS').length / rows.length; }
|
||
function meanScore(rows: Row[]): number { return mean(rows.map(r => r.score)); }
|
||
|
||
// Align all cells to a common instance_id ordering so cluster resampling picks
|
||
// the SAME conversation across cells.
|
||
function main(): void {
|
||
const cells: Record<string, Row[]> = {};
|
||
for (const [c, f] of Object.entries(CELL_FILES)) cells[c] = load(f);
|
||
|
||
const ids = cells.A.map(r => r.instance_id);
|
||
const convOf: Record<string, number> = {};
|
||
for (const r of cells.A) convOf[r.instance_id] = r.conv;
|
||
|
||
// index each cell by instance_id for aligned lookup
|
||
const byId: Record<string, Record<string, Row>> = {};
|
||
for (const [c, rows] of Object.entries(cells)) {
|
||
byId[c] = {};
|
||
for (const r of rows) byId[c][r.instance_id] = r;
|
||
}
|
||
for (const c of Object.keys(cells)) {
|
||
for (const id of ids) if (!byId[c][id]) throw new Error(`cell ${c} missing ${id}`);
|
||
}
|
||
|
||
const convs = [...new Set(ids.map(id => convOf[id]))];
|
||
const idsByConv: Record<number, string[]> = {};
|
||
for (const id of ids) (idsByConv[convOf[id]] ??= []).push(id);
|
||
|
||
// point estimates
|
||
const point: Record<string, { mean: number; pass: number; n: number }> = {};
|
||
for (const c of Object.keys(cells)) {
|
||
point[c] = { mean: meanScore(cells[c]), pass: passRate(cells[c]), n: cells[c].length };
|
||
}
|
||
|
||
// statistic vector computed from a set of instance ids on a given metric
|
||
const cellStat = (c: string, sampleIds: string[], metric: 'mean' | 'pass'): number => {
|
||
const rows = sampleIds.map(id => byId[c][id]);
|
||
return metric === 'mean' ? meanScore(rows) : passRate(rows);
|
||
};
|
||
|
||
const B = 10000;
|
||
// seeded RNG (mulberry32) for reproducibility
|
||
let seed = 0x9e3779b9;
|
||
const rng = (): number => {
|
||
seed |= 0; seed = (seed + 0x6d2b79f5) | 0;
|
||
let t = Math.imul(seed ^ (seed >>> 15), 1 | seed);
|
||
t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t;
|
||
return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
|
||
};
|
||
|
||
// accumulate bootstrap distributions
|
||
const dist: Record<string, number[]> = {};
|
||
const push = (k: string, v: number) => (dist[k] ??= []).push(v);
|
||
|
||
for (let b = 0; b < B; b++) {
|
||
// resample conversation clusters with replacement
|
||
const sampleIds: string[] = [];
|
||
for (let i = 0; i < convs.length; i++) {
|
||
const conv = convs[Math.floor(rng() * convs.length)];
|
||
sampleIds.push(...idsByConv[conv]);
|
||
}
|
||
for (const metric of ['mean', 'pass'] as const) {
|
||
const A = cellStat('A', sampleIds, metric);
|
||
const Bc = cellStat('B', sampleIds, metric);
|
||
const C = cellStat('C', sampleIds, metric);
|
||
const D = cellStat('D', sampleIds, metric);
|
||
push(`A_${metric}`, A); push(`B_${metric}`, Bc); push(`C_${metric}`, C); push(`D_${metric}`, D);
|
||
push(`store_${metric}`, (A + Bc) / 2 - (C + D) / 2); // raw − reconciled
|
||
push(`prompt_${metric}`, (A + C) / 2 - (Bc + D) / 2); // v2 − incumbent
|
||
push(`interaction_${metric}`, (A - Bc) - (C - D)); // prompt gap: raw − reconciled
|
||
push(`prompt_within_raw_${metric}`, A - Bc);
|
||
push(`prompt_within_recon_${metric}`, C - D);
|
||
push(`store_within_v2_${metric}`, A - C);
|
||
push(`store_within_incumbent_${metric}`, Bc - D);
|
||
}
|
||
}
|
||
|
||
const ci = (k: string): { lo: number; hi: number; se: number } => {
|
||
const xs = [...dist[k]].sort((a, b) => a - b);
|
||
const lo = xs[Math.floor(0.025 * xs.length)];
|
||
const hi = xs[Math.floor(0.975 * xs.length)];
|
||
const m = mean(xs);
|
||
const se = Math.sqrt(mean(xs.map(x => (x - m) ** 2)));
|
||
return { lo, hi, se };
|
||
};
|
||
|
||
// confound (optional)
|
||
let confound: unknown = null;
|
||
const confP = path.join(outDir, 'beam-e2-confound.json');
|
||
if (fs.existsSync(confP)) confound = JSON.parse(fs.readFileSync(confP, 'utf-8'));
|
||
|
||
const round = (x: number) => Math.round(x * 10000) / 10000;
|
||
const fmtCi = (k: string) => { const c = ci(k); return { lo: round(c.lo), hi: round(c.hi), se: round(c.se) }; };
|
||
|
||
const final = {
|
||
experiment: 'E2 — BEAM 2×2 store×prompt causal ablation (contradiction_resolution)',
|
||
n_questions: cells.A.length,
|
||
n_conversations: convs.length,
|
||
bootstrap: { method: 'conversation-cluster', B, seed_rng: 'mulberry32' },
|
||
external_anchor: { published_cellA_subset_mean: 0.5875, published_cellA_subset_pass: 0.8714, mem0_mean: 0.3571 },
|
||
design: {
|
||
A: { store: 'raw', prompt: 'v2' },
|
||
B: { store: 'raw', prompt: 'incumbent' },
|
||
C: { store: 'reconciled', prompt: 'v2' },
|
||
D: { store: 'reconciled', prompt: 'incumbent' },
|
||
},
|
||
cells: Object.fromEntries(Object.keys(cells).map(c => [c, {
|
||
mean: round(point[c].mean), pass: round(point[c].pass), n: point[c].n,
|
||
mean_ci: fmtCi(`${c}_mean`), pass_ci: fmtCi(`${c}_pass`),
|
||
}])),
|
||
effects: {
|
||
mean: {
|
||
store_raw_minus_reconciled: { point: round((point.A.mean + point.B.mean) / 2 - (point.C.mean + point.D.mean) / 2), ci: fmtCi('store_mean') },
|
||
prompt_v2_minus_incumbent: { point: round((point.A.mean + point.C.mean) / 2 - (point.B.mean + point.D.mean) / 2), ci: fmtCi('prompt_mean') },
|
||
interaction: { point: round((point.A.mean - point.B.mean) - (point.C.mean - point.D.mean)), ci: fmtCi('interaction_mean') },
|
||
prompt_within_raw: { point: round(point.A.mean - point.B.mean), ci: fmtCi('prompt_within_raw_mean') },
|
||
prompt_within_reconciled: { point: round(point.C.mean - point.D.mean), ci: fmtCi('prompt_within_recon_mean') },
|
||
store_within_v2: { point: round(point.A.mean - point.C.mean), ci: fmtCi('store_within_v2_mean') },
|
||
store_within_incumbent: { point: round(point.B.mean - point.D.mean), ci: fmtCi('store_within_incumbent_mean') },
|
||
},
|
||
pass: {
|
||
store_raw_minus_reconciled: { point: round((point.A.pass + point.B.pass) / 2 - (point.C.pass + point.D.pass) / 2), ci: fmtCi('store_pass') },
|
||
prompt_v2_minus_incumbent: { point: round((point.A.pass + point.C.pass) / 2 - (point.B.pass + point.D.pass) / 2), ci: fmtCi('prompt_pass') },
|
||
interaction: { point: round((point.A.pass - point.B.pass) - (point.C.pass - point.D.pass)), ci: fmtCi('interaction_pass') },
|
||
prompt_within_raw: { point: round(point.A.pass - point.B.pass), ci: fmtCi('prompt_within_raw_pass') },
|
||
prompt_within_reconciled: { point: round(point.C.pass - point.D.pass), ci: fmtCi('prompt_within_recon_pass') },
|
||
store_within_v2: { point: round(point.A.pass - point.C.pass), ci: fmtCi('store_within_v2_pass') },
|
||
store_within_incumbent: { point: round(point.B.pass - point.D.pass), ci: fmtCi('store_within_incumbent_pass') },
|
||
},
|
||
},
|
||
confound,
|
||
};
|
||
|
||
const outP = path.join(outDir, 'beam-e2-FINAL.json');
|
||
fs.writeFileSync(outP, JSON.stringify(final, null, 2) + '\n', 'utf-8');
|
||
|
||
// console summary
|
||
console.log('\n=== E2 2×2 (contradiction_resolution, n=' + cells.A.length + ') ===');
|
||
console.log('cell mean pass');
|
||
for (const c of ['A', 'B', 'C', 'D']) {
|
||
console.log(`${c} ${final.design[c as 'A'].store.padEnd(11)}${final.design[c as 'A'].prompt.padEnd(10)} ${point[c].mean.toFixed(4)} ${(point[c].pass * 100).toFixed(1)}%`);
|
||
}
|
||
const e = final.effects;
|
||
console.log('\nMAIN EFFECTS (mean nugget):');
|
||
console.log(` STORE (raw−recon): ${e.mean.store_raw_minus_reconciled.point.toFixed(4)} 95%CI[${e.mean.store_raw_minus_reconciled.ci.lo},${e.mean.store_raw_minus_reconciled.ci.hi}]`);
|
||
console.log(` PROMPT (v2−incumb): ${e.mean.prompt_v2_minus_incumbent.point.toFixed(4)} 95%CI[${e.mean.prompt_v2_minus_incumbent.ci.lo},${e.mean.prompt_v2_minus_incumbent.ci.hi}]`);
|
||
console.log(` INTERACTION: ${e.mean.interaction.point.toFixed(4)} 95%CI[${e.mean.interaction.ci.lo},${e.mean.interaction.ci.hi}]`);
|
||
console.log('MAIN EFFECTS (pass rate):');
|
||
console.log(` STORE (raw−recon): ${(e.pass.store_raw_minus_reconciled.point * 100).toFixed(1)}pp 95%CI[${(e.pass.store_raw_minus_reconciled.ci.lo * 100).toFixed(1)},${(e.pass.store_raw_minus_reconciled.ci.hi * 100).toFixed(1)}]`);
|
||
console.log(` PROMPT (v2−incumb): ${(e.pass.prompt_v2_minus_incumbent.point * 100).toFixed(1)}pp 95%CI[${(e.pass.prompt_v2_minus_incumbent.ci.lo * 100).toFixed(1)},${(e.pass.prompt_v2_minus_incumbent.ci.hi * 100).toFixed(1)}]`);
|
||
console.log(` INTERACTION: ${(e.pass.interaction.point * 100).toFixed(1)}pp 95%CI[${(e.pass.interaction.ci.lo * 100).toFixed(1)},${(e.pass.interaction.ci.hi * 100).toFixed(1)}]`);
|
||
console.log(`\n→ ${outP}`);
|
||
}
|
||
|
||
main();
|