mlx-bun 0.0.5 → 0.0.6

package/src/eval/tasks/mmlu.ts

@@ -1,17 +1,15 @@
-// MMLU 5-shot — multi-domain multiple-choice knowledge. Port of
-// optiq/eval/mmlu.py: per-subject dev exemplars, A/B/C/D format, greedy
-// decode, extract the predicted answer letter, compare to the gold index.
+// MMLU 5-shot — multi-domain multiple-choice knowledge. Faithful port of
+// optiq/eval/mmlu.py: per-subject dev exemplars, A/B/C/D format, and scoring by
+// **last-token logit argmax over the " A"/" B"/" C"/" D" tokens** on the RAW
+// 5-shot prompt — NO chat template, NO generation. This is the standard
+// (Hendrycks/Unsloth) MMLU methodology; the earlier chat-wrap+generate variant
+// scored ~5 pts low vs optiq's published number, so we match optiq exactly.
 //
-// The Python reference scores via first-letter logit argmax (no generation).
-// The mlx-bun harness routes everything through generateText() (decoded
-// completion). For an instruct model the faithful argmax-on-raw-text path
-// underscores on a small noisy sample, so we drive the chat template: the
-// 5-shot exemplars become the priming context and an explicit instruction
-// asks for ONLY the single letter A/B/C/D. Extraction is robust to chatty
-// formatting (</think> wrappers, "Answer: X", "(X)", "**X**", bare letter).
-// Prompt format, header, exemplar selection, and gold mapping are verbatim.
+// Questions: by DEFAULT we score optiq's EXACT frozen draw (mmlu_optiq_frozen, the
+// same ~969 his published 52.4 used) so base/adapter are directly comparable to his
+// table. MLX_BUN_MMLU_FROZEN=0 falls back to our own sampling of the full 14k set.
 
-import { generateText, loadJsonl, sampleIndices, type TaskModel } from "../runner";
+import { loadJsonl, sampleIndices, type TaskModel } from "../runner";
 
 interface MmluRow {
   question: string;
@@ -42,12 +40,8 @@ function formatExample(question: string, choices: string[], answerIdx?: number):
   return s;
 }
 
-/** 5-shot prompt: subject header → n demos → open test question.
- *
- * Built for the chat template: the header + worked exemplars prime the model,
- * and a trailing instruction pins the output to a single bare letter. The
- * exemplars still close their "Answer: X" line so the model sees the expected
- * shape, and the final question is left open ("Answer:") for the model. */
+/** 5-shot prompt (verbatim optiq _build_prompt): subject header → n demos →
+ *  open test question ending in "Answer:". Fed raw (no chat template). */
 function buildPrompt(row: MmluRow, devExamples: MmluRow[], nShots: number): string {
   let prompt = `The following are multiple choice questions (with answers) about ${formatSubject(row.subject)}.\n\n`;
   for (const ex of devExamples.slice(0, nShots)) {
@@ -55,48 +49,46 @@ function buildPrompt(row: MmluRow, devExamples: MmluRow[], nShots: number): stri
     prompt += "\n";
   }
   prompt += formatExample(row.question, row.choices);
-  prompt += "\n\nRespond with ONLY the single letter (A, B, C, or D) of the correct answer. Do not explain.";
   return prompt;
 }
 
