262 lines
10 KiB
TypeScript
262 lines
10 KiB
TypeScript
#!/usr/bin/env tsx
|
|
/**
|
|
* BEAM 1M — judge/metric plumbing SMOKE (no-context cell).
|
|
*
|
|
* PURPOSE: validate the full graded-nugget-judge + metric pipeline end-to-end
|
|
* on a cheap cell, per the build-phase brief. This is NOT a scored result — the
|
|
* "no-context" cell gives the answerer zero memories, so it should abstain on
|
|
* almost everything and score near-floor except on abstention questions (whose
|
|
* gold answer IS "I don't have enough information"). The point is to confirm the
|
|
* plumbing works and the metric distinguishes abilities, before any expensive
|
|
* substrate ingest.
|
|
*
|
|
* PIPELINE per question:
|
|
* 1. no-context answer : gpt-4o, buildAnswerGenerationPrompt(question, [])
|
|
* 2. graded judge : gpt-4o, each rubric nugget -> {0,0.5,1}, mean = score
|
|
* 3. metrics : Avg Score (micro) + Pass Rate (>=0.5), overall + per-ability
|
|
*
|
|
* SAMPLING: deterministic — the first N (default 5) questions per memory_ability
|
|
* encountered in the canonical's instance_id sort order (= 50 questions total).
|
|
*
|
|
* COST: hard-capped (default $4, under the $5 authorized). gpt-4o answerer+judge.
|
|
* OPENAI_API_KEY is read from waggle-os/.env (loadDotEnv).
|
|
*
|
|
* Usage:
|
|
* tsx benchmarks/harness/scripts/beam-smoke.ts \
|
|
* [--per-ability 5] [--model gpt-4o] [--budget 4] [--tau] \
|
|
* [--data benchmarks/data/beam/beam-1M.jsonl]
|
|
*/
|
|
|
|
import fs from 'node:fs';
|
|
import path from 'node:path';
|
|
import readline from 'node:readline';
|
|
import url from 'node:url';
|
|
import process from 'node:process';
|
|
|
|
import { createBeamOpenAiClient } from '../src/beam-openai-client.js';
|
|
import { buildAnswerGenerationPrompt, judgeQuestion } from '../src/beam-nugget-judge.js';
|
|
import type { BeamLlmResult } from '../src/beam-nugget-judge.js';
|
|
import { computeBeamMetrics, formatBeamMetrics } from '../src/beam-metrics.js';
|
|
import type { BeamQuestionResult } from '../src/beam-metrics.js';
|
|
|
|
const ALL_ABILITIES = [
|
|
'abstention', 'contradiction_resolution', 'event_ordering', 'information_extraction',
|
|
'instruction_following', 'knowledge_update', 'multi_session_reasoning',
|
|
'preference_following', 'summarization', 'temporal_reasoning',
|
|
];
|
|
|
|
interface CompactInstance {
|
|
instance_id: string;
|
|
question: string;
|
|
memory_ability: string;
|
|
rubric: string[];
|
|
expected: string[];
|
|
}
|
|
|
|
interface Args {
|
|
perAbility: number;
|
|
model: string;
|
|
budget: number;
|
|
computeTau: boolean;
|
|
dataPath: string;
|
|
}
|
|
|
|
function parseArgs(): Args {
|
|
const argv = process.argv.slice(2);
|
|
const here = url.fileURLToPath(import.meta.url);
|
|
const repoRoot = path.resolve(path.dirname(here), '..', '..', '..');
|
|
let perAbility = 5;
|
|
let model = 'gpt-4o';
|
|
let budget = 4;
|
|
let computeTau = false;
|
|
let dataPath = path.join(repoRoot, 'benchmarks', 'data', 'beam', 'beam-1M.jsonl');
|
|
for (let i = 0; i < argv.length; i++) {
|
|
const f = argv[i];
|
|
const next = argv[i + 1];
|
|
if (f === '--per-ability' && next) { perAbility = parseInt(next, 10); i++; }
|
|
else if (f === '--model' && next) { model = next; i++; }
|
|
else if (f === '--budget' && next) { budget = parseFloat(next); i++; }
|
|
else if (f === '--tau') { computeTau = true; }
|
|
else if (f === '--data' && next) { dataPath = path.resolve(next); i++; }
|
|
}
|
|
return { perAbility, model, budget, computeTau, dataPath };
|
|
}
|
|
|
|
/**
|
|
* Stream the (large, ~3 GB) canonical and collect the first `perAbility`
|
|
* instances per memory_ability. Early-terminates once every ability is full,
|
|
* so only a handful of conversations are ever parsed. Drops the giant `context`
|
|
* field immediately — the no-context cell does not use it.
|
|
*/
|
|
async function sampleInstances(dataPath: string, perAbility: number): Promise<CompactInstance[]> {
|
|
if (!fs.existsSync(dataPath)) {
|
|
throw new Error(`BEAM canonical not found at ${dataPath}. Build it via build-beam-canonical.ts --chat-size 1M`);
|
|
}
|
|
const buckets = new Map<string, CompactInstance[]>();
|
|
for (const a of ALL_ABILITIES) buckets.set(a, []);
|
|
const full = (): boolean => ALL_ABILITIES.every(a => (buckets.get(a)?.length ?? 0) >= perAbility);
|
|
|
|
const rl = readline.createInterface({ input: fs.createReadStream(dataPath, 'utf-8'), crlfDelay: Infinity });
|
|
try {
|
|
for await (const line of rl) {
|
|
const trimmed = line.trim();
|
|
if (!trimmed) continue;
|
|
let row: Record<string, unknown>;
|
|
try { row = JSON.parse(trimmed); } catch { continue; }
|
|
const ability = String(row.memory_ability ?? '');
|
|
const bucket = buckets.get(ability);
|
|
if (!bucket || bucket.length >= perAbility) {
|
|
if (full()) break;
|
|
continue;
|
|
}
|
|
bucket.push({
|
|
instance_id: String(row.instance_id ?? ''),
|
|
question: String(row.question ?? ''),
|
|
memory_ability: ability,
|
|
rubric: Array.isArray(row.rubric) ? (row.rubric as unknown[]).map(String) : [],
|
|
expected: Array.isArray(row.expected) ? (row.expected as unknown[]).map(String) : [],
|
|
});
|
|
if (full()) break;
|
|
}
|
|
} finally {
|
|
rl.close();
|
|
}
|
|
return ALL_ABILITIES.flatMap(a => buckets.get(a) ?? []);
|
|
}
|
|
|
|
function stripAnswerPrefix(text: string): string {
|
|
return text.includes('ANSWER:') ? text.split('ANSWER:').pop()!.trim() : text.trim();
|
|
}
|
|
|
|
async function main(): Promise<void> {
|
|
const args = parseArgs();
|
|
const startedAt = new Date();
|
|
console.log(`[beam-smoke] model=${args.model} per-ability=${args.perAbility} budget=$${args.budget} tau=${args.computeTau}`);
|
|
console.log(`[beam-smoke] data=${args.dataPath}`);
|
|
|
|
const client = createBeamOpenAiClient({ model: args.model });
|
|
|
|
console.log('[beam-smoke] sampling instances (streaming canonical)…');
|
|
const instances = await sampleInstances(args.dataPath, args.perAbility);
|
|
console.log(`[beam-smoke] sampled ${instances.length} instances across ${ALL_ABILITIES.length} abilities`);
|
|
|
|
let costUsd = 0;
|
|
let answerCalls = 0;
|
|
let judgeCalls = 0;
|
|
let inputTokens = 0;
|
|
let outputTokens = 0;
|
|
const acc = (r: BeamLlmResult): void => {
|
|
costUsd += r.costUsd; inputTokens += r.inputTokens; outputTokens += r.outputTokens;
|
|
};
|
|
|
|
const perQuestion: BeamQuestionResult[] = [];
|
|
const records: Record<string, unknown>[] = [];
|
|
let budgetStopped = false;
|
|
|
|
for (const inst of instances) {
|
|
if (costUsd >= args.budget) {
|
|
budgetStopped = true;
|
|
console.warn(`[beam-smoke] budget cap $${args.budget} reached — stopping at ${perQuestion.length} questions`);
|
|
break;
|
|
}
|
|
|
|
// 1. Generate no-context answer.
|
|
const ans = await client.chat({
|
|
system: '',
|
|
user: buildAnswerGenerationPrompt(inst.question, []),
|
|
maxTokens: 400,
|
|
});
|
|
acc(ans); answerCalls++;
|
|
const answer = stripAnswerPrefix(ans.text);
|
|
|
|
// 2. Judge nuggets.
|
|
const { judgement, llmResults } = await judgeQuestion(
|
|
client,
|
|
{ question: inst.question, rubric: inst.rubric, memoryAbility: inst.memory_ability, answer },
|
|
{ computeTau: args.computeTau },
|
|
);
|
|
for (const r of llmResults) { acc(r); judgeCalls++; }
|
|
|
|
perQuestion.push({
|
|
instanceId: inst.instance_id,
|
|
memoryAbility: inst.memory_ability,
|
|
score: judgement.score,
|
|
...(judgement.error ? { error: judgement.error } : {}),
|
|
});
|
|
records.push({
|
|
instance_id: inst.instance_id,
|
|
memory_ability: inst.memory_ability,
|
|
question: inst.question,
|
|
answer,
|
|
answer_failure_mode: ans.failureMode,
|
|
score: judgement.score,
|
|
judgment: judgement.judgment,
|
|
nugget_scores: judgement.nuggetScores,
|
|
...(judgement.scoreWithTau !== undefined ? { score_with_tau: judgement.scoreWithTau } : {}),
|
|
...(judgement.eventOrdering ? { event_ordering: judgement.eventOrdering } : {}),
|
|
n_nuggets: inst.rubric.length,
|
|
});
|
|
process.stdout.write(
|
|
` [${perQuestion.length}/${instances.length}] ${inst.memory_ability.padEnd(24)} ` +
|
|
`score=${judgement.score.toFixed(2)} (${judgement.judgment}) nuggets=${inst.rubric.length} $${costUsd.toFixed(3)}\n`,
|
|
);
|
|
}
|
|
|
|
const metrics = computeBeamMetrics(perQuestion);
|
|
const finishedAt = new Date();
|
|
|
|
// Write outputs.
|
|
const here = url.fileURLToPath(import.meta.url);
|
|
const repoRoot = path.resolve(path.dirname(here), '..', '..', '..');
|
|
const outDir = path.join(repoRoot, 'benchmarks', 'results', 'beam');
|
|
fs.mkdirSync(outDir, { recursive: true });
|
|
const ts = startedAt.toISOString().replace(/[:.]/g, '-');
|
|
const jsonlPath = path.join(outDir, `beam-1m-smoke-nocontext-${ts}.jsonl`);
|
|
const summaryPath = path.join(outDir, `beam-1m-smoke-nocontext-${ts}.summary.json`);
|
|
fs.writeFileSync(jsonlPath, records.map(r => JSON.stringify(r)).join('\n') + '\n', 'utf-8');
|
|
|
|
const summary = {
|
|
run: {
|
|
cell: 'no-context',
|
|
dataset: 'beam-1m',
|
|
model: args.model,
|
|
judge_model: args.model,
|
|
protocol: 'mem0-nugget-graded (0/0.5/1 avg-score, pass>=0.5)',
|
|
per_ability: args.perAbility,
|
|
compute_tau: args.computeTau,
|
|
startedAt: startedAt.toISOString(),
|
|
finishedAt: finishedAt.toISOString(),
|
|
budgetStopped,
|
|
},
|
|
metrics: {
|
|
overall_avg_score: metrics.overall.avgScore,
|
|
overall_pass_rate_pct: metrics.overall.accuracy,
|
|
total: metrics.overall.total,
|
|
correct: metrics.overall.correct,
|
|
errors: metrics.overall.errors,
|
|
by_ability: metrics.byAbility,
|
|
},
|
|
cost: {
|
|
total_usd: costUsd,
|
|
answer_calls: answerCalls,
|
|
judge_calls: judgeCalls,
|
|
input_tokens: inputTokens,
|
|
output_tokens: outputTokens,
|
|
},
|
|
};
|
|
fs.writeFileSync(summaryPath, JSON.stringify(summary, null, 2) + '\n', 'utf-8');
|
|
|
|
console.log('\n════════ BEAM 1M smoke (no-context) — metric plumbing ════════');
|
|
console.log(formatBeamMetrics(metrics));
|
|
console.log('──────────────────────────────────────────────────────────────');
|
|
console.log(`cost=$${costUsd.toFixed(4)} answer_calls=${answerCalls} judge_calls=${judgeCalls} ` +
|
|
`in=${inputTokens} out=${outputTokens}`);
|
|
console.log(`jsonl: ${jsonlPath}`);
|
|
console.log(`summary: ${summaryPath}`);
|
|
}
|
|
|
|
main().catch(err => {
|
|
console.error('[beam-smoke] FATAL:', err);
|
|
process.exit(1);
|
|
});
|