mlx-bun 0.0.5 → 0.0.7

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (57) hide show
  1. package/README.md +11 -2
  2. package/docs/dag/training-inference-map.html +575 -0
  3. package/docs/design/memory-system.md +193 -295
  4. package/docs/design/orpo-dynamic-lambda.md +179 -0
  5. package/docs/design/orpo-future-enhancements.md +42 -0
  6. package/docs/design/orpo-training.md +13 -0
  7. package/docs/design/parity-tier-dag.md +143 -0
  8. package/docs/design/w4a16-compute-precision-spike.md +232 -0
  9. package/docs/investigations/orpo-base-uf-experiment-and-directions.md +120 -0
  10. package/docs/investigations/orpo-flash-cce-pin-leak.md +167 -0
  11. package/docs/investigations/orpo-uf-testing-handoff.md +126 -0
  12. package/docs/investigations/steel-flash-cce-handoff.md +58 -0
  13. package/docs/investigations/trainer-validation-experiment.md +173 -0
  14. package/docs/reference/memory.md +77 -0
  15. package/docs/reference/orpo-quickstart.md +44 -5
  16. package/docs/reference/training.md +9 -5
  17. package/package.json +3 -3
  18. package/src/cli.ts +651 -18
  19. package/src/eval/runner.ts +76 -6
  20. package/src/eval/tasks/bfcl.ts +49 -15
  21. package/src/eval/tasks/gsm8k.ts +21 -9
  22. package/src/eval/tasks/hashhop.ts +33 -4
  23. package/src/eval/tasks/humaneval.ts +30 -8
  24. package/src/eval/tasks/ifeval.ts +52 -18
  25. package/src/eval/tasks/mmlu.ts +85 -74
  26. package/src/jobs/runner.ts +6 -0
  27. package/src/memory/chunk.ts +25 -0
  28. package/src/memory/cluster.ts +24 -0
  29. package/src/memory/db.ts +114 -0
  30. package/src/memory/ingest.ts +22 -0
  31. package/src/memory/pipeline.ts +63 -0
  32. package/src/memory/prompts.ts +23 -0
  33. package/src/memory/schedule.ts +163 -0
  34. package/src/memory/synthesize.ts +29 -0
  35. package/src/memory/tools.ts +399 -0
  36. package/src/memory/vault.ts +694 -0
  37. package/src/memory/wikify.ts +23 -0
  38. package/src/mlx/metal-kernel.ts +13 -0
  39. package/src/model/flash-attention.ts +5 -3
  40. package/src/pi-provider.ts +29 -13
  41. package/src/pi-session.ts +81 -0
  42. package/src/pi-terminal.ts +13 -24
  43. package/src/pi-web.ts +246 -166
  44. package/src/quantize/job.ts +17 -2
  45. package/src/quantize/quantizer.ts +26 -2
  46. package/src/server.ts +213 -44
  47. package/src/tokenizer.ts +6 -2
  48. package/src/train/dataset.ts +11 -3
  49. package/src/train/flash-cce.ts +18 -5
  50. package/src/train/job.ts +41 -4
  51. package/src/train/loss.ts +10 -4
  52. package/src/train/trainer.ts +62 -3
  53. package/src/train/watch.ts +542 -0
  54. package/src/web/app.html +264 -14
  55. package/src/web/skills/memory/SKILL.md +97 -0
  56. package/src/web/skills/web-research/SKILL.md +3 -1
  57. package/src/web/skills.ts +46 -8
package/src/cli.ts CHANGED
@@ -17,6 +17,7 @@ import pkg from "../package.json" with { type: "json" };
17
17
  import { renderHelp } from "./tui";
18
18
  import { isSupportedModelRecord } from "./model/support";
19
19
  import { perfKernelEnabled } from "./model/fused-decode-kernel";
20
+ import type { JobEvent } from "./jobs/types";
20
21
 
21
22
  const argv = process.argv.slice(2);
22
23
  // The appliance path: naked `mlx-bun` (or only options, e.g.
@@ -39,6 +40,8 @@ Commands:
39
40
  fit Will a model fit this machine? Memory + speed assessment
40
41
  scan Re-index the Hugging Face cache
41
42
  harness Connect your own pi install to the local mlx-bun server
43
+ train Fine-tune a LoRA adapter (SFT / DPO / ORPO) on your data
44
+ memory Set up + manage your local AI's personal wiki memory
42
45
  benchmark Measure decode/prefill speed of OUR stack on this machine
43
46
  evals Show recent benchmark runs (all stacks)
44
47
  help Show help for a command (also: mlx-bun <command> --help)
@@ -51,6 +54,7 @@ Examples:
51
54
  mlx-bun # download (first run) + serve + open the chat UI
52
55
  mlx-bun pi # start (download if needed) and chat with an agent
53
56
  mlx-bun serve 12B # serve the 12B; status page at http://localhost:8090/
57
+ mlx-bun train e4b --data ./prefs # ORPO LoRA fine-tune on {prompt,chosen,rejected} data
54
58
  mlx-bun get mlx-community/gemma-4-12B-it-OptiQ-4bit`;
55
59
 
56
60
  const SERVER_FLAGS = `Server options:
@@ -235,9 +239,84 @@ Options:
235
239
 
236
240
  Table view: when, model, bench kind, KV mode, decode tok/s, TTFT,