-/** Model output → predicted letter index 0..3, else null. Robust to chatty
- * instruct formatting: strips a leading <think>…</think> block, then prefers
- * an explicit "Answer: X" / "(X)" / "**X**" marker, falling back to the first
- * standalone A/B/C/D token (one not glued inside a longer word). */
-export function extractAnswer(textIn: string): number | null {
-  let text = textIn;
-  if (text.includes("</think>")) text = text.split("</think>").pop() ?? text;
-
-  const toIdx = (ch: string): number => LETTERS.indexOf(ch.toUpperCase() as (typeof LETTERS)[number]);
-
-  // Explicit answer markers, in priority order.
-  const markers = [
-    /(?:answer|correct answer|the answer is)\s*(?:is|:)?\s*\(?\*{0,2}([A-Da-d])\*{0,2}\)?/i,
-    /\*\*\s*([A-Da-d])\s*\*\*/, // **X**
-    /\(\s*([A-Da-d])\s*\)/, // (X)
-  ];
-  for (const re of markers) {
-    const m = text.match(re);
-    if (m) return toIdx(m[1]!);
+/** Forward the prompt (prefill + last token) and return the index 0..3 of the
+ *  highest-logit answer letter — optiq's argmax-on-last-token scoring. */
+function lastTokenArgmax(model: TaskModel["model"], ids: number[], letterIds: number[]): number {
+  const cache = model.makeCache();
+  try {
+    if (ids.length > 1) model.forward(ids.slice(0, -1), cache).dispose(); // prefill (cache only)
+    const logits = model.forward([ids[ids.length - 1]!], cache); // [1, 1, V]
+    const v = logits.toFloat32();
+    logits.dispose();
+    let best = 0, bestVal = -Infinity;
+    for (let i = 0; i < letterIds.length; i++) {
+      const val = v[letterIds[i]!]!;
+      if (val > bestVal) { bestVal = val; best = i; }
+    }
+    return best;
+  } finally {
+    for (const c of cache) c.dispose();
   }
-
-  // Fallback: first standalone letter (word-boundary, not embedded in a word).
-  const standalone = text.match(/\b([A-D])\b/);
-  if (standalone) return toIdx(standalone[1]!);
-
-  return null;
 }
 
 export async function evaluateMmlu(
   tm: TaskModel,
-  opts: { nSamples?: number; nShots?: number; maxTokens?: number; seed?: number } = {},
+  opts: { nSamples?: number; nShots?: number; seed?: number; frozen?: boolean } = {},
 ): Promise<MmluResult> {
   const nShots = opts.nShots ?? 5;
-  const maxTokens = opts.maxTokens ?? 16;
   const seed = opts.seed ?? 42;
 
-  const testRows = loadJsonl<MmluRow>("mmlu_test");
-  const devRows = loadJsonl<MmluRow>("mmlu_dev");
+  // Optiq-parity mode (DEFAULT): score optiq's EXACT frozen draw — the same ~969
+  // questions + dev exemplars his published table used — so base/adapter numbers are
+  // DIRECTLY comparable to optiq's. Any remaining gap is then pure forward-logit
+  // parity, not sampling. Set MLX_BUN_MMLU_FROZEN=0 (or frozen:false) to fall back to
+  // our own stratified sampling of the full 14k test set (a DIFFERENT subset: JS RNG
+  // ≠ numpy RandomState, so it does not match optiq's questions even at the same seed).
+  const useFrozen = opts.frozen ?? (process.env.MLX_BUN_MMLU_FROZEN !== "0");
+
+  const testRows = loadJsonl<MmluRow>(useFrozen ? "mmlu_optiq_frozen" : "mmlu_test");
+  const devRows = loadJsonl<MmluRow>(useFrozen ? "mmlu_optiq_dev" : "mmlu_dev");
 
   // Per-subject dev exemplars for prompt construction.
   const devBySubject = new Map<string, MmluRow[]>();
@@ -106,38 +98,57 @@ export async function evaluateMmlu(
     devBySubject.set(ex.subject, list);
   }
 
-  // Stratified sample: roughly equal questions per subject (matches mmlu.py),
-  // deterministic via the shared seeded subsampler.
-  const testBySubject = new Map<string, number[]>();
-  for (let i = 0; i < testRows.length; i++) {
-    const subj = testRows[i]!.subject;
-    const list = testBySubject.get(subj) ?? [];
-    list.push(i);
-    testBySubject.set(subj, list);
-  }
-  const subjects = [...testBySubject.keys()].sort();
-  const nSamples = opts.nSamples ?? 1000;
-  const perSubject = Math.max(1, Math.floor(nSamples / subjects.length));
-
-  let sampled: number[] = [];
-  for (const subj of subjects) {
-    const pool = testBySubject.get(subj)!;
-    const take = Math.min(perSubject, pool.length);
-    for (const j of sampleIndices(pool.length, take, seed)) sampled.push(pool[j]!);
-  }
-  if (sampled.length > nSamples) {
-    const keep = sampleIndices(sampled.length, nSamples, seed);
-    sampled = keep.map((k) => sampled[k]!);
+  let sampled: number[];
+  if (useFrozen) {
+    // Already optiq's exact draw — score ALL of it, in his order, NO resampling
+    // (re-sampling here would re-introduce exactly the divergence we're removing).
+    sampled = testRows.map((_, i) => i);
+  } else {
+    // Stratified sample: roughly equal questions per subject (mirrors mmlu.py).
+    const testBySubject = new Map<string, number[]>();
+    for (let i = 0; i < testRows.length; i++) {
+      const subj = testRows[i]!.subject;
+      const list = testBySubject.get(subj) ?? [];
+      list.push(i);
+      testBySubject.set(subj, list);
+    }
+    const subjects = [...testBySubject.keys()].sort();
+    const nSamples = opts.nSamples ?? 1000;
+    const perSubject = Math.max(1, Math.floor(nSamples / subjects.length));
+    sampled = [];
+    for (const subj of subjects) {
+      const pool = testBySubject.get(subj)!;
+      const take = Math.min(perSubject, pool.length);
+      for (const j of sampleIndices(pool.length, take, seed)) sampled.push(pool[j]!);
+    }
+    if (sampled.length > nSamples) {
+      const keep = sampleIndices(sampled.length, nSamples, seed);
+      sampled = keep.map((k) => sampled[k]!);
+    }
   }
 
+  // Answer-letter token ids — " A".." D" first token, matching how Hendrycks
+  // scores the token after "Answer:". Strip a leading BOS (≈ add_special_tokens=False).
+  const bos = tm.tokenizer.bosTokenId;
+  const stripBos = (ids: number[]) => (bos !== null && ids[0] === bos ? ids.slice(1) : ids);
+  const letterIds = LETTERS.map((L) => {
+    const ids = stripBos(tm.tokenizer.encode(` ${L}`));
+    return ids[0]!;
+  });
+
+  // Re-assert the mounted adapter for our direct forward()s: generation tasks reset
+  // loraState.active to [] when they finish, and MMLU bypasses generate() entirely,
+  // so without this the argmax would silently score BASE even with --adapter set.
+  if (tm.activeAdapters) tm.model.loraState.active = tm.activeAdapters;
+
   let nCorrect = 0;
   for (let k = 0; k < sampled.length; k++) {
     const item = testRows[sampled[k]!]!;
     const devEx = devBySubject.get(item.subject) ?? [];
-    const out = await generateText(tm, buildPrompt(item, devEx, nShots), { maxTokens, useChat: true });
-    const pred = extractAnswer(out);
-    if (pred !== null && pred === item.answer) nCorrect++;
-    if ((k + 1) % 10 === 0 || k + 1 === sampled.length)
+    const ids = stripBos(tm.tokenizer.encode(buildPrompt(item, devEx, nShots)));
+    const pred = lastTokenArgmax(tm.model, ids, letterIds);
+    if (pred === item.answer) nCorrect++;
+    if ((k + 1) % 25 === 0 || k + 1 === sampled.length)
       process.stderr.write(`\r  mmlu ${k + 1}/${sampled.length}  acc=${((nCorrect / (k + 1)) * 100).toFixed(1)}%`);
   }
   process.stderr.write("\n");

package/src/jobs/runner.ts

@@ -149,6 +149,7 @@ interface QueuedSpawn {
   jobId: string;
   entry: string;
   bin: string;
+  onComplete?: (jobId: string, code: number) => void;
 }
 
 let gpuLeaseHolder: string | null = null;
@@ -170,6 +171,9 @@ export interface SubprocessOpts {
   entry?: string;
   /** Override the runtime binary (tests). Defaults to "bun". */
   bin?: string;
+  /** Called on the server (parent) after the child exits — used to invalidate
+   *  caches (e.g. the Library) so a finished quantize surfaces immediately. */
+  onComplete?: (jobId: string, code: number) => void;
 }
 
 /** Create a 'queued' row and either spawn it now (lease free) or leave it
@@ -188,6 +192,7 @@ export function submitSubprocess(
     jobId: row.id,
     entry: opts.entry ?? JOB_ENTRY_PATH,
     bin: opts.bin ?? "bun",
+    onComplete: opts.onComplete,
   };
   if (gpuLeaseHolder === null) {
     spawnNow(item);
@@ -243,6 +248,7 @@ function spawnNow(item: QueuedSpawn): void {
       }
     }
     releaseLease(jobId);
+    try { item.onComplete?.(jobId, code); } catch {}
     drainQueue();
   })();
 }

package/src/memory/chunk.ts

@@ -0,0 +1,25 @@
+// mlx-bun memory — synthesis stage 2: CHUNK.  ⚠️ M1 STUB.
+//
+// Segment each new conversation into topic chunks (a local-model call — this is
+// the segmentation track distilled from gold; eval by boundary/label accuracy).
+// Re-chunk only conversations whose updated_at > chunked_at. Mirrors lucien's
+// chunk-recent.ts. The prompt is the future "chunk" LoRA seam (prompts.ts).
+//
+// See docs/design/memory-system.md → "The nightly pipeline".
+
+import type { SynthesisEvent } from "./pipeline";
+
+export interface ChunkResult {
+  conversations: number;
+  chunks: number;
+}
+
+/** Conversations → topic chunks (local model). STUB. */
+export async function chunkConversations(
+  _model: unknown,
+  _onEvent?: (e: SynthesisEvent) => void,
+): Promise<ChunkResult> {
+  // TODO(M1): for each conv with updated_at > chunked_at, run the chunk prompt
+  // on the local model, write chunks, set chunked_at.
+  throw new Error("memory.chunk is not implemented yet (M1)");
+}

package/src/memory/cluster.ts

@@ -0,0 +1,24 @@
+// mlx-bun memory — synthesis stage 3: CLUSTER.  ⚠️ M1 STUB.
+//
+// Assign each new chunk to an existing bucket or propose a new one (local-model
+// call). Buckets are emergent (Meta/Buckets.md), each mapping to a target
+// article. Mirrors lucien's cluster-assign-recent.ts.
+//
+// See docs/design/memory-system.md → "The nightly pipeline".
+
+import type { SynthesisEvent } from "./pipeline";
+
+export interface ClusterResult {
+  assigned: number;
+  newBuckets: number;
+}
+
+/** Chunks → buckets (local model). STUB. */
+export async function clusterChunks(
+  _model: unknown,
+  _onEvent?: (e: SynthesisEvent) => void,
+): Promise<ClusterResult> {
+  // TODO(M1): for each unbucketed chunk, classify into Meta/Buckets.md taxonomy
+  // (or propose a new bucket → new article stem); persist chunk.bucket.
+  throw new Error("memory.cluster is not implemented yet (M1)");
+}

package/src/memory/db.ts

@@ -0,0 +1,114 @@
+// mlx-bun memory — synthesis bookkeeping store (bun:sqlite).
+//
+// ⚠️ M1 STUB schema-wise the tables are real, but no stage writes them yet.
+// The read path (vault.ts / tools.ts) does NOT touch this — it's only used by
+// the synthesis pipeline (pipeline.ts), which is itself stubbed.
+//
+// Reuses mlx-bun's existing per-domain sqlite pattern (cf. JobStore in
+// src/jobs/db.ts, EvalDB in src/evaldb.ts): a small class over bun:sqlite with
+// WAL + schema-in-constructor + an additive migrate() hook, living under
+// ~/.cache/mlx-bun/ alongside jobs.sqlite / evals.sqlite / registry.sqlite. It's
+// derived state — rebuildable from the pi transcripts + the vault articles — so
+// it belongs in the rebuildable cache, NOT in the (precious, git-tracked) vault.
+//
+// The schema mirrors lucien's incremental design: per-source watermarks (don't
+// re-ingest), chunked_at > updated_at (re-chunk only changed convs), and a
+// synthesized-chunk ledger (never fold the same chunk into an article twice).
+
+import { Database } from "bun:sqlite";
+import { mkdirSync } from "node:fs";
+
+export const DEFAULT_MEMORY_DB = `${process.env.HOME}/.cache/mlx-bun/memory.sqlite`;
+
+export interface ConversationRow {
+  /** Stable conv id: first 8 hex of the session UUID (the conv:HASH citation). */
+  conv: string;
+  source: string; // "pi-web" | "pi-terminal"
+  title: string;
+  /** Full normalized transcript JSON ({role, content}[]). */
+  transcript: string;
+  updated_at: number;
+  chunked_at: number | null;
+}
+
+export interface ChunkRow {
+  id: string;
+  conv: string;
+  ordinal: number;
+  text: string;
+  bucket: string | null;
+}
+
+export interface BucketRow {
+  name: string;
+  article: string; // target article stem
+  description: string;
+}
+
+const SCHEMA = `
+CREATE TABLE IF NOT EXISTS conversations (
+  conv       TEXT PRIMARY KEY,
+  source     TEXT NOT NULL,
+  title      TEXT NOT NULL DEFAULT '',
+  transcript TEXT NOT NULL,
+  updated_at INTEGER NOT NULL,
+  chunked_at INTEGER
+);
+CREATE TABLE IF NOT EXISTS chunks (
+  id      TEXT PRIMARY KEY,
+  conv    TEXT NOT NULL REFERENCES conversations(conv),
+  ordinal INTEGER NOT NULL,
+  text    TEXT NOT NULL,
+  bucket  TEXT
+);
+CREATE INDEX IF NOT EXISTS idx_chunks_conv ON chunks(conv);
+CREATE TABLE IF NOT EXISTS buckets (
+  name        TEXT PRIMARY KEY,
+  article     TEXT NOT NULL,
+  description TEXT NOT NULL DEFAULT ''
+);
+-- Per-source ingest watermark: only process what's new since last run.
+CREATE TABLE IF NOT EXISTS watermarks (
+  source     TEXT PRIMARY KEY,
+  cursor     TEXT NOT NULL,
+  updated_at INTEGER NOT NULL
+);
+-- Idempotency ledger: a (bucket, chunk) pair is folded into an article once.
+CREATE TABLE IF NOT EXISTS synthesized_bucket_chunks (
+  bucket   TEXT NOT NULL,
+  chunk_id TEXT NOT NULL,
+  PRIMARY KEY (bucket, chunk_id)
+);
+`;
+
+/**
+ * Synthesis bookkeeping DB. Same shape as JobStore/EvalDB: open in the
+ * constructor (WAL on, schema ensured, additive migrate), expose `db`, and a
+ * `close()`. Only the synthesis pipeline constructs one — opening it is what
+ * *creates* memory.sqlite, so the read path never does. Stage queries land in
+ * M1; the tables they expect already exist here.
+ */
+export class MemoryStore {
+  readonly db: Database;
+  readonly dbPath: string;
+
+  constructor(dbPath: string = DEFAULT_MEMORY_DB) {
+    if (dbPath !== ":memory:") {
+      try { mkdirSync(dbPath.slice(0, dbPath.lastIndexOf("/")), { recursive: true }); } catch {}
+    }
+    this.dbPath = dbPath;
+    this.db = new Database(dbPath);
+    this.db.exec("PRAGMA journal_mode = WAL");
+    this.db.exec(SCHEMA);
+    this.migrate();
+  }
+
+  /** Additive-only migrations (mirrors JobStore.migrate). Re-opening is a no-op. */
+  private migrate(): void {
+    // No columns to add today; the hook stays so future fields are additive.
+  }
+
+  close(): void {
+    this.db.close();
+  }
+}

package/src/memory/ingest.ts

@@ -0,0 +1,22 @@
+// mlx-bun memory — synthesis stage 1: INGEST.  ⚠️ M1 STUB.
+//
+// pi sessions (~/.mlx-bun/sessions + ~/.mlx-bun/pi, pi's own JSONL) →
+// memory.sqlite, normalized to {conv, source, title, transcript}, watermark-
+// gated so each run only processes what's new. No web scraping, no Playwright:
+// mlx-bun synthesizes ITS OWN sessions only. Mirrors lucien's ingest-recent.ts.
+//
+// See docs/design/memory-system.md → "Ingestion — mlx-bun sessions only".
+
+import type { SynthesisEvent } from "./pipeline";
+
+export interface IngestResult {
+  scanned: number;
+  ingested: number;
+}
+
+/** Read new pi sessions into memory.sqlite. STUB: returns zero, does nothing. */
+export async function ingestSessions(_onEvent?: (e: SynthesisEvent) => void): Promise<IngestResult> {
+  // TODO(M1): SessionManager.listAll() → normalize → upsert into conversations,
+  // advancing the per-source watermark. pi's session format is the contract.
+  throw new Error("memory.ingest is not implemented yet (M1)");
+}

package/src/memory/pipeline.ts

@@ -0,0 +1,63 @@
+// mlx-bun memory — synthesis pipeline orchestrator.  ⚠️ M1 STUB.
+//
+// Ties the stages together: ingest → chunk → cluster → synthesize → wikify →
+// changelog (the lucien nightly chain, reimplemented on the local model). What
+// `mlx-bun memory synthesize` and the nightly launchd job invoke.
+//
+// CURRENTLY A NO-OP: it logs each stage as stubbed and writes nothing, so the
+// scheduled job runs cleanly tonight and every night without touching your
+// vault. When M1 lands, flip runSynthesis to actually call the stage modules
+// (each currently throws "not implemented") and move it onto the src/jobs/*
+// background runner with SSE progress, per the design doc.
+//
+// See docs/design/memory-system.md → "The nightly pipeline".
+
+export type SynthesisStage = "ingest" | "chunk" | "cluster" | "synthesize" | "wikify" | "changelog";
+
+export interface SynthesisEvent {
+  type: "stage" | "log" | "done";
+  stage?: SynthesisStage;
+  message: string;
+}
+
+export interface SynthesisOptions {
+  /** Only synthesize conversations newer than this (ISO date / conv cursor). */
+  since?: string;
+  /** Override the synthesis model (default e4b). */
+  model?: string;
+  /** Plan only — never write the vault. */
+  dryRun?: boolean;
+}
+
+export interface SynthesisSummary {
+  /** False while stubbed; flip to true when the stages are implemented. */
+  implemented: boolean;
+  stages: SynthesisStage[];
+  note: string;
+}
+
+const STAGES: SynthesisStage[] = ["ingest", "chunk", "cluster", "synthesize", "wikify", "changelog"];
+
+/**
+ * Run the synthesis pipeline. STUB: walks the stage list, emits a `stage` event
+ * for each (marked skipped), and returns implemented:false. Deterministic and
+ * side-effect-free, so it's safe for the nightly job to call today.
+ */
+export async function runSynthesis(
+  opts: SynthesisOptions = {},
+  onEvent?: (e: SynthesisEvent) => void,
+): Promise<SynthesisSummary> {
+  onEvent?.({
+    type: "log",
+    message: `memory synthesis is not implemented yet (M1)${opts.dryRun ? " · dry-run" : ""} — running as a no-op, nothing will be written.`,
+  });
+  for (const stage of STAGES) {
+    onEvent?.({ type: "stage", stage, message: `${stage}: stubbed (skipped)` });
+  }
+  onEvent?.({ type: "done", message: "no articles were written (synthesis is stubbed)." });
+  return {
+    implemented: false,
+    stages: STAGES,
+    note: "M1 stub — pipeline + scheduling wired; stage modules not implemented.",
+  };
+}

package/src/memory/prompts.ts

@@ -0,0 +1,23 @@
+// mlx-bun memory — per-stage synthesis prompts.  ⚠️ M1 STUBS.
+//
+// Each stage's prompt is the seam where a future specialized memory LoRA plugs
+// in (today: base model + prompt; later: adapterScoped per stage). Editorial
+// conventions are NOT duplicated here — they live in the vault's Meta/ pages
+// (Editorial_Guidelines.md, Article_Conventions.md, Buckets.md,
+// Topics_to_Ignore.md) and are read in at synthesis time so the system governs
+// itself in its own substrate.
+//
+// See docs/design/memory-system.md → "Synthesis runs on the local model".
+
+export const CHUNK_PROMPT = `TODO(M1): segment a conversation transcript into topic-coherent chunks.`;
+
+export const CLUSTER_PROMPT = `TODO(M1): assign a chunk to an existing bucket from Meta/Buckets.md, or propose a new bucket.`;
+
+export const SYNTHESIZE_PROMPT = `TODO(M1): integrate new chunks into the bucket's article. Preserve prose and every conv: citation; integrate, don't overwrite; respect the ≥70% word-count floor and the editorial gate.`;
+
+/** Placeholder registry so the pipeline can reference prompts by stage name. */
+export const STAGE_PROMPTS: Record<string, string> = {
+  chunk: CHUNK_PROMPT,
+  cluster: CLUSTER_PROMPT,
+  synthesize: SYNTHESIZE_PROMPT,
+};