mlx-bun 0.4.0 → 1.0.0

package/src/eval/tasks/hashhop.ts

@@ -1,139 +0,0 @@
-// HashHop — long-context multi-hop key/value retrieval. Port of
-// optiq/eval/hashhop.py. The generator (Magic AI's HashHop, codelion's
-// `hashhop` package) isn't on PyPI, so we reconstruct MultiHopEval.make_one
-// from the documented format: a chain h0=h1, h1=h2, …, h(N-1)='hN' (quoted
-// terminal) buried in many distractor assignments, serialized as a shuffled
-// dictionary. The model walks N lookups from h0 and reports hN.
-//
-// This is the long-context guard (≈12k-char dict → ~8–9k tokens): online-
-// softmax / mixed-4/8-bit-KV error compounds across hops, so it's where the
-// fused-attention kernel (opportunity A) has the most to prove.
-
-import { generateText, loadJsonl, type TaskModel } from "../runner";
-import { runtimeValue } from "../../runtime-config";
-
-const INSTR =
-  "You will be given a list of hash assignments and a starting hash. " +
-  "Walk the chain by repeated key->value lookup until you hit a terminal " +
-  "value (written with single quotes, e.g. `KEY = 'VALUE'`). Output ONLY " +
-  "the final 16-character hash inside the quotes, with no extra words or " +
-  "punctuation.";
-
-const ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";
-const HASH_LEN = 16;
-
-function mulberry32(seed: number): () => number {
-  let s = seed >>> 0;
-  return () => {
-    s = (s + 0x6d2b79f5) | 0;
-    let t = Math.imul(s ^ (s >>> 15), 1 | s);
-    t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t;
-    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
-  };
-}
-
-function randHash(rand: () => number): string {
-  let h = "";
-  for (let i = 0; i < HASH_LEN; i++) h += ALPHABET[Math.floor(rand() * ALPHABET.length)];
-  return h;
-}
-
-interface Sample { prompt: string; start: string; expected: string }
-
-/** One multi-hop instance: a length-`hops` chain padded with distractors to
- *  ~`nChars`, lines shuffled. start = h0, expected = h(hops). */
-function makeOne(rand: () => number, hops: number, nChars: number): Sample {
-  const chain: string[] = [];
-  for (let i = 0; i <= hops; i++) chain.push(randHash(rand));
-  const lines: string[] = [];
-  for (let i = 0; i < hops; i++)
-    lines.push(i === hops - 1 ? `${chain[i]} = '${chain[i + 1]}'` : `${chain[i]} = ${chain[i + 1]}`);
-
-  let chars = lines.reduce((a, l) => a + l.length + 1, 0);
-  while (chars < nChars) {
-    const a = randHash(rand);
-    const b = randHash(rand);
-    const line = rand() < 0.5 ? `${a} = '${b}'` : `${a} = ${b}`;
-    lines.push(line);
-    chars += line.length + 1;
-  }
-  for (let i = lines.length - 1; i > 0; i--) {
-    const j = Math.floor(rand() * (i + 1));
-    [lines[i], lines[j]] = [lines[j]!, lines[i]!];
-  }
-  return { prompt: lines.join("\n"), start: chain[0]!, expected: chain[hops]! };
-}
-
-function buildPrompt(s: Sample): string {
-  return `${INSTR}\n\nDictionary:\n${s.prompt}\n\nStarting hash: ${s.start}\nFinal hash:`;
-}
-
-/** First 16-char alphabetic run in the output (port of _extract_pred). */
-function extractPred(out: string): string {
-  const m = out.match(/[A-Za-z]{16}/);
-  return m ? m[0] : out.trim();
-}
-
-export interface HashhopResult {
-  nTotal: number;
-  nCorrect: number;
-  accuracy: number; // overall, 0..1
-  byHops: Record<number, number>;
-}
-
-interface HashhopFrozen { raw_problem: string; query: string; expected: string }
-
-export async function evaluateHashhop(
-  tm: TaskModel,
-  opts: { nPerHop?: number; hopsList?: number[]; chars?: number; maxTokens?: number; seed?: number; frozen?: boolean } = {},
-): Promise<HashhopResult> {
-  const nPerHop = opts.nPerHop ?? 25;
-  const hopsList = opts.hopsList ?? [1, 2, 3, 4];
-  const maxTokens = opts.maxTokens ?? 32;
-
-  // Optiq-parity mode (DEFAULT): use optiq's EXACT generated problems (his
-  // MultiHopEval samples — seed 42, 25/hop × hops 1-4 = 100, emitted in hop order),
-  // fed through OUR chat template + extractor. Our generator below is a
-  // RECONSTRUCTION producing DIFFERENT problems, so it cannot reproduce optiq's
-  // number — the frozen problems are the only way to match. MLX_BUN_HASHHOP_FROZEN=0
-  // uses our generator.
-  const useFrozen = opts.frozen ?? (runtimeValue("MLX_BUN_HASHHOP_FROZEN") !== "0");
-  if (useFrozen) {
-    const rows = loadJsonl<HashhopFrozen>("hashhop_optiq_frozen");
-    const hopCorrect: Record<number, number> = {};
-    let total = 0;
-    let correct = 0;
-    for (let k = 0; k < rows.length; k++) {
-      const item = rows[k]!;
-      const hops = hopsList[Math.floor(k / nPerHop)] ?? 0; // optiq emits in hop order
-      const out = await generateText(tm, item.raw_problem, { maxTokens, useChat: true });
-      if (extractPred(out) === item.expected) { correct++; hopCorrect[hops] = (hopCorrect[hops] ?? 0) + 1; }
-      total++;
-      process.stderr.write(`\r  hashhop ${k + 1}/${rows.length}  acc=${(correct / (k + 1) * 100).toFixed(0)}%`);
-    }
-    process.stderr.write("\n");
-    const byHops: Record<number, number> = {};
-    for (const h of hopsList) byHops[h] = (hopCorrect[h] ?? 0) / nPerHop;
-    return { nTotal: total, nCorrect: correct, accuracy: total ? correct / total : 0, byHops };
-  }
-
-  // Fallback: our own reconstructed generator (produces DIFFERENT problems).
-  const chars = opts.chars ?? 12000;
-  const rand = mulberry32(opts.seed ?? 42);
-  const byHops: Record<number, number> = {};
-  let total = 0;
-  let correct = 0;
-  for (const hops of hopsList) {
-    let hc = 0;
-    for (let k = 0; k < nPerHop; k++) {
-      const s = makeOne(rand, hops, chars);
-      const out = await generateText(tm, buildPrompt(s), { maxTokens, useChat: true });
-      if (extractPred(out) === s.expected) { hc++; correct++; }
-      total++;
-      process.stderr.write(`\r  hashhop hops=${hops} ${k + 1}/${nPerHop}  acc=${(hc / (k + 1) * 100).toFixed(0)}%`);
-    }
-    byHops[hops] = nPerHop ? hc / nPerHop : 0;
-    process.stderr.write("\n");
-  }
-  return { nTotal: total, nCorrect: correct, accuracy: total ? correct / total : 0, byHops };
-}

