Something went wrong. Try again.
Parametric learning in your code harness.
Something went wrong. Try again.
23 kB · 470 lines
TypeScript
at letta-code
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471import type { ExtensionAPI, ExtensionCommandContext } from "@earendil-works/pi-coding-agent";import { execFile as execFileCb } from "node:child_process";import { existsSync, readFileSync } from "node:fs";import { mkdir, writeFile } from "node:fs/promises";import path from "node:path";import { fileURLToPath } from "node:url";import { promisify } from "node:util";import { machineDataFile } from "../lib/storage.ts";import { parseOptions, readJsonl, rel, shellSplit } from "../lib/cli.ts";import { registerMachineReporter, setMachineReportMirror } from "../lib/report.ts";import { pythonBin } from "../lib/python.ts";
const execFile = promisify(execFileCb);// The 64K base; /machine student in machine.ts defaults to the 256K peft// entry (DEFAULT_FT_MODEL). They differ on purpose.const DEFAULT_MODEL = "thinkingmachines/Inkling-Small";const PACKAGE_ROOT = path.resolve(path.dirname(fileURLToPath(import.meta.url)), "..");
type EvalSummary = { target?: string; renderer_model?: string; renderer?: string; effort?: number | null; num_examples?: number; num_correct?: number; accuracy?: number; results?: Array<{ index: number; correct: boolean; expected: string; output: string }>;};
function isInklingModel(model?: string): boolean { return Boolean(model?.startsWith("thinkingmachines/Inkling"));}
function compactMetrics(row: Record<string, unknown>, limit = 28): string { const interesting = Object.entries(row) .filter(([k, v]) => typeof v === "number" && (k.includes("loss") || k.includes("score") || k.includes("reward") || k.includes("kl") || k.includes("nll") || k.startsWith("progress/") || k.startsWith("time/total")) ) .slice(0, limit); if (interesting.length === 0) return JSON.stringify(row, null, 2).slice(0, 2000); return interesting.map(([k, v]) => `- ${k}: ${typeof v === "number" ? Number(v.toFixed(6)) : v}`).join("\n");}
function latestMetrics(metricsPath: string): string | undefined { if (!existsSync(metricsPath)) return undefined; try { const last = readJsonl(metricsPath).slice(-1)[0]; return last ? compactMetrics(last) : undefined; } catch (error) { return `Could not parse ${metricsPath}: ${error instanceof Error ? error.message : String(error)}`; }}
function renderTemplate(name: string, replacements: Record<string, string>): string { let text = readFileSync(path.join(PACKAGE_ROOT, "templates", name), "utf8"); for (const [key, value] of Object.entries(replacements)) { text = text.replaceAll(key, value); } const leftover = text.match(/__[A-Z_]+__/g); if (leftover?.length) { throw new Error(`unsubstituted template placeholders in ${name}: ${[...new Set(leftover)].join(", ")}`); } return text;}
function validateDataset(filePath: string): string { if (!existsSync(filePath)) return `File not found: ${filePath}`; const lines = readFileSync(filePath, "utf8").split(/\n/).filter((line) => line.trim()); let ok = 0; let assistantMessages = 0; let userMessages = 0; const issues: string[] = []; const samples: string[] = [];
lines.forEach((line, idx) => { try { const obj = JSON.parse(line); if (!Array.isArray(obj.messages)) { issues.push(`line ${idx + 1}: missing messages[]`); return; } if (obj.messages.length < 2) issues.push(`line ${idx + 1}: fewer than 2 messages`); for (const [j, msg] of obj.messages.entries()) { if (!msg || typeof msg !== "object") { issues.push(`line ${idx + 1} message ${j}: not an object`); continue; } if (!["system", "user", "assistant", "tool"].includes(msg.role)) { issues.push(`line ${idx + 1} message ${j}: unexpected role ${JSON.stringify(msg.role)}`); } if (msg.role === "assistant") assistantMessages++; if (msg.role === "user") userMessages++; if (!(typeof msg.content === "string" || Array.isArray(msg.content))) { issues.push(`line ${idx + 1} message ${j}: content should be string or content-part array`); } } ok++; if (samples.length < 3) { const preview = obj.messages .map((m: { role?: string; content?: unknown }) => `${m.role}: ${typeof m.content === "string" ? m.content : JSON.stringify(m.content)}`) .join("\n") .slice(0, 1200); samples.push(`### Example ${idx + 1}\n\n\`\`\`text\n${preview}\n\`\`\``); } } catch (error) { issues.push(`line ${idx + 1}: invalid JSON (${error instanceof Error ? error.message : String(error)})`); } });
const warnings: string[] = []; if (assistantMessages === 0) warnings.push("No assistant messages found; SFT will have nothing obvious to train on."); if (userMessages === 0) warnings.push("No user messages found; expected chat JSONL may be malformed."); if (lines.length < 20) warnings.push("Very small dataset; use this only for smoke tests.");
return [ `# Dataset validation: ${filePath}`, `- JSONL rows: ${lines.length}`, `- Valid rows: ${ok}`, `- User messages: ${userMessages}`, `- Assistant messages: ${assistantMessages}`, issues.length ? `\n## Issues\n${issues.slice(0, 50).map((x) => `- ${x}`).join("\n")}` : "\n## Issues\nNone found by the lightweight validator.", warnings.length ? `\n## Warnings\n${warnings.map((x) => `- ${x}`).join("\n")}` : "", `\n## Samples\n${samples.join("\n\n")}`, ].filter(Boolean).join("\n");}
async function validateDatasetWithPython(filePath: string, model: string, maxExamples: number, maxLength: number, cwd = path.dirname(filePath)): Promise<string> { const script = path.join(PACKAGE_ROOT, "scripts", "validate-dataset.py"); const { stdout } = await execFile(pythonBin(cwd, PACKAGE_ROOT), [script, filePath, model, String(maxExamples), String(maxLength)], { timeout: 240_000, maxBuffer: 12 * 1024 * 1024, }); const result = JSON.parse(stdout.trim().split(/\n/).slice(-1)[0] ?? "{}"); if (!result.ok && result.stage !== "done") { const installCommand = isInklingModel(model) ? "uv pip install -U 'tinker-cookbook[inkling]'" : "uv pip install -U tinker-cookbook"; return [ `# Python-backed validation could not run`, `- Stage: ${result.stage ?? "unknown"}`, `- Error: ${result.error ?? "unknown"}`, "", "Install/upgrade dependencies with:", "```bash", installCommand, "```", ].join("\n"); }
const readiness = String(result.readiness ?? "UNKNOWN"); const icon = readiness === "READY" ? "✅" : readiness === "SMOKE ONLY" ? "⚠️" : "❌"; const fmtStats = (s: { count?: number; min?: number; median?: number; p90?: number; p95?: number; max?: number; mean?: number } | undefined) => s?.count ? `count=${s.count}, min=${s.min}, p50=${s.median}, p90=${s.p90}, p95=${s.p95}, max=${s.max}, mean=${Number(s.mean).toFixed(1)}` : "n/a"; const recommendation = readiness === "READY" ? "Run `/train smoke --yes`, inspect metrics, then scale only after defining an eval." : readiness === "SMOKE ONLY" ? "Run only a tiny smoke test for now. Inspect warnings before larger runs." : "Fix the data issues above before creating a training run.";
return [ `# Machine data readiness report`, `## ${icon} ${readiness}`, `- File: \`${filePath}\``, `- Model: \`${result.model}\``, `- Recommended renderer: \`${result.renderer}\``, `- Rows: ${result.rows} (${result.valid_rows} valid JSONL conversation rows)`, `- Checked with tokenizer/renderer: ${result.checked_examples}`, "", `## Recommendation\n${recommendation}`, (result.serious ?? []).length ? `\n## Must fix\n${(result.serious ?? []).map((x: string) => `- ${x}`).join("\n")}` : "", (result.warnings ?? []).length ? `\n## Warnings\n${(result.warnings ?? []).map((x: string) => `- ${x}`).join("\n")}` : "", `\n## Token stats\n- Input tokens: ${fmtStats(result.token_lengths)}\n- Trainable assistant tokens: ${fmtStats(result.trainable_tokens)}`, ].filter(Boolean).join("\n");}
function readEvalSummary(filePath: string): EvalSummary { return JSON.parse(readFileSync(filePath, "utf8")) as EvalSummary;}
function formatEvalSummary(filePath: string): string { const result = readEvalSummary(filePath); const accuracy = typeof result.accuracy === "number" ? `${(result.accuracy * 100).toFixed(1)}%` : "n/a"; const examples = result.num_examples ?? result.results?.length ?? 0; const correct = result.num_correct ?? result.results?.filter((r) => r.correct).length ?? 0; return [ `- File: \`${filePath}\``, `- Target: \`${result.target ?? "unknown"}\``, `- Renderer: \`${result.renderer ?? "unknown"}\``, result.effort !== null && result.effort !== undefined ? `- Inkling effort: ${result.effort}` : "", `- Accuracy: ${accuracy} (${correct}/${examples})`, ].filter(Boolean).join("\n");}
function compareEvalSummaries(baselinePath: string, candidatePath: string): string { const baseline = readEvalSummary(baselinePath); const candidate = readEvalSummary(candidatePath); const bAcc = baseline.accuracy ?? 0; const cAcc = candidate.accuracy ?? 0; return [ `# Eval comparison`, `- Baseline: \`${baselinePath}\` — ${(bAcc * 100).toFixed(1)}% (${baseline.num_correct}/${baseline.num_examples})`, `- Candidate: \`${candidatePath}\` — ${(cAcc * 100).toFixed(1)}% (${candidate.num_correct}/${candidate.num_examples})`, `- Delta: ${cAcc - bAcc >= 0 ? "+" : ""}${((cAcc - bAcc) * 100).toFixed(1)} points`, ].join("\n");}
async function writeFiles(cwd: string, files: Array<{ rel: string; content: string }>, force: boolean): Promise<string[]> { const written: string[] = []; for (const file of files) { const target = path.join(cwd, file.rel); if (existsSync(target) && !force) throw new Error(`${target} already exists. Re-run with --force to overwrite.`); await mkdir(path.dirname(target), { recursive: true }); await writeFile(target, file.content); written.push(file.rel); } return written;}
export default function (pi: ExtensionAPI) { const sendReport = registerMachineReporter(pi);
pi.on("session_start", (_event, ctx) => { setMachineReportMirror(!ctx.hasUI); });
pi.registerCommand("train", { description: "Validate data, scaffold SFT, smoke-test, and evaluate Tinker runs", handler: async (args, ctx: ExtensionCommandContext) => { const [subcommandRaw, ...rest] = shellSplit(args); const subcommand = subcommandRaw ?? "help";
if (subcommand === "help" || subcommand === "--help" || subcommand === "-h") { sendReport("machine train", [ "```text", "/train validate [jsonl]", "/train sft [jsonl]", "/train smoke --yes", "/train eval init", "/train eval baseline --yes", "/train eval checkpoint tinker://... --yes", "/train eval compare eval_results/baseline.json eval_results/checkpoint.json", "```", "", `Defaults to \`${machineDataFile(ctx.cwd, "train")}\` and \`thinkingmachines/Inkling-Small\`.`, ].join("\n")); return; }
if (subcommand === "validate") { const { positional, options } = parseOptions(rest); const file = positional[0]; const dataFile = file ? path.resolve(ctx.cwd, file) : machineDataFile(ctx.cwd, "train"); const model = String(options.model ?? DEFAULT_MODEL); const lightweight = validateDataset(dataFile); if (options.quick === true) { sendReport("Dataset validation", lightweight); return; } try { const python = await validateDatasetWithPython(dataFile, model, Number(options.examples ?? 200), Number(options["max-length"] ?? 32768), ctx.cwd); sendReport("Dataset validation", `${python}\n\n---\n\n${lightweight}`); } catch (error) { sendReport("Dataset validation", `${lightweight}\n\n## Python-backed validation failed\n${error instanceof Error ? error.message : String(error)}`, "warning"); } return; }
if (subcommand === "sft") { const { positional, options } = parseOptions(rest); const dataFileArg = positional[0]; const dataFile = dataFileArg ? path.resolve(ctx.cwd, dataFileArg) : machineDataFile(ctx.cwd, "train"); if (!existsSync(dataFile)) { sendReport("SFT scaffold", `Data file not found: ${dataFile}`, "error"); return; } const model = String(options.model ?? DEFAULT_MODEL); const maxSteps = String(options.steps ?? options.max_steps ?? "20"); const batchSize = String(options["batch-size"] ?? "8"); const learningRate = String(options.lr ?? options["learning-rate"] ?? "2e-4"); const numEpochs = String(options.epochs ?? "1"); const testSize = String(options["test-size"] ?? "0"); const maxLength = String(options["max-length"] ?? "32768"); const logPath = String(options.log ?? `logs/sft-${new Date().toISOString().replace(/[:.]/g, "-").slice(0, 19)}`); try { const written = await writeFiles(ctx.cwd, [ { rel: "train_sft.py", content: renderTemplate("train_sft.py", { "__MODEL__": JSON.stringify(model), "__MAX_LENGTH__": maxLength, "__BATCH_SIZE__": batchSize, "__DATA_FILE__": JSON.stringify(dataFile), "__TEST_SIZE__": testSize, "__LOG_PATH__": JSON.stringify(logPath), "__LEARNING_RATE__": learningRate, "__NUM_EPOCHS__": numEpochs, "__MAX_STEPS__": maxSteps, }), }, { rel: "tinker.yaml", content: `task: sft\ndata: ${dataFile}\nmodel: ${model}\nlog_path: ${logPath}\nnum_epochs: ${numEpochs}\nmax_steps: ${maxSteps}\nbatch_size: ${batchSize}\nlearning_rate: ${learningRate}\n`, }, ], options.force === true); sendReport("SFT scaffold created", [ `Wrote ${written.map((x) => `\`${x}\``).join(", ")} from \`${rel(ctx.cwd, dataFile)}\`.`, "", "Next:", "```text", dataFileArg ? `/train validate ${dataFileArg}` : "/train validate", "/train smoke --yes", "```", ].join("\n"), "success"); } catch (error) { sendReport("SFT scaffold", error instanceof Error ? error.message : String(error), "warning"); } return; }
if (subcommand === "smoke") { const { positional, options } = parseOptions(rest); const script = path.resolve(ctx.cwd, positional[0] ?? "train_sft.py"); if (!existsSync(script)) { sendReport("Smoke test", `Script not found: ${script}. Run \`/train sft\` first.`, "error"); return; } const ok = options.yes === true || !ctx.hasUI ? true : await ctx.ui.confirm("Run smoke test?", "This may create a real Tinker training client and incur small API usage."); if (!ok) { sendReport("Smoke test", "Cancelled.", "warning"); return; } try { ctx.ui.setStatus("machine", "smoke running"); const { stdout, stderr } = await execFile(pythonBin(ctx.cwd, PACKAGE_ROOT), [script, "max_steps=2"], { cwd: ctx.cwd, timeout: Number(options.timeout ?? 1_800_000), maxBuffer: 10 * 1024 * 1024, }); const combined = `${stdout}\n${stderr}`.trim(); const logMatch = combined.match(/log_path='([^']+)'|log_path=([^,\s)]+)/); const logDir = logMatch ? path.resolve(ctx.cwd, (logMatch[1] ?? logMatch[2] ?? "").replace(/^['"]|['"]$/g, "")) : undefined; const metrics = logDir ? latestMetrics(path.join(logDir, "metrics.jsonl")) : undefined; sendReport("Smoke test passed", [ "The 2-step command completed.", metrics ? `\n## Latest metrics\n${metrics}` : "", `\n## Output tail\n\`\`\`text\n${combined.split(/\n/).slice(-80).join("\n")}\n\`\`\``, ].filter(Boolean).join("\n"), "success"); } catch (error: unknown) { const err = error as { message?: string; stdout?: string; stderr?: string }; sendReport("Smoke test failed", [ err?.message ?? String(error), err?.stdout ? `\n## stdout\n\`\`\`text\n${String(err.stdout).split(/\n/).slice(-80).join("\n")}\n\`\`\`` : "", err?.stderr ? `\n## stderr\n\`\`\`text\n${String(err.stderr).split(/\n/).slice(-80).join("\n")}\n\`\`\`` : "", ].filter(Boolean).join("\n"), "error"); } finally { ctx.ui.setStatus("machine", undefined); } return; }
if (subcommand === "eval") { const [actionRaw, ...evalRest] = rest; const action = actionRaw ?? "help"; const { positional, options } = parseOptions(evalRest);
if (action === "help" || action === "--help") { sendReport("machine eval", [ "- `/train eval init` — create `eval.py` and `data/eval.jsonl`.", "- `/train eval baseline --yes` — evaluate the base model.", "- `/train eval checkpoint tinker://... --yes` — evaluate a sampler checkpoint.", "- `/train eval compare a.json b.json` — compare results.", ].join("\n")); return; }
if (action === "init") { try { const written = await writeFiles(ctx.cwd, [ { rel: "eval.py", content: readFileSync(path.join(PACKAGE_ROOT, "templates", "eval.py"), "utf8") }, { rel: path.join("data", "eval.jsonl"), content: [ JSON.stringify({ messages: [{ role: "user", content: "What is --as?" }], expected: "rewrite", match: "contains" }), JSON.stringify({ messages: [{ role: "user", content: "Where do tagged traces go?" }], expected: "traces.jsonl", match: "contains" }), ].join("\n") + "\n", }, ], options.force === true); sendReport("Eval initialized", `Wrote ${written.map((x) => `\`${x}\``).join(", ")}.`, "success"); } catch (error) { sendReport("Eval init", error instanceof Error ? error.message : String(error), "warning"); } return; }
if (action === "baseline" || action === "checkpoint") { const script = path.resolve(ctx.cwd, String(options.script ?? "eval.py")); if (!existsSync(script)) { sendReport("Eval", `Missing \`${script}\`. Run \`/train eval init\` first.`, "warning"); return; } const model = String(options.model ?? options["base-model"] ?? DEFAULT_MODEL); const data = path.resolve(ctx.cwd, String(options.data ?? "data/eval.jsonl")); if (!existsSync(data)) { sendReport("Eval", `Eval data not found: ${data}`, "error"); return; } const checkpoint = action === "checkpoint" ? positional[0] : undefined; if (action === "checkpoint" && !checkpoint) { sendReport("Eval checkpoint", "Usage: `/train eval checkpoint tinker://.../sampler_weights/...`", "warning"); return; } const out = path.resolve(ctx.cwd, String(options.out ?? (action === "baseline" ? "eval_results/baseline.json" : `eval_results/${String(checkpoint).split("/").slice(-1)[0] || "checkpoint"}.json`))); const ok = options.yes === true || !ctx.hasUI ? true : await ctx.ui.confirm("Run eval?", "This will sample from Tinker and may incur API usage."); if (!ok) { sendReport("Eval", "Cancelled.", "warning"); return; } const pythonArgs = action === "baseline" ? [script, "--base-model", model, "--data", data, "--out", out] : [script, "--model-path", String(checkpoint), "--renderer-model", model, "--data", data, "--out", out]; if (options.effort) pythonArgs.push("--effort", String(options.effort)); if (options.limit) pythonArgs.push("--limit", String(options.limit)); try { ctx.ui.setStatus("machine", `eval ${action}`); const { stdout, stderr } = await execFile(pythonBin(ctx.cwd, PACKAGE_ROOT), pythonArgs, { cwd: ctx.cwd, timeout: Number(options.timeout ?? 1_800_000), maxBuffer: 12 * 1024 * 1024, }); sendReport("Eval completed", [ formatEvalSummary(out), `\n## Output tail\n\`\`\`text\n${`${stdout}\n${stderr}`.trim().split(/\n/).slice(-60).join("\n")}\n\`\`\``, ].join("\n"), "success"); } catch (error: unknown) { const err = error as { message?: string; stdout?: string; stderr?: string }; sendReport("Eval failed", [ err?.message ?? String(error), err?.stdout ? `\n## stdout\n\`\`\`text\n${String(err.stdout).split(/\n/).slice(-80).join("\n")}\n\`\`\`` : "", err?.stderr ? `\n## stderr\n\`\`\`text\n${String(err.stderr).split(/\n/).slice(-80).join("\n")}\n\`\`\`` : "", ].filter(Boolean).join("\n"), "error"); } finally { ctx.ui.setStatus("machine", undefined); } return; }
if (action === "compare") { const baseline = positional[0] ? path.resolve(ctx.cwd, positional[0]) : path.resolve(ctx.cwd, "eval_results/baseline.json"); const candidate = positional[1] ? path.resolve(ctx.cwd, positional[1]) : ""; if (!candidate) { sendReport("Eval compare", "Usage: `/train eval compare eval_results/baseline.json eval_results/checkpoint.json`", "warning"); return; } try { sendReport("Eval comparison", compareEvalSummaries(baseline, candidate)); } catch (error) { sendReport("Eval comparison", `Could not compare eval results: ${error instanceof Error ? error.message : String(error)}`, "error"); } return; }
sendReport("Eval", `Unknown eval action: ${action}. Run \`/train eval help\`.`, "warning"); return; }
sendReport("Unknown train command", `Unknown subcommand: ${subcommand}\n\nRun \`/train help\`.`, "warning"); }, });}