moving

benchmarks/harness/scripts/build-beam-canonical.ts

@@ -0,0 +1,611 @@
+#!/usr/bin/env tsx
+/**
+ * Canonical BEAM archive builder — Track A (manifest-v8.2-final.md).
+ *
+ * BEAM data ships inside the mohammadtavakoli78/BEAM GitHub repo under chats/.
+ * No separate download step is required if you have cloned the repo.
+ *
+ * Actual directory layout (discovered by inspection):
+ *   <beam-repo>/chats/<size>/              100K | 500K | 1M | 10M
+ *     <N>/                                numbered conversation directories (1-based)
+ *       chat.json                          [{batch_number, time_anchor, turns: [[{role,id,time_anchor,index,question_type,content}, ...], ...]}]
+ *       probing_questions/
+ *         probing_questions.json           {<category>: [{question, <answer_field>, difficulty, ...}, ...], ...}
+ *       topic.json                         {topic, description, ...}
+ *
+ * Chat-size alias: repo uses "100K" for what we call "128K" (~130K tokens each).
+ *
+ * Reads:   <beam-chats-path>/<chat-size-dir>/
+ * Writes:  benchmarks/data/beam/beam-<chat-size>.jsonl
+ *          benchmarks/data/beam/beam-<chat-size>.meta.json
+ *
+ * Canonicalisation guarantees (required for dataset_version hash determinism):
+ *   1. Include every (conversation × memory_ability × question) triple.
+ *   2. Sort by instance_id ascending.
+ *   3. JSON.stringify each row (no spaces, explicit key order) + '\n'. No BOM.
+ *   4. SHA-256 of the final byte stream.
+ *
+ * Per-instance JSONL schema:
+ *   {
+ *     "instance_id":        "beam_<chatSize>_<convId>_<ability>_q<idx>",
+ *     "conversation_id":    "beam_<convId>",
+ *     "question":           "<probing question>",
+ *     "expected":           ["<reference answer>"],
+ *     "context":            "<flat role: content lines>",
+ *     "memory_ability":     "<category>",
+ *     "chat_size":          "<128K|500K|1M|10M>",
+ *     "conversation_index": <int>
+ *   }
+ *
+ * Source: mohammadtavakoli78/BEAM (GitHub)
+ * Paper:  Tavakoli et al. 2024 "Beyond a Million Tokens: Benchmarking and
+ *         Enhancing Long-Term Memory in LLMs" (arXiv:2510.27246, ICLR 2026).
+ *
+ * Zero LLM calls. Zero npm packages beyond Node.js built-ins.
+ *
+ * Usage:
+ *   tsx build-beam-canonical.ts --beam-chats-path /path/to/BEAM/chats [--chat-size 128K]
+ *   # --beam-chats-path defaults to <repo-root>/benchmarks/harness/scripts/../../../BEAM/chats
+ *   #                                i.e. a sibling BEAM clone next to waggle-os
+ */
+
+import crypto from 'node:crypto';
+import fs from 'node:fs';
+import path from 'node:path';
+import process from 'node:process';
+import url from 'node:url';
+
+// ---------------------------------------------------------------------------
+// Constants
+// ---------------------------------------------------------------------------
+
+const VALID_CHAT_SIZES = ['128K', '500K', '1M', '10M'] as const;
+type ChatSize = (typeof VALID_CHAT_SIZES)[number];
+
+/**
+ * Map our canonical chat-size names to the directory names used in the BEAM repo.
+ * 128K ≈ 100K (actual token count ~130K).
+ */
+const CHAT_SIZE_DIR_MAP: Record<ChatSize, string[]> = {
+  '128K': ['100K'],
+  '500K': ['500K'],
+  '1M':   ['1M'],
+  '10M':  ['10M'],
+};
+
+/**
+ * All 10 BEAM memory ability categories.
+ * Source: Table 1 in arXiv:2510.27246 and repo README.
+ */
+const MEMORY_ABILITIES = [
+  'abstention',
+  'contradiction_resolution',
+  'event_ordering',
+  'information_extraction',
+  'instruction_following',
+  'knowledge_update',
+  'multi_session_reasoning',
+  'preference_following',
+  'summarization',
+  'temporal_reasoning',
+] as const;
+type MemoryAbility = (typeof MEMORY_ABILITIES)[number];
+
+// ---------------------------------------------------------------------------
+// BEAM data schema (verified by inspection of actual BEAM repo files)
+// ---------------------------------------------------------------------------
+
+/** One message inside a BEAM conversation turn. */
+interface BeamMessage {
+  role: string;
+  content: string;
+  id?: number;
+  time_anchor?: string | null;
+  index?: string;
+  question_type?: string;
+}
+
+/** One batch (session) in chat.json. */
+interface BeamBatch {
+  batch_number?: number;
+  time_anchor?: string | null;
+  turns: BeamMessage[][];  // list of turn groups; each group is list of msgs
+}
+
+/**
+ * Per-category probing question.
+ * Answer field varies by category — we try all known variants.
+ */
+interface BeamProbingQuestion {
+  question?: string;
+  // Category-specific answer fields (verified by inspection):
+  answer?: string;             // event_ordering, information_extraction, knowledge_update, multi_session_reasoning, temporal_reasoning
+  ideal_response?: string;     // abstention
+  ideal_answer?: string;       // contradiction_resolution
+  expected_compliance?: string; // instruction_following, preference_following
+  ideal_summary?: string;       // summarization
+  // Extra fields (stored for provenance, not used in eval directly):
+  difficulty?: string;
+  rubric?: string[];
+  [key: string]: unknown;
+}
+
+// ---------------------------------------------------------------------------
+// Output schema
+// ---------------------------------------------------------------------------
+
+interface CanonicalInstance {
+  instance_id: string;
+  conversation_id: string;
+  question: string;
+  expected: string[];
+  context: string;
+  memory_ability: string;
+  chat_size: string;
+  conversation_index: number;
+  /**
+   * schema v2 (2026-07-06): the ordered list of rubric "nuggets" for this
+   * probing question. This is the criterion set the OFFICIAL BEAM metric
+   * scores — each nugget is judged 0 / 0.5 / 1 by the graded LLM judge and
+   * the per-question score is their mean (see beam-nugget-judge.ts). Carried
+   * additively; v1 archives (e.g. the committed beam-128K.jsonl, dataset
+   * hash 9311bba4…) do not have it. Appended LAST in FIELD_ORDER so the
+   * leading columns are byte-identical to v1 for a human diff — note that
+   * ANY added field changes the SHA-256 dataset_version, so a v1 archive
+   * rebuilt with this code becomes a v2 hash. We do not rebuild 128K here.
+   */
+  rubric: string[];
+}
+
+const FIELD_ORDER: readonly (keyof CanonicalInstance)[] = [
+  'instance_id',
+  'conversation_id',
+  'question',
+  'expected',
+  'context',
+  'memory_ability',
+  'chat_size',
+  'conversation_index',
+  'rubric',
+];
+
+/** Canonical schema version. Bumped to 2 when `rubric` nuggets were added. */
+const SCHEMA_VERSION = 2;
+
+// ---------------------------------------------------------------------------
+// Turn flattening
+// ---------------------------------------------------------------------------
+
+/**
+ * Flatten BEAM's nested batch/turn-group structure into a flat list of
+ * {role, content} messages, preserving temporal order.
+ *
+ * chat.json structure:
+ *   [ {batch_number, time_anchor, turns: [ [msg, msg, ...], [msg, ...] ]} ]
+ *
+ * We flatten: batches → turn groups → individual messages.
+ * We only emit messages with non-empty content.
+ */
+function flattenBeamTurns(batches: BeamBatch[]): Array<{ role: string; content: string }> {
+  const out: Array<{ role: string; content: string }> = [];
+  for (const batch of batches) {
+    if (!Array.isArray(batch.turns)) continue;
+    for (const turnGroup of batch.turns) {
+      if (!Array.isArray(turnGroup)) continue;
+      for (const msg of turnGroup) {
+        const role = String(msg.role ?? 'unknown').toLowerCase();
+        const content = String(msg.content ?? '').trim();
+        if (content) {
+          out.push({ role, content });
+        }
+      }
+    }
+  }
+  return out;
+}
+
+/**
+ * Convert flat message list to context string.
+ * Format: "user: ...\nassistant: ...\n"
+ */
+function buildContext(msgs: Array<{ role: string; content: string }>): string {
+  return msgs.map(m => `${m.role}: ${m.content}`).join('\n');
+}
+
+// ---------------------------------------------------------------------------
+// Answer normalisation
+// ---------------------------------------------------------------------------
+
+/**
+ * Extract the reference answer from a probing question, trying all known
+ * per-category answer field names.
+ * Returns null if no answer field is found.
+ */
+function normaliseAnswer(pq: BeamProbingQuestion): string | null {
+  return (
+    pq.answer ??
+    pq.ideal_response ??
+    pq.ideal_answer ??
+    pq.expected_compliance ??
+    pq.ideal_summary ??
+    null
+  );
+}
+
+/**
+ * Extract the ordered list of rubric "nuggets" from a probing question.
+ * Ported verbatim from mem0's `extract_rubric_nuggets` (benchmarks/beam/run.py):
+ * the `rubric` field may be a list[str] (the BEAM 1M/10M shape), a dict with a
+ * `nuggets` list, or a bare scalar. Empty/whitespace nuggets are dropped.
+ */
+function extractRubricNuggets(pq: BeamProbingQuestion): string[] {
+  const raw = (pq as Record<string, unknown>).rubric;
+  const clean = (arr: unknown[]): string[] =>
+    arr
+      .map(n =>
+        n !== null && typeof n === 'object'
+          ? String((n as Record<string, unknown>).description ??
+                   (n as Record<string, unknown>).text ??
+                   JSON.stringify(n))
+          : String(n),
+      )
+      .map(s => s.trim())
+      .filter(s => s.length > 0);
+
+  if (Array.isArray(raw)) return clean(raw);
+  if (raw !== null && typeof raw === 'object') {
+    const nuggets = (raw as Record<string, unknown>).nuggets;
+    if (Array.isArray(nuggets)) return clean(nuggets);
+  }
+  if (raw !== undefined && raw !== null && String(raw).trim().length > 0) {
+    return [String(raw).trim()];
+  }
+  return [];
+}
+
+// ---------------------------------------------------------------------------
+// Serialisation
+// ---------------------------------------------------------------------------
+
+function serializeCanonical(inst: CanonicalInstance): string {
+  const ordered: Record<string, unknown> = {};
+  for (const key of FIELD_ORDER) {
+    ordered[key] = inst[key];
+  }
+  return JSON.stringify(ordered);
+}
+
+// ---------------------------------------------------------------------------
+// Directory discovery
+// ---------------------------------------------------------------------------
+
+/**
+ * Locate the chat-size directory under beamChatsPath.
+ * Maps our canonical size name to the BEAM repo directory name.
+ */
+function locateChatSizeDir(beamChatsPath: string, chatSize: ChatSize): string | null {
+  const candidates = CHAT_SIZE_DIR_MAP[chatSize] ?? [chatSize];
+  for (const dirName of candidates) {
+    const full = path.join(beamChatsPath, dirName);
+    if (fs.existsSync(full) && fs.statSync(full).isDirectory()) {
+      return full;
+    }
+  }
+  return null;
+}
+
+/**
+ * Return all numbered conversation directories inside chatSizeDir.
+ * These are directories whose names are numeric strings (1, 2, 3, ...).
+ */
+function discoverConversationDirs(chatSizeDir: string): string[] {
+  const entries = fs.readdirSync(chatSizeDir, { withFileTypes: true });
+  return entries
+    .filter(e => e.isDirectory() && /^\d+$/.test(e.name))
+    .sort((a, b) => Number(a.name) - Number(b.name))
+    .map(e => path.join(chatSizeDir, e.name));
+}
+
+// ---------------------------------------------------------------------------
+// CLI arg parsing
+// ---------------------------------------------------------------------------
+
+function parseArgs(): { beamChatsPath: string; chatSize: ChatSize } {
+  const argv = process.argv.slice(2);
+  let beamChatsPath = '';
+  let chatSize: ChatSize = '128K';
+
+  for (let i = 0; i < argv.length; i++) {
+    const flag = argv[i];
+    const next = argv[i + 1];
+    if ((flag === '--beam-chats-path' || flag === '--beam-data-path') && next) {
+      // Accept both --beam-chats-path (new) and --beam-data-path (old compat)
+      // If user passes BEAM root (contains chats/ subdir), auto-append chats/
+      let p = next;
+      if (fs.existsSync(path.join(p, 'chats'))) {
+        p = path.join(p, 'chats');
+      }
+      beamChatsPath = p;
+      i++;
+    } else if (flag === '--chat-size' && next) {
+      const val = next as ChatSize;
+      if (!VALID_CHAT_SIZES.includes(val)) {
+        console.error(
+          `[build-beam-canonical] unknown --chat-size "${val}". Valid: ${VALID_CHAT_SIZES.join(', ')}`,
+        );
+        process.exit(1);
+      }
+      chatSize = val;
+      i++;
+    }
+  }
+
+  return { beamChatsPath, chatSize };
+}
+
+// ---------------------------------------------------------------------------
+// Main
+// ---------------------------------------------------------------------------
+
+function main(): void {
+  const { beamChatsPath: beamChatsPathArg, chatSize } = parseArgs();
+
+  const here = url.fileURLToPath(import.meta.url);
+  // Script lives at benchmarks/harness/scripts/build-beam-canonical.ts
+  // repo root = 3 levels up: scripts/ → harness/ → benchmarks/ → repo root
+  const scriptDir = path.dirname(here);
+  const repoRoot = path.resolve(scriptDir, '..', '..', '..');
+  const dataDir = path.resolve(repoRoot, 'benchmarks', 'data');
+
+  // Auto-discover beamChatsPath if not provided:
+  // Try <repo-root>/../BEAM/chats (sibling clone convention)
+  let beamChatsPath = beamChatsPathArg;
+  if (!beamChatsPath) {
+    const siblingGuess = path.resolve(repoRoot, '..', 'BEAM', 'chats');
+    if (fs.existsSync(siblingGuess)) {
+      beamChatsPath = siblingGuess;
+      console.log(`[build-beam-canonical] auto-discovered BEAM chats at ${beamChatsPath}`);
+    } else {
+      console.error('[build-beam-canonical] --beam-chats-path is required (or clone BEAM as sibling of waggle-os).\n');
+      console.error('Clone with:  git clone https://github.com/mohammadtavakoli78/BEAM.git');
+      console.error('Then re-run: tsx build-beam-canonical.ts --beam-chats-path /path/to/BEAM/chats');
+      process.exit(2);
+    }
+  }
+
+  if (!fs.existsSync(beamChatsPath)) {
+    console.error(`[build-beam-canonical] BEAM chats path not found: ${beamChatsPath}`);
+    process.exit(2);
+  }
+
+  // ------------------------------------------------------------------
+  // Step 1: locate chat-size directory
+  // ------------------------------------------------------------------
+
+  const chatSizeDir = locateChatSizeDir(beamChatsPath, chatSize);
+  if (!chatSizeDir) {
+    const tried = (CHAT_SIZE_DIR_MAP[chatSize] ?? [chatSize]).map(d => path.join(beamChatsPath, d));
+    console.error(
+      `[build-beam-canonical] could not find chat-size directory for "${chatSize}" under ${beamChatsPath}.`,
+    );
+    console.error(`Tried: ${tried.join(', ')}`);
+    process.exit(2);
+  }
+
+  console.log(`[build-beam-canonical] chat-size directory: ${chatSizeDir}`);
+
+  // ------------------------------------------------------------------
+  // Step 2: scan conversation directories
+  // ------------------------------------------------------------------
+
+  const convDirs = discoverConversationDirs(chatSizeDir);
+  if (convDirs.length === 0) {
+    console.error(
+      `[build-beam-canonical] no numbered conversation directories found under ${chatSizeDir}.`,
+    );
+    process.exit(2);
+  }
+
+  console.log(`[build-beam-canonical] found ${convDirs.length} conversation directories`);
+
+  const all: CanonicalInstance[] = [];
+  const skipStats = {
+    missingChat: 0,
+    missingProbing: 0,
+    missingQuestion: 0,
+    missingAnswer: 0,
+    noTurns: 0,
+  };
+
+  for (let convIdx = 0; convIdx < convDirs.length; convIdx++) {
+    const convDir = convDirs[convIdx];
+    const convId = path.basename(convDir);  // e.g. "1", "2", ...
+
+    // ── Load chat.json ──────────────────────────────────────────────
+    const chatJsonPath = path.join(convDir, 'chat.json');
+    if (!fs.existsSync(chatJsonPath)) {
+      console.warn(`[build-beam-canonical] skipping ${convDir}: no chat.json`);
+      skipStats.missingChat++;
+      continue;
+    }
+
+    let batches: BeamBatch[];
+    try {
+      batches = JSON.parse(fs.readFileSync(chatJsonPath, 'utf-8')) as BeamBatch[];
+    } catch (err) {
+      console.warn(`[build-beam-canonical] skipping ${chatJsonPath}: ${String(err)}`);
+      skipStats.missingChat++;
+      continue;
+    }
+
+    const flatMsgs = flattenBeamTurns(batches);
+    if (flatMsgs.length === 0) {
+      console.warn(`[build-beam-canonical] skipping ${convDir}: 0 messages after flatten`);
+      skipStats.noTurns++;
+      continue;
+    }
+    const context = buildContext(flatMsgs);
+
+    // ── Load probing_questions/probing_questions.json ───────────────
+    const pqPath = path.join(convDir, 'probing_questions', 'probing_questions.json');
+    if (!fs.existsSync(pqPath)) {
+      console.warn(`[build-beam-canonical] skipping ${convDir}: no probing_questions.json`);
+      skipStats.missingProbing++;
+      continue;
+    }
+
+    let pqData: Record<string, BeamProbingQuestion[]>;
+    try {
+      pqData = JSON.parse(fs.readFileSync(pqPath, 'utf-8')) as Record<string, BeamProbingQuestion[]>;
+    } catch (err) {
+      console.warn(`[build-beam-canonical] skipping ${pqPath}: ${String(err)}`);
+      skipStats.missingProbing++;
+      continue;
+    }
+
+    const conversationId = `beam_${convId}`;
+    const safeChatSize = chatSize.replace(/[^a-zA-Z0-9]/g, '');
+
+    // ── Iterate categories ──────────────────────────────────────────
+    for (const [category, questions] of Object.entries(pqData)) {
+      if (!Array.isArray(questions)) continue;
+
+      for (let qi = 0; qi < questions.length; qi++) {
+        const pq = questions[qi];
+        const questionText = pq.question ?? null;
+        if (!questionText) {
+          skipStats.missingQuestion++;
+          continue;
+        }
+        const rubric = extractRubricNuggets(pq);
+        // `expected` keeps the single normalised reference answer for the
+        // legacy substring scorer. If a question has no single-answer field
+        // but does carry rubric nuggets, fall back to the joined rubric
+        // (mem0's ground_truth_answer convention) rather than dropping it.
+        let answerText = normaliseAnswer(pq);
+        if (answerText === null) {
+          if (rubric.length > 0) {
+            answerText = rubric.join(' | ');
+          } else {
+            skipStats.missingAnswer++;
+            continue;
+          }
+        }
+
+        const instanceId = `beam_${safeChatSize}_${convId}_${category}_q${qi}`;
+
+        all.push({
+          instance_id: instanceId,
+          conversation_id: conversationId,
+          question: questionText,
+          expected: [answerText],
+          context,
+          memory_ability: category,
+          chat_size: chatSize,
+          conversation_index: convIdx,
+          rubric,
+        });
+      }
+    }
+  }
+
+  // ------------------------------------------------------------------
+  // Step 3: validate
+  // ------------------------------------------------------------------
+
+  if (all.length === 0) {
+    console.error('[build-beam-canonical] extracted 0 instances.');
+    console.error(`Scanned ${convDirs.length} conversation dirs. Skip stats: ${JSON.stringify(skipStats)}`);
+    process.exit(2);
+  }
+
+  // ------------------------------------------------------------------
+  // Step 4: sort + distribution
+  // ------------------------------------------------------------------
+
+  all.sort((a, b) => a.instance_id.localeCompare(b.instance_id));
+
+  const byAbility: Record<string, number> = {};
+  for (const ma of MEMORY_ABILITIES) byAbility[ma] = 0;
+  for (const inst of all) {
+    byAbility[inst.memory_ability] = (byAbility[inst.memory_ability] ?? 0) + 1;
+  }
+
+  // ------------------------------------------------------------------
+  // Step 5: write outputs
+  // ------------------------------------------------------------------
+
+  const outDir = path.join(dataDir, 'beam');
+  if (!fs.existsSync(outDir)) fs.mkdirSync(outDir, { recursive: true });
+
+  const safeChatSize = chatSize.replace(/[^a-zA-Z0-9]/g, '');
+  const jsonlPath = path.join(outDir, `beam-${safeChatSize}.jsonl`);
+
+  // Stream the write + hash incrementally, one line at a time. The full body
+  // for the 1M/10M tracks (~3 GB for 1M, since the ~4 MB context is repeated
+  // per question) exceeds V8's max string length (~512 MB), so it can never
+  // be materialised as a single `join()`ed string. Writing `line + '\n'` for
+  // each row in sorted order produces byte-identical output to the old
+  // `all.map(serializeCanonical).join('\n') + '\n'` (a trailing newline after
+  // the final row), so the SHA-256 dataset_version stays deterministic and
+  // matches what the join-based path would have produced.
+  const hasher = crypto.createHash('sha256');
+  const fd = fs.openSync(jsonlPath, 'w');
+  try {
+    for (const inst of all) {
+      const line = serializeCanonical(inst) + '\n';
+      fs.writeSync(fd, line, null, 'utf-8');
+      hasher.update(line, 'utf-8');
+    }
+  } finally {
+    fs.closeSync(fd);
+  }
+  const hash = hasher.digest('hex');
+
+  const metaPath = path.join(outDir, `beam-${safeChatSize}.meta.json`);
+  const withRubric = all.filter(i => i.rubric.length > 0).length;
+  const totalNuggets = all.reduce((s, i) => s + i.rubric.length, 0);
+
+  const meta = {
+    dataset_version: hash,
+    schema_version: SCHEMA_VERSION,
+    instances_with_rubric: withRubric,
+    total_nuggets: totalNuggets,
+    instance_count: all.length,
+    chat_size: chatSize,
+    built_at: new Date().toISOString(),
+    source: 'mohammadtavakoli78/BEAM (GitHub)',
+    source_reference:
+      'Tavakoli, Salemi, Ye, Abdalla, Zamani, Mitchell 2024, "Beyond a Million Tokens: ' +
+      'Benchmarking and Enhancing Long-Term Memory in LLMs" (arXiv:2510.27246, ICLR 2026)',
+    beam_chats_path: beamChatsPath,
+    conversations_processed: convDirs.length,
+    chat_size_dir_alias: CHAT_SIZE_DIR_MAP[chatSize]?.[0] ?? chatSize,
+    canonicalisation: {
+      sort_order: 'instance_id ascending',
+      field_order: FIELD_ORDER,
+      line_terminator: '\\n',
+      trailing_newline: true,
+      encoding: 'utf-8',
+      no_bom: true,
+    },
+    distribution_by_memory_ability: byAbility,
+    skip_stats: skipStats,
+  };
+  fs.writeFileSync(metaPath, JSON.stringify(meta, null, 2) + '\n', 'utf-8');
+
+  // ------------------------------------------------------------------
+  // Step 6: report
+  // ------------------------------------------------------------------
+
+  console.log('[build-beam-canonical] distribution by memory_ability:');
+  for (const [k, v] of Object.entries(byAbility)) {
+    console.log(`  ${k}: ${v}`);
+  }
+  console.log('[build-beam-canonical] skip_stats:', skipStats);
+  console.log(`[build-beam-canonical] wrote ${jsonlPath} (${all.length} instances)`);
+  console.log(`[build-beam-canonical] wrote ${metaPath}`);
+  console.log(`[build-beam-canonical] dataset_version (SHA-256): ${hash}`);
+}
+
+main();