237
241
  peak memory, commit. Runs are written by ./benchmark.sh.`,
242
+
243
+ train: `mlx-bun train — fine-tune a LoRA adapter on your data
244
+
245
+ Usage: mlx-bun train <model> --data <dir> [options]
246
+
247
+ <model> Model to fine-tune (query or snapshot path; auto-picks
248
+ the default model if omitted)
249
+
250
+ Data: <dir> must contain train.jsonl (+ optional valid.jsonl). Rows are
251
+ {prompt, chosen, rejected} for dpo/orpo, or {messages|text} for sft.
252
+
253
+ Options:
254
+ --data <dir> Dataset dir with train.jsonl (required)
255
+ --method <m> sft | dpo | orpo [default: orpo]
256
+ --adapter <dir> Output adapter dir
257
+ [default: ~/.cache/mlx-bun/mlx-bun-finetunes/<method>-<model>]
258
+ --iters <n> Training iterations [default: 100]
259
+ --lr <f> Learning rate [default: orpo 1e-5 · dpo 5e-5 · sft 2e-4]
260
+ --rank <n> LoRA rank [default: orpo 16 · else 8]
261
+ --scale <f> LoRA scale [default: orpo 2.0 · else 1.0]
262
+ --seq <n> Max sequence length [default: gemma 8192 · else 4096]
263
+ --batch <n> Batch size [default: 1]
264
+ --lambda <f> ORPO odds-ratio weight [default: 0.1]
265
+ --seg <n> Layers per segment (segmented backward; orpo default 2)
266
+ --save-every <n> Crash-safe mountable checkpoint every n steps
267
+ --resume <dir> Warm-start LoRA weights from a checkpoint/adapter dir
268
+ --no-flash Disable the flash-CCE Metal head (use the MLX fused head)
269
+ --no-prefix Disable prefix-sharing (two-forward branches)
270
+ --no-segment Disable the segmented backward (hold all activations)
271
+ --dry-run Inspect the dataset + print the resolved plan, don't train
272
+
273
+ ORPO defaults run the full stack: the [M,vocab]-free flash-CCE head +
274
+ prefix-sharing + segmented backward (each falls back + logs if a row's
275
+ preconditions aren't met). Gemma/e4b sets its required training env flags
276
+ automatically. Adapters are mountable directly (mlx-bun serve --adapter).
277
+ Run long jobs detached from your own shell: nohup mlx-bun train … &`,
278
+ memory: `mlx-bun memory — your local AI's personal wiki
279
+
280
+ Usage: mlx-bun memory <subcommand>
281
+
282
+ A local, durable memory for the assistant: a wiki of Markdown articles
283
+ (~/.mlx-bun/wiki) it reads to remember your projects, people, and history
284
+ across sessions. It's yours — git-tracked, editable in any tool (Obsidian
285
+ opens it as a vault), and it never leaves the machine. A read-only Reference/
286
+ folder symlinks mlx-bun's own docs into memory on day one. Once set up, it
287
+ loads automatically into every \`mlx-bun pi\` and \`mlx-bun serve\` session.
288
+
289
+ Subcommands:
290
+ init, setup Create the wiki + walk through setup (idempotent);
291
+ offers to import an existing vault and install the
292
+ nightly synthesis job
293
+ status Path, article count, git + schedule state
294
+ open, browse [article]
295
+ Open the wiki, or a specific article, in Obsidian
296
+ (falls back to Finder / the default Markdown app)
297
+ list List article titles + read-only Reference docs
298
+ search <query> Search articles from the terminal
299
+ toc <article> Print an article's headings + anchors
300
+ section <article> <anchor>
301
+ Print one article section
302
+ links <article> Show resolved outbound + inbound wikilinks
303
+ read <article> Print an article (stem, e.g. Archie_Project)
304
+ synthesize Run the synthesis pipeline now (--since, --model,
305
+ --dry-run). STUBBED — a no-op until M1 lands.
306
+ schedule Install the nightly launchd job (--at HH:MM [03:00])
307
+ unschedule Remove the nightly launchd job
308
+
309
+ The read path is live: the assistant reads a wiki you set up by hand or
310
+ import during \`memory init\`. Ask the assistant to "open my memory" or run
311
+ \`mlx-bun memory open\` to browse it in Obsidian/Finder, or
312
+ \`mlx-bun memory open <article>\` to jump to a specific page. Synthesis
313
+ (conversations → articles) is stubbed; scheduling is real, so the nightly
314
+ job is wired and ready and starts producing articles the moment synthesis
315
+ lands.`,
238
316
  };
239
317
 
240
318
  HELP.bench = HELP.benchmark!;
319
+ HELP.setup = HELP.memory!;
241
320
 
