mlx-bun 0.0.4 → 0.0.6

package/src/eval/tasks/ifeval.ts

@@ -67,7 +67,12 @@ function reEscape(s: string): string {
 // \b\w+\b — JS \w is ASCII-ish like Python's default re (no re.UNICODE here).
 function checkLengthWords(response: string, kw: Kw): boolean {
   const n = (response.match(/\b\w+\b/g) ?? []).length;
-  const rel = kwStr(kw, "relation", "at least");
+  // Python `kw.get("relation", "at least")`: the default applies ONLY when the key is
+  // ABSENT — a present-but-None `relation` (as the frozen kwargs supply for letter/
+  // paragraph/capital checks) returns None, and optiq's `"least" in None` RAISES →
+  // _verify_response returns False. Replicate exactly: rel is null → `.includes` throws
+  // below → verifyResponse catches → FAIL. (kwStr's None→default was the leniency.)
+  const rel = ("relation" in kw ? kw["relation"] : "at least") as string;
   const target = kwInt(kw, "num_words", 0);
   return rel.includes("least") ? n >= target : n <= target;
 }
@@ -75,14 +80,24 @@ function checkLengthWords(response: string, kw: Kw): boolean {
 function checkLengthSentences(response: string, kw: Kw): boolean {
   // Python: re.split(r"[.!?]+\s+", response.strip())
   const n = response.trim().split(/[.!?]+\s+/).length;
-  const rel = kwStr(kw, "relation", "at least");
+  // Python `kw.get("relation", "at least")`: the default applies ONLY when the key is
+  // ABSENT — a present-but-None `relation` (as the frozen kwargs supply for letter/
+  // paragraph/capital checks) returns None, and optiq's `"least" in None` RAISES →
+  // _verify_response returns False. Replicate exactly: rel is null → `.includes` throws
+  // below → verifyResponse catches → FAIL. (kwStr's None→default was the leniency.)
+  const rel = ("relation" in kw ? kw["relation"] : "at least") as string;
   const target = kwInt(kw, "num_sentences", 0);
   return rel.includes("least") ? n >= target : n <= target;
 }
 
 function checkLengthParagraphs(response: string, kw: Kw): boolean {
   const n = response.split("\n\n").filter((p) => p.trim()).length;
-  const rel = kwStr(kw, "relation", "at least");
+  // Python `kw.get("relation", "at least")`: the default applies ONLY when the key is
+  // ABSENT — a present-but-None `relation` (as the frozen kwargs supply for letter/
+  // paragraph/capital checks) returns None, and optiq's `"least" in None` RAISES →
+  // _verify_response returns False. Replicate exactly: rel is null → `.includes` throws
+  // below → verifyResponse catches → FAIL. (kwStr's None→default was the leniency.)
+  const rel = ("relation" in kw ? kw["relation"] : "at least") as string;
   const target = kwInt(kw, "num_paragraphs", 0);
   return rel.includes("least") ? n >= target : n <= target;
 }
@@ -101,7 +116,12 @@ function checkKeywordsForbidden(response: string, kw: Kw): boolean {
 
 function checkKeywordFrequency(response: string, kw: Kw): boolean {
   const kwWord = kwStr(kw, "keyword", "").toLowerCase();
-  const rel = kwStr(kw, "relation", "at least");
+  // Python `kw.get("relation", "at least")`: the default applies ONLY when the key is
+  // ABSENT — a present-but-None `relation` (as the frozen kwargs supply for letter/
+  // paragraph/capital checks) returns None, and optiq's `"least" in None` RAISES →
+  // _verify_response returns False. Replicate exactly: rel is null → `.includes` throws
+  // below → verifyResponse catches → FAIL. (kwStr's None→default was the leniency.)
+  const rel = ("relation" in kw ? kw["relation"] : "at least") as string;
   const target = kwInt(kw, "frequency", 0);
   const n = (response.toLowerCase().match(new RegExp(`\\b${reEscape(kwWord)}\\b`, "g")) ?? []).length;
   return rel.includes("least") ? n >= target : n <= target;
@@ -109,7 +129,12 @@ function checkKeywordFrequency(response: string, kw: Kw): boolean {
 
 function checkLetterFrequency(response: string, kw: Kw): boolean {
   const letter = kwStr(kw, "letter", "").toLowerCase();
-  const rel = kwStr(kw, "relation", "at least");
+  // Python `kw.get("relation", "at least")`: the default applies ONLY when the key is
+  // ABSENT — a present-but-None `relation` (as the frozen kwargs supply for letter/
+  // paragraph/capital checks) returns None, and optiq's `"least" in None` RAISES →
+  // _verify_response returns False. Replicate exactly: rel is null → `.includes` throws
+  // below → verifyResponse catches → FAIL. (kwStr's None→default was the leniency.)
+  const rel = ("relation" in kw ? kw["relation"] : "at least") as string;
   // Python: int(kw.get("let_frequency", kw.get("frequency", 0)))
   const target = kw["let_frequency"] !== undefined && kw["let_frequency"] !== null
     ? kwInt(kw, "let_frequency", 0)
@@ -124,7 +149,12 @@ function checkCapitalWordsCount(response: string, kw: Kw): boolean {
   // Python: w.isupper() and len(w) > 1 — str.isupper() requires ≥1 cased char,
   // all cased chars uppercase. Approximate: has a letter, no lowercase letter.
   const n = words.filter((w) => w.length > 1 && /[A-Za-z]/.test(w) && w === w.toUpperCase()).length;
-  const rel = kwStr(kw, "relation", "at least");
+  // Python `kw.get("relation", "at least")`: the default applies ONLY when the key is
+  // ABSENT — a present-but-None `relation` (as the frozen kwargs supply for letter/
+  // paragraph/capital checks) returns None, and optiq's `"least" in None` RAISES →
+  // _verify_response returns False. Replicate exactly: rel is null → `.includes` throws
+  // below → verifyResponse catches → FAIL. (kwStr's None→default was the leniency.)
+  const rel = ("relation" in kw ? kw["relation"] : "at least") as string;
   // Python: int(kw.get("capital_frequency", kw.get("frequency", 0)))
   const target = kw["capital_frequency"] !== undefined && kw["capital_frequency"] !== null
     ? kwInt(kw, "capital_frequency", 0)
@@ -251,11 +281,12 @@ const VERIFIERS = new Map<string, Verifier>([
 /**
  * Verify one response against its instruction list.
  * Returns { pass, unhandled }. An instruction_id with no registered verifier is
- * counted as a FAIL (NOT silently skipped) and recorded in `unhandled`. This is
- * stricter than ifeval.py (which treats unhandled as pass); per the mlx-bun
- * harness rule, anything we cannot port must fail rather than bias high.
+ * treated as a PASS (does not affect `pass`) and only recorded in `unhandled` for
+ * visibility — matching ifeval.py's `continue` policy EXACTLY, so our strict_acc
+ * reproduces optiq's. (Our VERIFIERS map mirrors his `_VERIFIERS` 1:1, so unhandled
+ * is empty for the standard IFEval set anyway; this only matters for parity safety.)
  */
-function verifyResponse(
+export function verifyResponse(
   response: string,
   instructionIds: string[],
   kwargsList: Kw[],
@@ -268,8 +299,7 @@ function verifyResponse(
     const verifier = VERIFIERS.get(iid);
     if (verifier === undefined) {
       unhandled.push(iid);
-      pass = false; // count unported instruction as a failure
-      continue;
+      continue; // ifeval.py: treat an unhandled instruction as PASS (don't touch `pass`)
     }
     try {
       if (!verifier(response, kw)) pass = false;
@@ -281,7 +311,7 @@ function verifyResponse(
 }
 
 /** Loose-mode preprocessing — mirrors ifeval.py `_loose_clean`. */
-function looseClean(response: string): string {
+export function looseClean(response: string): string {
   let s = response.trim();
   // Python: re.sub(r"^(Sure|Here|Of course)[,!.]\s*[^\n]*\n+", "", s, count=1)
   s = s.replace(/^(Sure|Here|Of course)[,!.]\s*[^\n]*\n+/, "");
@@ -293,22 +323,26 @@ function looseClean(response: string): string {
 }
 
 /** Strip a leading thinking block, matching ifeval.py's split on </think>. */
-function stripThinking(response: string): string {
+export function stripThinking(response: string): string {
   const idx = response.indexOf("</think>");
   return idx === -1 ? response : response.slice(idx + "</think>".length);
 }
 
 export async function evaluateIfeval(
   tm: TaskModel,
-  opts: { nSamples?: number; maxTokens?: number; seed?: number } = {},
+  opts: { nSamples?: number; maxTokens?: number; seed?: number; frozen?: boolean } = {},
 ): Promise<IfevalResult> {
   // ifeval.py: max_tokens default 512.
   const maxTokens = opts.maxTokens ?? 512;
   const seed = opts.seed ?? 42;
 
-  const rows = loadJsonl<IfevalRow>("ifeval");
+  // Optiq-parity mode (DEFAULT): score optiq's EXACT prompt set (he runs the full
+  // IFEval split — 541 prompts) through OUR verifiers + OUR chat template.
+  // MLX_BUN_IFEVAL_FROZEN=0 reverts to our own copy + optional sampling.
+  const useFrozen = opts.frozen ?? (process.env.MLX_BUN_IFEVAL_FROZEN !== "0");
+  const rows = loadJsonl<IfevalRow>(useFrozen ? "ifeval_optiq_frozen" : "ifeval");
   const nSamples = opts.nSamples ?? rows.length;
-  const idx = nSamples < rows.length
+  const idx = (!useFrozen && nSamples < rows.length)
     ? sampleIndices(rows.length, nSamples, seed)
     : Array.from({ length: rows.length }, (_, i) => i);
 
@@ -341,7 +375,7 @@ export async function evaluateIfeval(
   if (unhandledCounts.size) {
     const top = [...unhandledCounts.entries()].sort((a, b) => b[1] - a[1]);
     process.stderr.write(
-      "  ifeval unported instruction ids (counted as FAIL): " +
+      "  ifeval unported instruction ids (treated as PASS, per optiq): " +
         top.map(([k, v]) => `${k}(${v})`).join(", ") +
         "\n",
     );

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/eval/ultrafeedback.ts

@@ -0,0 +1,60 @@
+// UltraFeedback (binarized) → mlx-bun preference rows. The binarized schema
+// (HuggingFaceH4/ultrafeedback_binarized) gives each row a `prompt` plus
+// `chosen` / `rejected` as message lists; ORPO wants {prompt, chosen, rejected}
+// where chosen/rejected are the assistant response *strings*. Pure transform so
+// it is unit-testable without downloading the dataset.
+
+export interface PreferenceRow {
+  prompt: string;
+  chosen: string;
+  rejected: string;
+}
+
+type Msg = { role?: string; content?: string };
+
+/** Last assistant message content in a binarized message list (the response). */
+function lastAssistant(messages: unknown): string | null {
+  if (!Array.isArray(messages)) return null;
+  for (let i = messages.length - 1; i >= 0; i--) {
+    const m = messages[i] as Msg;
+    if (m && m.role === "assistant" && typeof m.content === "string" && m.content.trim())
+      return m.content;
+  }
+  return null;
+}
+
+/** Convert one binarized UltraFeedback row → a preference row, or null if it is
+ *  unusable (missing fields, or chosen == rejected so there is no signal). */
+export function toPreferenceRow(row: Record<string, unknown>): PreferenceRow | null {
+  const prompt = typeof row.prompt === "string" ? row.prompt : null;
+  const chosen = lastAssistant(row.chosen);
+  const rejected = lastAssistant(row.rejected);
+  if (!prompt || !chosen || !rejected) return null;
+  if (chosen.trim() === rejected.trim()) return null;
+  return { prompt, chosen, rejected };
+}
+
+/** Rough token estimate (≈4 chars/token) for a cheap length pre-filter before
+ *  tokenization — keeps the curated set within a target context without loading
+ *  a tokenizer. */
+export function approxTokens(text: string): number {
+  return Math.ceil(text.length / 4);
+}
+
+/** Filter + convert a stream of binarized rows. Drops unusable rows and ones
+ *  whose prompt+chosen or prompt+rejected approx-exceeds `maxApproxTokens`. */
+export function curate(
+  rows: Record<string, unknown>[],
+  maxApproxTokens = 2048,
+): PreferenceRow[] {
+  const out: PreferenceRow[] = [];
+  for (const row of rows) {
+    const p = toPreferenceRow(row);
+    if (!p) continue;
+    const base = approxTokens(p.prompt);
+    if (base + approxTokens(p.chosen) > maxApproxTokens) continue;
+    if (base + approxTokens(p.rejected) > maxApproxTokens) continue;
+    out.push(p);
+  }
+  return out;
+}

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/lora.ts

@@ -37,7 +37,7 @@ export interface AdapterInfo {
  *  materialized mlx arrays. Tensor names come from our header parser;
  *  arrays from mlx's native loader via the proven map-get pattern
  *  (weights.ts). The map and header mmap are freed before returning. */
-function loadAdapterTensors(file: string): Map<string, MlxArray> {
+export function loadAdapterTensors(file: string): Map<string, MlxArray> {
   const sf = SafetensorsFile.open(file);
   const names = [...sf.tensors.keys()];
   sf.mmap.unmap();
@@ -74,18 +74,21 @@ function loadAdapterTensors(file: string): Map<string, MlxArray> {
 }
 
 /** Adapter scale from its config: mlx-lm writes lora_parameters.scale;
- *  PEFT writes lora_alpha + r (scale = alpha / r). Reference: mount.py. */
-async function readAdapterScale(dir: string): Promise<{ scale: number; rank: number | null }> {
+ *  PEFT writes lora_alpha + r (scale = alpha / r). Reference: mount.py.
+ *  `rsLora` (recorded at train time) means the effective per-layer scale is
+ *  α/√rank — the caller divides by √(that layer's rank). */
+async function readAdapterScale(dir: string): Promise<{ scale: number; rank: number | null; rsLora: boolean }> {
   for (const name of ["optiq_lora_config.json", "adapter_config.json"]) {
     const f = Bun.file(`${dir}/${name}`);
     if (await f.exists()) {
       const cfg = (await f.json()) as Record<string, any>;
       const lp = cfg.lora_parameters;
+      const rsLora = Boolean(cfg.rs_lora ?? lp?.rs_lora ?? false);
       if (lp && typeof lp === "object")
-        return { scale: Number(lp.scale ?? 20.0), rank: lp.rank ?? null };
+        return { scale: Number(lp.scale ?? 20.0), rank: lp.rank ?? null, rsLora };
       const alpha = Number(cfg.lora_alpha ?? 16);
       const r = Number(cfg.r ?? 8);
-      return { scale: r ? alpha / r : 1.0, rank: r || null };
+      return { scale: r ? alpha / r : 1.0, rank: r || null, rsLora };
     }
   }
   throw new Error(`no adapter_config.json in ${dir}`);
@@ -200,7 +203,7 @@ export class AdapterManager {
     if (!existsSync(path)) throw new Error(`adapter dir not found: ${path}`);
 
     const weightsFile = adapterWeightsFile(path);
-    const { scale, rank: configRank } = await readAdapterScale(path);
+    const { scale, rank: configRank, rsLora } = await readAdapterScale(path);
     const tensors = loadAdapterTensors(weightsFile);
     const targets = this.#model.loraTargets();
 
@@ -243,7 +246,9 @@ export class AdapterManager {
             `vs base [in ${linear.inFeatures}, out ${linear.outFeatures}]; ` +
             `was this adapter trained for a different base model?`,
           );
-        validated.push({ linear, lw: { a, b, scale, rank } });
+        // rsLoRA: effective per-layer scale is α/√rank (matches training).
+        const effScale = rsLora ? scale / Math.sqrt(rank) : scale;
+        validated.push({ linear, lw: { a, b, scale: effScale, rank } });
       }
       if (validated.length === 0)
         throw new Error(

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)");
+}