package/src/eval/tasks/humaneval.ts

@@ -1,275 +0,0 @@
-// HumanEval — code-generation pass@1. Port of optiq/eval/humaneval.py:
-// raw function-completion (or instruct code-block when a chat template
-// exists), greedy decode, stop-sequence truncation, then execute the
-// assembled program (prompt + completion + test + check(entry)) in a
-// macOS sandbox-exec jail (falls back to a plain subprocess) against the
-// oracle venv python. A problem passes iff the program exits 0 within the
-// timeout.
-//
-// SECURITY: generated code is UNTRUSTED model output. It is NEVER eval'd
-// in-process — each candidate runs as a separate, sandboxed python
-// subprocess with a hard wall-clock timeout and a memory cap, in a fresh
-// temp dir that is removed afterwards.
-
-import { spawnSync } from "node:child_process";
-import { mkdtempSync, rmSync, writeFileSync } from "node:fs";
-import { tmpdir } from "node:os";
-import { join } from "node:path";
-import { generateText, loadJsonl, sampleIndices, type TaskModel } from "../runner";
-import { runtimeValue } from "../../runtime-config";
-
-// Oracle venv python — the same interpreter the Python reference uses.
-// Machine-specific: overridable via MLX_BUN_ORACLE_PYTHON (the committed default is
-// the reference box; other laptops set the env, never re-commit the path).
-const ORACLE_PYTHON = runtimeValue("MLX_BUN_ORACLE_PYTHON") ?? "/Users/joshrossi/Code/mlx-lm/.venv/bin/python";
-
-interface HumanevalRow {
-  task_id: string;
-  prompt: string;
-  canonical_solution: string;
-  test: string;
-  entry_point: string;
-}
-
-export interface HumanevalResult {
-  nPass: number;
-  nTotal: number;
-  accuracy: number; // pass@1, 0..1
-}
-
-// Stop sequences — mirror humaneval.py's _STOP_SEQS exactly.
-const STOP_SEQS = ["\nclass ", "\ndef ", "\n#", "\nif __name__", "\nprint(", "\n```"];
-
-/** Last fenced ```python|py? ... ``` block, or null (port of _extract_code_block). */
-function extractCodeBlock(text: string): string | null {
-  const re = /```(?:python|py)?\n([\s\S]*?)\n```/g;
-  let last: string | null = null;
-  for (const m of text.matchAll(re)) last = m[1]!;
-  return last;
-}
-
-/** Escape a string for use as a literal inside a RegExp (re.escape). */
-function reEscape(s: string): string {
-  return s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
-}
-
-/**
- * Strip the `def entry_point(...):` header + docstring, keep just the body
- * (port of _strip_function_def). Falls through if no such def is found.
- */
-function stripFunctionDef(body: string, entryPoint: string): string {
-  const pat = new RegExp(`def\\s+${reEscape(entryPoint)}\\s*\\([^)]*\\)[^:]*:\\s*\\n`);
-  const m = body.match(pat);
-  if (!m || m.index === undefined) return body;
-  const after = body.slice(m.index + m[0].length);
-  // Skip leading blank lines + a possible docstring.
-  let rest = after.replace(/^\n+/, "");
-  const doc = rest.match(/^\s*(["']{3})([\s\S]*?)\1\s*\n/);
-  if (doc) rest = rest.slice(doc[0].length);
-  return rest;
-}
-
-/**
- * Reduce the model response to a function body suitable for concatenating
- * onto the HumanEval prompt (port of _truncate_completion).
- */
-export function truncateCompletion(text: string, entryPoint = ""): string {
-  let out = text;
-  if (out.includes("</think>")) out = out.split("</think>").slice(1).join("</think>");
-  if (out.includes("<channel|>")) out = out.split("<channel|>").slice(1).join("<channel|>");
-
-  // Prefer the last fenced code block (conversational models).
-  const block = extractCodeBlock(out);
-  if (block !== null) {
-    return entryPoint ? stripFunctionDef(block, entryPoint) : block;
-  }
-
-  // Legacy / completion-mode path: strip a leading fence, then cut at the
-  // first stop sequence.
-  out = out.replace(/^\s+/, "");
-  if (out.startsWith("```python\n")) out = out.slice("```python\n".length);
-  else if (out.startsWith("```\n")) out = out.slice("```\n".length);
-
-  let earliest = out.length;
-  for (const stop of STOP_SEQS) {
-    const i = out.indexOf(stop);
-    if (i !== -1 && i < earliest) earliest = i;
-  }
-  return out.slice(0, earliest);
-}
-
-/** Concatenate prompt + completion + test driver (port of _build_program). */
-function buildProgram(prompt: string, completion: string, test: string, entryPoint: string): string {
-  return prompt + completion + "\n\n" + test + "\n\n" + `check(${entryPoint})\n`;
-}
-
-// sandbox-exec SBPL profile (port of _SBPL_PROFILE): deny network, deny
-// most filesystem writes except the temp dir; reads stay open so python
-// can find its stdlib.
-const SBPL_PROFILE = (tmp: string) => `(version 1)
-(deny default)
-(allow process-fork)
-(allow process-exec)
-(allow file-read*)
-(allow file-write* (subpath "${tmp}"))
-(allow file-write-data
-       (literal "/dev/null")
-       (literal "/dev/dtracehelper"))
-(allow sysctl-read)
-(allow mach-lookup)
-(allow ipc-posix-shm)
-`;
-
-// rlimit + no-socket preamble prepended to the script (port of the
-// sandbox-exec preamble). Caps address space / data segment and disables
-// raw socket creation so a buggy program can't exhaust RAM or hit the net.
-const preamble = (memoryMb: number) => `import resource
-try:
-    resource.setrlimit(resource.RLIMIT_AS,   (${memoryMb} * 1024 * 1024,) * 2)
-except (ValueError, OSError):
-    pass
-try:
-    resource.setrlimit(resource.RLIMIT_DATA, (${memoryMb} * 1024 * 1024,) * 2)
-except (ValueError, OSError):
-    pass
-import socket as _s
-_orig_socket = _s.socket
-def _no_socket(*a, **k):
-    raise OSError("network access disabled in sandbox")
-_s.socket = _no_socket
-`;
-
-const HAS_SANDBOX_EXEC = (() => {
-  try {
-    return spawnSync("which", ["sandbox-exec"]).status === 0;
-  } catch {
-    return false;
-  }
-})();
-
-interface RunResult { ok: boolean; timedOut: boolean; stderr?: string }
-
-/**
- * Run `program` as a python script in a fresh temp dir, sandboxed if
- * sandbox-exec is available (else a plain rlimit-capped subprocess), with
- * a hard wall-clock timeout. Returns whether it exited 0 (pass). The code
- * is NEVER eval'd in-process — always a separate subprocess.
- */
-function runProgram(program: string, timeoutSec: number, memoryMb: number): RunResult {
-  const sbRoot = mkdtempSync(join(tmpdir(), "mlxbun_he_"));
-  try {
-    const scriptPath = join(sbRoot, "script.py");
-    writeFileSync(scriptPath, preamble(memoryMb) + "\n" + program);
-
-    const env = {
-      PATH: "/usr/bin:/bin",
-      PYTHONIOENCODING: "utf-8",
-      HOME: sbRoot,
-      TMPDIR: sbRoot,
-      MPLBACKEND: "Agg",
-      MPLCONFIGDIR: sbRoot,
-      XDG_CACHE_HOME: sbRoot,
-      XDG_CONFIG_HOME: sbRoot,
-    };
-
-    let argv: string[];
-    if (HAS_SANDBOX_EXEC) {
-      const profilePath = join(sbRoot, "policy.sb");
-      writeFileSync(profilePath, SBPL_PROFILE(sbRoot));
-      argv = ["sandbox-exec", "-f", profilePath, ORACLE_PYTHON, "-I", scriptPath];
-    } else {
-      argv = [ORACLE_PYTHON, "-I", scriptPath];
-    }
-
-    const res = spawnSync(argv[0]!, argv.slice(1), {
-      cwd: sbRoot,
-      env,
-      timeout: timeoutSec * 1000, // hard wall-clock kill
-      killSignal: "SIGKILL",
-      stdio: ["ignore", "ignore", runtimeValue("MLX_BUN_EVAL_DEBUG") === "1" ? "pipe" : "ignore"],
-      encoding: "utf8",
-    });
-
-    // spawnSync sets signal === "SIGKILL" (and status === null) when the
-    // timeout fires; ETIMEDOUT shows up in res.error on some platforms.
-    const timedOut =
-      res.signal === "SIGKILL" ||
-      (res.error as NodeJS.ErrnoException | undefined)?.code === "ETIMEDOUT";
-    const ok = !timedOut && res.status === 0;
-    return { ok, timedOut, stderr: typeof res.stderr === "string" ? res.stderr : undefined };
-  } finally {
-    rmSync(sbRoot, { recursive: true, force: true });
-  }
-}
-
-export async function evaluateHumaneval(
-  tm: TaskModel,
-  opts: {
-    nSamples?: number;
-    maxTokens?: number;
-    seed?: number;
-    timeoutSec?: number;
-    memoryMb?: number;
-    frozen?: boolean;
-  } = {},
-): Promise<HumanevalResult> {
-  const maxTokens = opts.maxTokens ?? 512;
-  const timeoutSec = opts.timeoutSec ?? 10;
-  const memoryMb = opts.memoryMb ?? 512;
-
-  // Optiq-parity mode (DEFAULT): optiq's EXACT 164-problem set (he runs the full
-  // HumanEval test split), scored through OUR pipeline. MLX_BUN_HUMANEVAL_FROZEN=0
-  // reverts to our own copy + optional sampling.
-  const useFrozen = opts.frozen ?? (runtimeValue("MLX_BUN_HUMANEVAL_FROZEN") !== "0");
-  const rows = loadJsonl<HumanevalRow>(useFrozen ? "humaneval_optiq_frozen" : "humaneval");
-  const idx = useFrozen
-    ? Array.from({ length: rows.length }, (_, i) => i)
-    : sampleIndices(rows.length, opts.nSamples ?? rows.length, opts.seed ?? 42);
-
-  // When a chat template exists, ask the model (instruct-style) for the
-  // complete function inside a ```python block — humaneval.py's use_chat
-  // path. Otherwise feed the raw signature for completion.
-  const useChat = tm.template !== null;
-
-  let nPass = 0;
-  for (let k = 0; k < idx.length; k++) {
-    const ex = rows[idx[k]!]!;
-
-    let body: string;
-    if (useChat) {
-      body =
-        "Complete the following Python function. Output the " +
-        "complete function (header, docstring, body) inside a " +
-        "single ```python ... ``` code block. No commentary " +
-        "outside the code block.\n\n" +
-        "```python\n" + ex.prompt + "```";
-    } else {
-      body = ex.prompt;
-    }
-
-    const raw = await generateText(tm, body, { maxTokens, useChat });
-    const completion = truncateCompletion(raw, ex.entry_point);
-    const program = buildProgram(ex.prompt, completion, ex.test, ex.entry_point);
-
-    const rr = runProgram(program, timeoutSec, memoryMb);
-    if (rr.ok) nPass++;
-    else if (runtimeValue("MLX_BUN_EVAL_DEBUG") === "1") {
-      const err = rr.stderr ?? "";
-      const reason = rr.timedOut ? "TIMEOUT"
-        : completion.trim() === "" ? "EMPTY"
-        : /SyntaxError|IndentationError/.test(err) ? "SYNTAX"
-        : /AssertionError/.test(err) ? "WRONG-ANSWER"
-        : /Error|Exception|Traceback/.test(err) ? "RUNTIME-ERR"
-        : "OTHER";
-      const errSnip = err.trim().split("\n").slice(-1)[0]?.slice(0, 120) ?? "";
-      process.stderr.write(`\n  [FAIL ${ex.entry_point}] ${reason} complen=${completion.length} :: ${errSnip}\n`);
-    }
-
-    if ((k + 1) % 5 === 0 || k + 1 === idx.length)
-      process.stderr.write(
-        `\r  humaneval ${k + 1}/${idx.length}  pass@1=${((nPass / (k + 1)) * 100).toFixed(1)}%`,
-      );
-  }
-  process.stderr.write("\n");
-  return { nPass, nTotal: idx.length, accuracy: idx.length ? nPass / idx.length : 0 };
-}