242
321
  function printHelp(topic?: string): never {
243
322
  if (topic && HELP[topic]) console.log(renderHelp(HELP[topic]));
@@ -368,6 +447,63 @@ function openChatUi(url: string, hostPort: string): void {
368
447
  /** Server/runtime flags shared by every mode that loads a model
369
448
  * (serve, pi). Env levers are set here so they're in place before the
370
449
  * generate/compiled-decode modules read them. */
450
+ /** Resolve the decode ROUTE: a tier alias (--l1/--l2/--l3) sets the whole route,
451
+ * and an explicit per-fork flag (--perf-kernel/--kv-quant/…) overrides the alias.
452
+ * The tier IS the intent; the forks are the escape hatch. Sets the decode env
453
+ * levers and returns the kv-quant mode. See docs/design/parity-tier-dag.md.
454
+ * Each tier is a GUARANTEE about which reference you reproduce bit-for-bit:
455
+ * --l1 bit-for-bit IDENTICAL to mlx-lm — drop-in replacement for mlx-lm
456
+ * --l2 bit-for-bit IDENTICAL to mlx-optiq — drop-in replacement for mlx-optiq.
457
+ * optiq SHIPS the fused quantized decode (the quant-KV goldens track
458
+ * perf-kernel ON), so L2 turns it ON — not the unfused reference.
459
+ * --l3 best performance, NO bit-exact guarantee — correct, but gated by KL +
460
+ * tests, because mlx-optiq has no analogy for what L3 does.
461
+ * L2 ⊂ L3 (L3 = L2 + compiled). L1 is a separate target (bf16, unfused). */
462
+ function applyDecodeRoute(): { kvQuant?: "off" | "config" | number } {
463
+ const onOff = (name: string): boolean | null => {
464
+ const v = opt(name); if (v == null) return null;
465
+ if (v === "on" || v === "1" || v === "true") return true;
466
+ if (v === "off" || v === "0" || v === "false") return false;
467
+ console.error(`--${name} expects on|off (got "${v}")`); process.exit(1);
468
+ };
469
+ type Preset = { kv: "off" | "config"; perf: boolean; fusedSdpa: boolean; compiled: boolean; fusedDecode: boolean };
470
+ const TIERS: Record<string, Preset> = {
471
+ // A tier is the GUARANTEE; within it we default to the FAST kernel that still
472
+ // holds it, and expose the slow one as an opt-in. compiled-decode is BIT-EXACT
473
+ // with uncompiled (compiled-decode.test: on==off) → free speed, so it's ON in
474
+ // EVERY tier (--compiled-decode off is the slow, same-tier opt-in). perf-kernel
475
+ // + fused-sdpa match OPTIQ (not the unfused reference; fused-sdpa.test) and are
476
+ // quantized-KV only → they're the L2 bridge, no-ops on L1's bf16. NOTE: for the
477
+ // DECODE path L2 == L3 — every decode kernel matches mlx-lm OR optiq. L3's real
478
+ // territory is the no-oracle FEATURES (HLG sampler, expert offload, batched
479
+ // mixed-precision) gated by KL+tests, which live off this decode axis.
480
+ l1: { kv: "off", perf: false, fusedSdpa: false, compiled: true, fusedDecode: false }, // = mlx-lm bit-for-bit (bf16)
481
+ l2: { kv: "config", perf: true, fusedSdpa: true, compiled: true, fusedDecode: false }, // = mlx-optiq bit-for-bit
482
+ l3: { kv: "config", perf: true, fusedSdpa: true, compiled: true, fusedDecode: false }, // best perf (decode == L2 today)
483
+ };
484
+ const tier = flag("l1") ? "l1" : flag("l2") ? "l2" : flag("l3") ? "l3" : null;
485
+ const p = tier ? TIERS[tier]! : null;
486
+ const pick = (name: string, base: boolean | undefined): boolean | null => {
487
+ const ex = onOff(name); return ex !== null ? ex : (base ?? null); // explicit flag wins, else tier preset
488
+ };
489
+ const set = (env: string, val: boolean | null, invert = false) => {
490
+ if (val !== null) process.env[env] = (invert ? !val : val) ? "1" : "0";
491
+ };
492
+ set("MLX_BUN_COMPILED_DECODE", pick("compiled-decode", p?.compiled));
493
+ set("MLX_BUN_PERF_KERNEL", pick("perf-kernel", p?.perf));
494
+ set("MLX_BUN_FUSED_DECODE", pick("fused-decode", p?.fusedDecode));
495
+ set("MLX_BUN_NO_FUSED_SDPA", pick("fused-sdpa", p?.fusedSdpa), true); // inverted env
496
+ const kv = opt("kv-quant"); // explicit --kv-quant overrides the tier
497
+ if (kv === "off") return { kvQuant: "off" };
498
+ if (kv === "config") return { kvQuant: "config" };
499
+ if (kv) {
500
+ const bits = Number(kv);
501
+ if (![4, 8].includes(bits)) { console.error(`--kv-quant expects config|off|4|8 (got "${kv}")`); process.exit(1); }
502
+ return { kvQuant: bits };
503
+ }
504
+ return p ? { kvQuant: p.kv } : {};
505
+ }
506
+
371
507
  function serverRuntimeFlags(): { port: number; serverOptions: import("./server").ServerOptions } {
372
508
  const onOff = (name: string): boolean | null => {
373
509
  const v = opt(name);
@@ -377,21 +513,14 @@ function serverRuntimeFlags(): { port: number; serverOptions: import("./server")
377
513
  console.error(`--${name} expects on|off (got "${v}")`);
378
514
  process.exit(1);
379
515
  };
380
- const cd = onOff("compiled-decode");
381
- if (cd !== null) process.env.MLX_BUN_COMPILED_DECODE = cd ? "1" : "0";
382
- const pk = onOff("perf-kernel");
383
- if (pk !== null) process.env.MLX_BUN_PERF_KERNEL = pk ? "1" : "0";
384
- const fd = onOff("fused-decode");
385
- if (fd !== null) process.env.MLX_BUN_FUSED_DECODE = fd ? "1" : "0";
386
- const fs = onOff("fused-sdpa");
387
- if (fs !== null) process.env.MLX_BUN_NO_FUSED_SDPA = fs ? "0" : "1"; // inverted env
516
+ const route = applyDecodeRoute(); // --l1/--l2/--l3 tier alias, with per-fork flags overriding
388
517
  if (flag("force-wire")) process.env.MLX_BUN_FORCE_WIRE = "1";
389
518
 
390
519
  const serverOptions: import("./server").ServerOptions = {};
391
520
  const budgetGB = Number(opt("memory-budget", "0"));
392
521
  if (budgetGB > 0) serverOptions.memoryBudgetBytes = budgetGB * 1e9;
393
- const pcGB = Number(opt("prompt-cache", "0"));
394
- if (pcGB > 0) serverOptions.promptCacheBytes = pcGB * 2 ** 30;
522
+ const pcRaw = opt("prompt-cache"); // null if absent; an explicit value (incl. 0) wins → `--prompt-cache 0` DISABLES the cache
523
+ if (pcRaw !== null) serverOptions.promptCacheBytes = Math.max(0, Number(pcRaw)) * 2 ** 30;
395
524
  // --batch N: max concurrent requests batched through the mlx-lm-parity
396
525
  // engine (N=1 = today's serial path). --decode-concurrency is accepted as
397
526
  // an mlx_lm.server-compatible alias (drop-in).
@@ -404,14 +533,7 @@ function serverRuntimeFlags(): { port: number; serverOptions: import("./server")
404
533
  }
405
534
  serverOptions.batch = n;
406
535
  }
407
- const kv = opt("kv-quant");
408
- if (kv === "off") serverOptions.kvQuant = "off";
409
- else if (kv === "config") serverOptions.kvQuant = "config";
410
- else if (kv) {
411
- const bits = Number(kv);
412
- if (![4, 8].includes(bits)) { console.error(`--kv-quant expects config|off|4|8 (got "${kv}")`); process.exit(1); }
413
- serverOptions.kvQuant = bits;
414
- }
536
+ if (route.kvQuant !== undefined) serverOptions.kvQuant = route.kvQuant;
415
537
  const host = opt("host");
416
538
  if (host) serverOptions.hostname = host;
417
539
  // Server-wide default for the chat template's enable_thinking variable
@@ -733,6 +855,34 @@ switch (cmd) {
733
855
  break;
734
856
  }
735
857
 
858
+ case "gen":
859
+ case "generate": {
860
+ // RAW generate — a direct one-shot entry point (no server, no chat UI). The
861
+ // three entry points (raw / OpenAI API / chat UI) differ only in how params
862
+ // are POPULATED; here every param is explicit. Decode-path levers mirror
863
+ // serve so you can pin the route (mlx-lm compat = --perf-kernel off
864
+ // --fused-sdpa off --compiled-decode off). KV uses the bf16 bit-exact path.
865
+ applyDecodeRoute(); // --l1/--l2/--l3 + per-fork overrides set the decode env levers (KV stays bf16 via generateText)
866
+
867
+ const prompt = opt("prompt") ?? positional(1);
868
+ if (!prompt) {
869
+ console.error('usage: mlx-bun generate [query] --prompt "…" [--raw] [--max-tokens N]');
870
+ console.error(' sampling: --temperature N --top-p N --top-k N --seed N');
871
+ console.error(' decode path: --perf-kernel off --fused-sdpa off --compiled-decode off (= mlx-lm compat)');
872
+ process.exit(1);
873
+ }
874
+ const { loadTaskModel, generateText } = await import("./eval/runner");
875
+ const tm = await loadTaskModel(positional(0) ?? opt("query") ?? "");
876
+ const num = (n: string): number | undefined => { const v = opt(n); return v == null ? undefined : Number(v); };
877
+ const text = await generateText(tm, prompt, {
878
+ maxTokens: Number(opt("max-tokens", "256")),
879
+ useChat: !flag("raw"),
880
+ sampler: { temperature: num("temperature"), topP: num("top-p"), topK: num("top-k"), seed: num("seed") },
881
+ });
882
+ process.stdout.write(text.endsWith("\n") ? text : text + "\n");
883
+ break;
884
+ }
885
+
736
886
  case "bench":
737
887
  case "benchmark": {
738
888
  const { banner, step, box, style } = await import("./tui");
@@ -1070,6 +1220,489 @@ switch (cmd) {
1070
1220
  break;
1071
1221
  }
1072
1222
 
1223
+ case "train": {
1224
+ const { banner, step, box, style } = await import("./tui");
1225
+ const query = positional(0) ?? opt("query");
1226
+ const dataDir = opt("data");
1227
+ if (!dataDir) {
1228
+ console.error("usage: mlx-bun train <model> --data <dir> (see: mlx-bun help train)");
1229
+ process.exit(1);
1230
+ }
1231
+ if (!(await Bun.file(`${dataDir}/train.jsonl`).exists())) {
1232
+ console.error(`no train.jsonl in ${dataDir}`);
1233
+ process.exit(1);
1234
+ }
1235
+ const method = opt("method", "orpo")!;
1236
+ if (method !== "sft" && method !== "dpo" && method !== "orpo") {
1237
+ console.error(`--method must be sft | dpo | orpo (got "${method}")`);
1238
+ process.exit(1);
1239
+ }
1240
+ const isOrpo = method === "orpo";
1241
+ // Validated numeric flag: default if absent, hard-exit on a non-number.
1242
+ const numFlag = (name: string, dflt: number): number => {
1243
+ const v = opt(name);
1244
+ if (v == null) return dflt;
1245
+ const n = Number(v);
1246
+ if (!Number.isFinite(n)) { console.error(`--${name} expects a number (got "${v}")`); process.exit(1); }
1247
+ return n;
1248
+ };
1249
+
1250
+ banner(pkg.version);
1251
+ const sNative = step("native runtime");
1252
+ await ensureNative(sNative);
1253
+ sNative.done("native runtime ready");
1254
+
1255
+ const { m, picked } = await resolveModelAuto(query);
1256
+
1257
+ // Detect Gemma/e4b and set its required training env flags BEFORE the
1258
+ // trainer is imported (perfKernelEnabled / fused-gelu read them lazily at
1259
+ // forward time, so setting them here — like the launcher — takes effect).
1260
+ const isGemma = (await Bun.file(`${m.path}/config.json`).text()).toLowerCase().includes("gemma");
1261
+ if (isGemma) {
1262
+ process.env.MLX_BUN_PERF_KERNEL ??= "0";
1263
+ process.env.MLX_BUN_FUSED_GELU ??= "0";
1264
+ }
1265
+
1266
+ const modelTag = isGemma ? "e4b" : "cpm5";
1267
+ const adapter = opt("adapter") ?? `${process.env.HOME}/.cache/mlx-bun/mlx-bun-finetunes/${method}-${modelTag}`;
1268
+ const iters = numFlag("iters", 100);
1269
+ const seq = numFlag("seq", isGemma ? 8192 : 4096);
1270
+ const seg = flag("no-segment") ? 0 : numFlag("seg", isOrpo ? 2 : 0);
1271
+ const saveEvery = numFlag("save-every", 0);
1272
+ const flashOn = !flag("no-flash");
1273
+ const prefixOn = !flag("no-prefix");
1274
+ const resume = opt("resume") ?? "";
1275
+
1276
+ // Build the snake_case submit record — the exact shape the server hands to
1277
+ // the finetune job runner (src/train/job.ts parseConfig).
1278
+ const cfg: Record<string, unknown> = {
1279
+ model_dir: m.path,
1280
+ data_dir: dataDir,
1281
+ adapter_path: adapter,
1282
+ method,
1283
+ rank: numFlag("rank", isOrpo ? 16 : 8),
1284
+ scale: numFlag("scale", isOrpo ? 2.0 : 1.0),
1285
+ rank_scaling: "by_bits",
1286
+ num_layers: -1,
1287
+ iters,
1288
+ learning_rate: numFlag("lr", isOrpo ? 1e-5 : method === "dpo" ? 5e-5 : 2e-4),
1289
+ max_seq_length: seq,
1290
+ batch_size: numFlag("batch", 1),
1291
+ grad_accumulation_steps: numFlag("grad-accum", 1),
1292
+ seed: numFlag("seed", 0),
1293
+ steps_per_report: 1,
1294
+ steps_per_eval: saveEvery > 0 ? saveEvery : 1_000_000,
1295
+ save_checkpoints: saveEvery > 0,
1296
+ segment_size: seg,
1297
+ grad_clip_norm: numFlag("grad-clip", 1.0),
1298
+ val_max_examples: numFlag("val-size", 256),
1299
+ warm_start_adapter: resume,
1300
+ ...(isOrpo ? {
1301
+ orpo_lambda: numFlag("lambda", 0.1),
1302
+ orpo_lr_schedule: "cosine",
1303
+ orpo_warmup_iters: Math.min(10, Math.floor(iters / 10)),
1304
+ orpo_chunk_size: 512,
1305
+ orpo_flash_ce: flashOn,
1306
+ orpo_fused_ce: !flashOn,
1307
+ orpo_prefix_shared: prefixOn,
1308
+ } : {}),
1309
+ };
1310
+
1311
+ // Pre-flight: dataset counts + detected format (bail before loading the model).
1312
+ const { inspectDataset } = await import("./train/job");
1313
+ const ds = await inspectDataset(dataDir);
1314
+ if (!ds.ok) { console.error(`dataset: ${ds.error}`); process.exit(1); }
1315
+
1316
+ const planLines = [
1317
+ `${style.green("●")} ${style.bold(`train ${method}`)} ${style.dim(`· ${m.repoId}${picked ? " (auto-picked)" : ""}${isGemma ? " · e4b env set" : ""}`)}`,
1318
+ "",
1319
+ `data ${style.bold(`${ds.n_train} train`)}${ds.n_valid ? ` · ${ds.n_valid} valid` : ""} ${style.dim(`· format ${ds.format}`)}`,
1320
+ `loop ${style.dim(`iters ${iters} · lr ${cfg.learning_rate} · rank ${cfg.rank} · scale ${cfg.scale} · seq ${seq} · batch ${cfg.batch_size}`)}`,
1321
+ ];
1322
+ if (isOrpo)
1323
+ planLines.push(
1324
+ `head ${style.dim(flashOn ? "flash-CCE Metal ([M,vocab]-free)" : "MLX fused linear-CE")}`,
1325
+ `stack ${style.dim(`prefix-share ${prefixOn ? "on" : "off"} · segmented ${seg > 0 ? `${seg}/seg` : "off"} · λ ${cfg.orpo_lambda}`)}`,
1326
+ `stability ${style.dim(`grad-clip ${cfg.grad_clip_norm || "off"} · val-size ${cfg.val_max_examples}${(cfg.grad_accumulation_steps as number) > 1 ? ` · grad-accum ${cfg.grad_accumulation_steps} (eff batch ${(cfg.batch_size as number) * (cfg.grad_accumulation_steps as number)})` : ""}`)}`,
1327
+ );
1328
+ else planLines.push(`stack ${style.dim(`segmented ${seg > 0 ? `${seg}/seg` : "off"}`)}`);
1329
+ if (resume) planLines.push(`warm-start ${style.dim(`from ${resume} (weights only)`)}`);
1330
+ if (saveEvery > 0) planLines.push(`checkpoint ${style.dim(`every ${saveEvery} steps`)}`);
1331
+ planLines.push("", `adapter ${style.dim(adapter)}`);
1332
+ console.log();
1333
+ box(planLines);
1334
+ console.log();
1335
+ console.log(` ${style.dim("watch live (other tab):")} ${style.accent(`mlx-bun train-watch ${adapter}`)}`);
1336
+ console.log();
1337
+
1338
+ if (flag("dry-run")) { console.log(` ${style.dim("dry run — not training.")}`); break; }
1339
+
1340
+ // Run the finetune job runner IN-PROCESS (foreground), streaming metrics to
1341
+ // the terminal — the same runner the server drives via submitSubprocess.
1342
+ // Run long jobs detached from your own shell: nohup mlx-bun train … &
1343
+ const { finetuneRunner } = await import("./train/job");
1344
+ const { peakMemory, resetPeakMemory } = await import("./mlx/ffi");
1345
+ const losses: number[] = [];
1346
+ const stepMs: number[] = [];
1347
+ let lastStepT = Date.now();
1348
+ resetPeakMemory();
1349
+ const emit = (e: JobEvent) => {
1350
+ if (e.type === "stage" && e.message) console.log(` ${style.dim("·")} ${e.message}`);
1351
+ else if (e.type === "metric" && e.kind === "train") {
1352
+ const now = Date.now(); stepMs.push(now - lastStepT); lastStepT = now;
1353
+ losses.push(e.loss);
1354
+ const n = losses.length;
1355
+ if (n <= 3 || n % 10 === 0)
1356
+ console.log(` step ${n}/${iters}: loss ${style.bold(e.loss.toFixed(4))} ${style.dim(`(${(stepMs[stepMs.length - 1]! / 1000).toFixed(1)}s/step · peak ${gb(peakMemory())})`)}`);
1357
+ }
1358
+ };
1359
+ try {
1360
+ await finetuneRunner(emit, cfg);
1361
+ } catch (err) {
1362
+ console.error(`\n training failed: ${err instanceof Error ? err.message : String(err)}`);
1363
+ process.exit(1);
1364
+ }
1365
+ const finite = losses.every((l) => Number.isFinite(l));
1366
+ const sorted = stepMs.slice(1).sort((a, b) => a - b);
1367
+ const med = sorted[Math.floor(sorted.length / 2)] ?? stepMs[0] ?? 0;
1368
+ console.log();
1369
+ box([
1370
+ `${style.green("●")} ${style.bold("training complete")} ${style.dim(`· ${losses.length} steps`)}`,
1371
+ "",
1372
+ `loss ${style.bold(`${losses[0]?.toFixed(4) ?? "—"} → ${losses[losses.length - 1]?.toFixed(4) ?? "—"}`)}${finite ? "" : " (NON-FINITE!)"}`,
1373
+ `speed ${style.dim(`${(med / 1000).toFixed(1)}s/step median · peak ${gb(peakMemory())}`)}`,
1374
+ "",
1375
+ `adapter ${style.bold(adapter)}`,
1376
+ `serve it ${style.accent(`mlx-bun serve ${m.repoId} --adapter ${adapter}`)}`,
1377
+ ]);
1378
+ break;
1379
+ }
1380
+
1381
+ case "train-watch": {
1382
+ // Live mactop-style dashboard for a training run: tails <adapter>/metrics.jsonl.
1383
+ const dir = positional(0) ?? opt("adapter")
1384
+ ?? `${process.env.HOME}/.cache/mlx-bun/mlx-bun-finetunes/orpo-cpm5`;
1385
+ const { runWatch } = await import("./train/watch");
1386
+ await runWatch(dir);
1387
+ break;
1388
+ }
1389
+
1390
+ case "memory": {
1391
+ const { style, box, step, banner } = await import("./tui");
1392
+ const {
1393
+ vaultRoot, vaultStatus, setupVault, importArticlesFrom, commitVault,
1394
+ listMemoryDocuments, readArticle, resolveArticlePath, searchArticles,
1395
+ parseToc, extractSection, getArticleLinks,
1396
+ } = await import("./memory/vault");
1397
+ const { installSchedule, removeSchedule, scheduleStatus, parseAt } = await import("./memory/schedule");
1398
+ const root = vaultRoot();
1399
+ const fmtAt = (a: { hour: number; minute: number }) => `${String(a.hour).padStart(2, "0")}:${String(a.minute).padStart(2, "0")}`;
1400
+ const openMemoryTarget = async (targetPath = root): Promise<"obsidian" | "fallback"> => {
1401
+ // For articles, Obsidian's URL handler opens the exact file inside the
1402
+ // vault. For the vault root, `open -a Obsidian <folder>` opens the folder
1403
+ // as a vault. Both fall back to macOS `open` if Obsidian is unavailable.
1404
+ if (targetPath !== root) {
1405
+ const uri = `obsidian://open?path=${encodeURIComponent(targetPath)}`;
1406
+ const byUri = Bun.spawn(["open", uri], { stdout: "ignore", stderr: "ignore" });
1407
+ if ((await byUri.exited) === 0) return "obsidian";
1408
+ } else {
1409
+ const byApp = Bun.spawn(["open", "-a", "Obsidian", root], { stdout: "ignore", stderr: "ignore" });
1410
+ if ((await byApp.exited) === 0) return "obsidian";
1411
+ }
1412
+ const fallback = Bun.spawn(["open", targetPath], { stdout: "ignore", stderr: "ignore" });
1413
+ await fallback.exited;
1414
+ return "fallback";
1415
+ };
1416
+
1417
+ // Words after `memory`, minus flags: [subcommand, ...query].
1418
+ const words = argv.slice(1).filter((a) => !a.startsWith("-"));
1419
+ const sub = (words[0] ?? "status").toLowerCase();
1420
+ const rest = words.slice(1).join(" ").trim();
1421
+
1422
+ // TTY-only line prompt; in a non-interactive shell returns the default so
1423
+ // the wizard degrades to a non-destructive no-op rather than hanging.
1424
+ const ask = async (q: string, dflt = ""): Promise<string> => {
1425
+ if (!process.stdin.isTTY) return dflt;
1426
+ const { createInterface } = await import("node:readline/promises");
1427
+ const rl = createInterface({ input: process.stdin, output: process.stdout });
1428
+ try {
1429
+ return ((await rl.question(q)).trim() || dflt);
1430
+ } finally {
1431
+ rl.close();
1432
+ }
1433
+ };
1434
+ const confirmYN = async (q: string, defaultYes: boolean): Promise<boolean> => {
1435
+ const a = (await ask(`${q} ${defaultYes ? "[Y/n]" : "[y/N]"} `)).toLowerCase();
1436
+ if (!a) return defaultYes;
1437
+ return a === "y" || a === "yes";
1438
+ };
1439
+
1440
+ if (sub === "init" || sub === "setup") {
1441
+ banner(pkg.version);
1442
+ box([
1443
+ `${style.accent("●")} ${style.bold("Personal memory for your local AI")}`,
1444
+ "",
1445
+ "A wiki of Markdown articles your local assistant reads to remember",
1446
+ "your projects, people, and history across sessions. It's yours —",
1447
+ `git-tracked at ${style.dim(root)}, editable in any tool, and it`,
1448
+ "never leaves this machine.",
1449
+ ]);
1450
+ const sSetup = step("creating your wiki");
1451
+ const res = await setupVault(root);
1452
+ sSetup.done(res.alreadySetUp ? `already set up ${style.dim(`· ${root}`)}` : `wiki ready ${style.dim(`· ${root}`)}`);
1453
+
1454
+ // Offer to seed from an existing vault (e.g. a lucien ~/Dreaming) so the
1455
+ // assistant has real content to read on day one, before synthesis exists.
1456
+ const { homedir } = await import("node:os");
1457
+ const { join } = await import("node:path");
1458
+ const { readdir } = await import("node:fs/promises");
1459
+ const seedRoot = await ask(
1460
+ `\n Seed from an existing wiki? Enter its path, or blank to skip ${style.dim(`[${join(homedir(), "Dreaming")}]`)}: `,
1461
+ "",
1462
+ );
1463
+ if (seedRoot) {
1464
+ const srcArticles = join(seedRoot.replace(/^~(?=$|[/\\])/, homedir()), "articles");
1465
+ let count = 0;
1466
+ try { count = (await readdir(srcArticles)).filter((n) => n.endsWith(".md")).length; } catch { count = 0; }
1467
+ if (count === 0) {
1468
+ console.log(style.dim(` no articles found under ${srcArticles} — skipping import`));
1469
+ } else if (await confirmYN(` Import ${count} article(s) from ${srcArticles}?`, true)) {
1470
+ const sImp = step(`importing ${count} article(s)`);
1471
+ const imported = await importArticlesFrom(srcArticles, root);
1472
+ await commitVault(root, `Import ${imported.length} article(s) from ${seedRoot}`);
1473
+ sImp.done(`imported ${imported.length} article(s)`);
1474
+ }
1475
+ }
1476
+
1477
+ // Offer to install the nightly synthesis job (a launchd agent). The job
1478
+ // itself is a no-op today (synthesis is stubbed), but installing the
1479
+ // schedule now means it starts producing articles the moment M1 lands —
1480
+ // no reconfiguration. Persistent action → TTY-only, defaults to no.
1481
+ let scheduled: string | null = null;
1482
+ console.log(style.dim(
1483
+ "\n Synthesis (turning your conversations into articles) isn't implemented yet,\n" +
1484
+ " but you can install the nightly job now so it runs automatically once it is.",
1485
+ ));
1486
+ if (await confirmYN(" Install the nightly synthesis job (runs in the background)?", false)) {
1487
+ const atRaw = await ask(` At what time? 24h HH:MM ${style.dim("[03:00]")}: `, "03:00");
1488
+ const sSched = step(`installing nightly job at ${fmtAt(parseAt(atRaw))}`);
1489
+ const r = await installSchedule({ at: atRaw });
1490
+ scheduled = fmtAt(r.at);
1491
+ sSched.done(
1492
+ r.loaded
1493
+ ? `nightly job installed ${style.dim(`· ${scheduled} · ${r.plistPath}`)}`
1494
+ : `plist written ${style.dim(`· ${r.plistPath} · launchctl load failed (load it manually)`)}`,
1495
+ );
1496
+ }
1497
+
1498
+ const st = await vaultStatus(root);
1499
+ console.log();
1500
+ box([
1501
+ `${style.green("●")} ${style.bold("memory is set up")} ${style.dim(`· ${st.articleCount} article(s) · ${st.referenceCount} reference doc(s)`)}`,
1502
+ "",
1503
+ `chat now ${style.accent("mlx-bun pi")} ${style.dim("— the assistant now reads your memory automatically")}`,
1504
+ `inspect ${style.accent("mlx-bun memory status")} ${style.dim("·")} ${style.accent("mlx-bun memory search <q>")}`,
1505
+ `browse ${style.dim(`open ${root} in Obsidian, or any editor`)}`,
1506
+ scheduled
1507
+ ? `nightly ${style.green("scheduled")} ${style.dim(`· ${scheduled} daily · mlx-bun memory unschedule to undo`)}`
1508
+ : `nightly ${style.dim("not scheduled · mlx-bun memory schedule to set it up")}`,
1509
+ "",
1510
+ style.dim("Synthesis (conversations → articles) is stubbed for now; the scheduled"),
1511
+ style.dim("job no-ops until it lands. Add or import articles to read meanwhile."),
1512
+ ]);
1513
+ break;
1514
+ }
1515
+
1516
+ if (sub === "status") {
1517
+ const st = await vaultStatus(root);
1518
+ if (!st.exists) {
1519
+ console.log(`no memory wiki yet at ${style.dim(root)}`);
1520
+ console.log(`set one up: ${style.accent("mlx-bun memory init")}`);
1521
+ break;
1522
+ }
1523
+ const sched = await scheduleStatus();
1524
+ console.log();
1525
+ box([
1526
+ `${style.bold("memory")} ${style.dim(`· ${root}`)}`,
1527
+ "",
1528
+ `articles ${style.bold(String(st.articleCount))}`,
1529
+ `reference ${style.bold(String(st.referenceCount))} ${style.dim("read-only docs")}`,
1530
+ `git ${st.isGitRepo ? style.green("tracked") : style.dim("not a git repo")}`,
1531
+ `synthesis ${style.dim("stubbed (M1) — not implemented yet")}`,
1532
+ `last run ${style.dim("not available yet")}`,
1533
+ `nightly ${
1534
+ sched.installed
1535
+ ? `${sched.loaded ? style.green("scheduled") : style.accent("installed (not loaded)")} ${style.dim(`· ${sched.plistPath}`)}`
1536
+ : style.dim("not scheduled · mlx-bun memory schedule")
1537
+ }`,
1538
+ st.recentArticles.length
1539
+ ? `recent ${st.recentArticles.slice(0, 5).map((r) => r.article).join(", ")}`
1540
+ : `recent ${style.dim("none")}`,
1541
+ ]);
1542
+ break;
1543
+ }
1544
+
1545
+ if (sub === "open" || sub === "browse") {
1546
+ const st = await vaultStatus(root);
1547
+ if (!st.exists) {
1548
+ console.log(`no memory wiki yet at ${style.dim(root)}`);
1549
+ console.log(`set one up: ${style.accent("mlx-bun memory init")}`);
1550
+ break;
1551
+ }
1552
+ let targetPath = root;
1553
+ let label = "memory";
1554
+ if (rest) {
1555
+ try {
1556
+ targetPath = await resolveArticlePath(root, rest);
1557
+ label = rest.endsWith(".md") ? rest.slice(0, -3) : rest;
1558
+ } catch (err) {
1559
+ console.error(err instanceof Error ? err.message : String(err));
1560
+ console.error(`try: ${style.accent(`mlx-bun memory search ${rest}`)}`);
1561
+ process.exit(1);
1562
+ }
1563
+ }
1564
+ const where = await openMemoryTarget(targetPath);
1565
+ console.log(
1566
+ where === "obsidian"
1567
+ ? `opened ${label} in Obsidian: ${targetPath}`
1568
+ : `opened ${label}: ${targetPath}`,
1569
+ );
1570
+ break;
1571
+ }
1572
+
1573
+ if (sub === "synthesize") {
1574
+ const { runSynthesis } = await import("./memory/pipeline");
1575
+ const dryRun = flag("dry-run");
1576
+ console.log();
1577
+ const summary = await runSynthesis(
1578
+ { since: opt("since") ?? undefined, model: opt("model") ?? undefined, dryRun },
1579
+ (e) => {
1580
+ if (e.type === "stage") console.log(` ${style.dim("·")} ${e.message}`);
1581
+ else console.log(` ${e.message}`);
1582
+ },
1583
+ );
1584
+ console.log(style.dim(`\n ${summary.note}`));
1585
+ break;
1586
+ }
1587
+
1588
+ if (sub === "schedule") {
1589
+ const atRaw = opt("at") ?? "03:00";
1590
+ const r = await installSchedule({ at: atRaw });
1591
+ console.log();
1592
+ box([
1593
+ r.loaded
1594
+ ? `${style.green("●")} ${style.bold("nightly synthesis scheduled")} ${style.dim(`· ${fmtAt(r.at)} daily`)}`
1595
+ : `${style.accent("●")} ${style.bold("plist written, but launchctl load failed")}`,
1596
+ "",
1597
+ `plist ${style.dim(r.plistPath)}`,
1598
+ `runs ${style.dim("mlx-bun memory synthesize (a no-op until synthesis lands)")}`,
1599
+ `undo ${style.accent("mlx-bun memory unschedule")}`,
1600
+ ]);
1601
+ break;
1602
+ }
1603
+
1604
+ if (sub === "unschedule") {
1605
+ const removed = await removeSchedule();
1606
+ console.log(removed ? "nightly synthesis job removed" : "no nightly job was installed");
1607
+ break;
1608
+ }
1609
+
1610
+ if (sub === "list") {
1611
+ const docs = await listMemoryDocuments(root);
1612
+ if (docs.length === 0) { console.log("no articles or reference docs yet"); break; }
1613
+ for (const s of docs) console.log(` ${s}`);
1614
+ const refs = docs.filter((s) => s.startsWith("Reference/")).length;
1615
+ console.log(style.dim(`\n ${docs.length - refs} article(s), ${refs} read-only reference doc(s)`));
1616
+ break;
1617
+ }
1618
+
1619
+ if (sub === "search") {
1620
+ if (!rest) { console.error("usage: mlx-bun memory search <query>"); process.exit(1); }
1621
+ const { summaries, hits } = await searchArticles(root, rest, { limit: 12 });
1622
+ if (summaries.length === 0) { console.log(`no articles match "${rest}"`); break; }
1623
+ console.log();
1624
+ for (const s of summaries.slice(0, 15)) {
1625
+ const terms = s.matched_terms?.length ? style.dim(` [${s.matched_terms.join(", ")}]`) : "";
1626
+ console.log(` ${style.bold(s.article)} ${style.dim(`· ${s.occurrences} hit(s)`)}${terms}`);
1627
+ }
1628
+ if (hits.length) {
1629
+ console.log(style.dim("\n sample lines:"));
1630
+ for (const h of hits.slice(0, 6)) {
1631
+ const where = h.anchor ? `${h.article}#${h.anchor}` : h.article;
1632
+ console.log(` ${style.dim(`${where}:${h.line}`)} ${h.excerpt}`);
1633
+ }
1634
+ }
1635
+ console.log(style.dim(`\n read one: mlx-bun memory read ${summaries[0]!.article}`));
1636
+ break;
1637
+ }
1638
+
1639
+ if (sub === "toc") {
1640
+ if (!rest) { console.error("usage: mlx-bun memory toc <article>"); process.exit(1); }
1641
+ try {
1642
+ const { content } = await readArticle(root, rest);
1643
+ const toc = parseToc(content);
1644
+ if (toc.length === 0) { console.log("no headings"); break; }
1645
+ for (const h of toc) console.log(`${" ".repeat(Math.max(0, h.depth - 1))}- ${h.title} ${style.dim(`#${h.anchor}`)}`);
1646
+ } catch (err) {
1647
+ console.error(err instanceof Error ? err.message : String(err));
1648
+ console.error(`try: ${style.accent(`mlx-bun memory search ${rest}`)}`);
1649
+ process.exit(1);
1650
+ }
1651
+ break;
1652
+ }
1653
+
1654
+ if (sub === "section") {
1655
+ const [article, anchorRaw] = words.slice(1);
1656
+ const anchor = anchorRaw?.replace(/^#/, "");
1657
+ if (!article || !anchor) { console.error("usage: mlx-bun memory section <article> <anchor>"); process.exit(1); }
1658
+ try {
1659
+ const { content } = await readArticle(root, article);
1660
+ const section = extractSection(content, anchor);
1661
+ if (!section) { console.error(`no section #${anchor} in ${article}`); console.error(`try: ${style.accent(`mlx-bun memory toc ${article}`)}`); process.exit(1); }
1662
+ console.log(section);
1663
+ } catch (err) {
1664
+ console.error(err instanceof Error ? err.message : String(err));
1665
+ console.error(`try: ${style.accent(`mlx-bun memory search ${article}`)}`);
1666
+ process.exit(1);
1667
+ }
1668
+ break;
1669
+ }
1670
+
1671
+ if (sub === "links") {
1672
+ if (!rest) { console.error("usage: mlx-bun memory links <article>"); process.exit(1); }
1673
+ try {
1674
+ const { outbound, inbound } = await getArticleLinks(root, rest);
1675
+ console.log(`${style.bold("outbound")}:`);
1676
+ console.log(outbound.length ? outbound.map((s) => ` ${s}`).join("\n") : " (none)");
1677
+ console.log(`\n${style.bold("inbound")}:`);
1678
+ console.log(inbound.length ? inbound.map((s) => ` ${s}`).join("\n") : " (none)");
1679
+ } catch (err) {
1680
+ console.error(err instanceof Error ? err.message : String(err));
1681
+ console.error(`try: ${style.accent(`mlx-bun memory search ${rest}`)}`);
1682
+ process.exit(1);
1683
+ }
1684
+ break;
1685
+ }
1686
+
1687
+ if (sub === "read") {
1688
+ if (!rest) { console.error("usage: mlx-bun memory read <article>"); process.exit(1); }
1689
+ try {
1690
+ const { content } = await readArticle(root, rest);
1691
+ console.log(content);
1692
+ } catch (err) {
1693
+ console.error(err instanceof Error ? err.message : String(err));
1694
+ console.error(`try: ${style.accent(`mlx-bun memory search ${rest}`)}`);
1695
+ process.exit(1);
1696
+ }
1697
+ break;
1698
+ }
1699
+
1700
+
1701
+ console.error(`unknown: mlx-bun memory ${sub}`);
1702
+ printHelp("memory");
1703
+ break;
1704
+ }
1705
+
1073
1706
  default:
1074
1707
  printHelp(cmd);
1075
1708
  }