Something went wrong. Try again.
Parametric learning in your code harness.
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251// Fine-tuning a machine's model on good runs: pick the runs, write them as// training data, and run supervised fine-tuning on Tinker through// tinker-cookbook (scripts/train_sft.py). Whether the result is better is// decided elsewhere, by running it against the live version on held-out// cases (see sdk.ts, train()).
import { spawn } from "node:child_process";import { appendFile, mkdir, readFile, writeFile } from "node:fs/promises";import path from "node:path";import { fileURLToPath } from "node:url";import { caseOf } from "./converse.ts";import { conversation, stepFailed, uniformContent, withoutThinking, type ChatMessage, type ExportLine, type ToolSpec } from "./export.ts";import type { Tool } from "./harness.ts";import type { Machine } from "./machine.ts";import { TINKER_TRAIN_PRICES } from "./pricing.ts";import { rubricFingerprint } from "./rubric.ts";import type { Run } from "./types.ts";
export type DataSelection = { /** Lowest scorecard total to train on. Default 1: only runs that met every criterion. */ minScore?: number; /** Keep at most this many runs per case, best first (then newest). Default 2. */ bestOf?: number; /** Cases never to train on (the held-out inputs). */ exclude?: Set<string>; /** Only runs by these models, when given (e.g. the live model, for rejection sampling). */ models?: Set<string>;};
/** Graded runs worth training on: scored under the current rubric, at least minScore, the best few per case. */export function selectRuns(machine: Machine, runs: Run[], options: DataSelection = {}): Run[] { const minScore = options.minScore ?? 1; const bestOf = options.bestOf ?? 2; const fp = rubricFingerprint(machine.rubric); const byCase = new Map<string, Run[]>(); for (const run of runs) { if (run.machine !== machine.name || run.harness === "echo" || !run.scorecard) continue; if (run.conversation && !run.conversation.final) continue; if (run.scorecard.rubric !== fp || run.scorecard.total < minScore) continue; if (options.exclude?.has(caseOf(run)) || options.exclude?.has(run.input)) continue; if (options.models && !options.models.has(run.model)) continue; // Train only on runs made under the instructions the model will run with (not, say, an optimizer's rewrite). if (run.instructions !== undefined && run.instructions !== machine.instructions) continue; byCase.set(caseOf(run), [...(byCase.get(caseOf(run)) ?? []), run]); } const out: Run[] = []; for (const group of byCase.values()) { group.sort((a, b) => b.scorecard!.total - a.scorecard!.total || b.id.localeCompare(a.id)); out.push(...group.slice(0, bestOf)); } return out;}
/** * One training example per assistant step of a conversation: the context the * model saw at that step, ending with what it produced. Earlier turns of a * conversation carry no reasoning (harnesses do not send it back across * turns). Within the run, `keepStepReasoning` keeps earlier steps' thinking * in the context: a chat-completions harness sends each step's message back * whole, reasoning included, and Tinker's endpoints render it (checked against * their prompt_tokens for Qwen3.6 and Inkling-Small). Without it, earlier * steps lose their thinking. Each example is trained on its last message only * (train_sft.py's default). Steps whose tool calls all failed (e.g. an edit * whose snippet was not found) are left out, so the model is not taught to * repeat them; they stay in later steps' context. */export function stepExamples(messages: ChatMessage[], options: { skipFailedSteps?: boolean; tools?: Tool[]; keepStepReasoning?: boolean } = {}): ChatMessage[][] { const skipFailed = options.skipFailedSteps ?? true; const lastUser = messages.findLastIndex((m) => m.role === "user"); const out: ChatMessage[][] = []; for (let i = lastUser + 1; i < messages.length; i++) { if (messages[i].role !== "assistant" || (skipFailed && stepFailed(messages, i, options.tools))) continue; out.push([...messages.slice(0, i).map((m, j) => (options.keepStepReasoning && j > lastUser ? m : withoutThinking(m))), messages[i]]); } return out;}
/** Harnesses that send a step's whole reply back to the model, reasoning included (chat completions). */const RESENDS_REASONING = new Set(["tinker", "openrouter", "openai", "chat-completions"]);
/** * One example for a whole run: its conversation, with every assistant * message a target except steps whose tool calls all failed. When the harness * sends each step back with its reasoning (RESENDS_REASONING), every step's * prompt is a prefix of this sequence, so it trains the same targets on the * same contexts as stepExamples at a fraction of the tokens: on an * Inkling-Small dev run, 7,699 tokens against 121,144 for its 29 step * prompts, which were token for token prefixes of it. */export function wholeRunExample(messages: ChatMessage[], tools: Tool[] = []): ChatMessage[] { const lastUser = messages.findLastIndex((m) => m.role === "user"); return messages.map((m, i) => (m.role === "assistant" && i > lastUser ? { ...m, trainable: !stepFailed(messages, i, tools) } : m));}
/** Whether a base model's training data can be whole runs (see wholeRunExample): verified for Inkling's renderer. */export function wholeRunsFor(baseModel: string): boolean { return baseModel.startsWith("thinkingmachines/Inkling");}
/** * Training examples from runs, with the tools the model was offered: one per * assistant step (see stepExamples), or with `wholeRuns`, one per run whose * harness resends reasoning (see wholeRunExample). */export function trainingLines(machine: Machine, runs: Run[], options: { reasoning?: boolean; wholeRuns?: boolean } = {}): ExportLine[] { const reasoning = options.reasoning === false ? null : (machine.transformReasoning ?? ((r: string) => r)); const tools: ToolSpec[] = machine.tools.map((t) => ({ name: t.name, description: t.description, parameters: t.inputSchema })); const lines: ExportLine[] = []; for (const run of runs) { if (options.wholeRuns && RESENDS_REASONING.has(run.harness)) { const messages = wholeRunExample(conversation(run, machine.instructions, reasoning), machine.tools); lines.push({ messages, ...(tools.length ? { tools } : {}), machine: machine.name, run: run.id, score: run.scorecard!.total, whole_run: true }); continue; } for (const messages of stepExamples(conversation(run, machine.instructions, reasoning), { tools: machine.tools, keepStepReasoning: RESENDS_REASONING.has(run.harness) })) { lines.push({ messages, ...(tools.length ? { tools } : {}), machine: machine.name, run: run.id, score: run.scorecard!.total }); } } return uniformContent(lines);}
/** * A rough upper estimate of the tokens a training file holds: 2.5 characters * per token over the messages and tool declarations as JSON. (Measured on dev * runs with Qwen3's renderer: about 4 characters per token.) */export function estimateTokens(lines: ExportLine[]): number { return Math.ceil(lines.reduce((n, l) => n + JSON.stringify(l.messages).length + (l.tools ? JSON.stringify(l.tools).length : 0), 0) / 2.5);}
/** Estimated training cost in USD: tokens × epochs at Tinker's training price for the base model. */export function trainingCost(baseModel: string, tokens: number, epochs: number): number { const price = TINKER_TRAIN_PRICES[baseModel]; if (price === undefined) throw new Error(`No Tinker training price for ${baseModel}; cannot bound the cost of training it.`); return (tokens * epochs * price) / 1_000_000;}
export type TrainJob = { /** JSONL training data. */ data: string; /** Directory for the run's checkpoints and metrics, under .machine/training/ to show on the dashboard. */ logDir: string; baseModel: string; epochs: number; learningRate?: number; batchSize?: number; /** Thinking effort the runs sampled at, 0 to 0.99, for models whose prompt carries it (Inkling). Default 0.9, the endpoint's. */ effort?: number; log?: (line: string) => void; signal?: AbortSignal;};
export type TrainResult = { /** Sampler path of the final checkpoint: the model id to run it with. */ checkpoint: string; steps: number; /** Tokens processed across all steps (every epoch), as logged; what training is billed on. */ tokens?: number; /** Training loss of the first and last logged steps, when logged. */ firstLoss?: number; lastLoss?: number;};
/** Something that trains: tinkerSft by default; tests pass a fake. */export type Trainer = (job: TrainJob) => Promise<TrainResult>;
/** * Supervised fine-tuning on Tinker with scripts/train_sft.py, run by * `python` (default $MACHINE_PYTHON, else python3), which must have * tinker-cookbook installed. Needs TINKER_API_KEY. The trainer's own output * goes to <logDir>/train.log; progress is reported from its metrics. */export function tinkerSft(options: { python?: string } = {}): Trainer { return async (job) => { const python = options.python ?? process.env.MACHINE_PYTHON ?? "python3"; const script = fileURLToPath(new URL("../scripts/train_sft.py", import.meta.url)); const args = [script, job.data, job.logDir, "--model", job.baseModel, "--epochs", String(job.epochs)]; if (job.learningRate !== undefined) args.push("--lr", String(job.learningRate)); if (job.batchSize !== undefined) args.push("--batch-size", String(job.batchSize)); if (job.effort !== undefined) args.push("--effort", String(job.effort)); await mkdir(job.logDir, { recursive: true, mode: 0o700 }); const logFile = path.join(job.logDir, "train.log"); let reported = -1; const progress = setInterval(async () => { const steps = (await readFile(path.join(job.logDir, "metrics.jsonl"), "utf8").catch(() => "")).split("\n").filter((l) => l.trim()); if (steps.length === 0) return; const last = JSON.parse(steps.at(-1)!) as { step?: number; train_mean_nll?: number; progress?: number }; // About ten lines per run: one per tenth of the way through. const tenth = Math.floor((last.progress ?? 0) * 10); if (tenth > reported) { job.log?.(` step ${(last.step ?? steps.length - 1) + 1}${last.progress ? ` (${Math.round(last.progress * 100)}%)` : ""}: loss ${last.train_mean_nll?.toFixed(3) ?? "?"}`); reported = tenth; } }, 5_000); try { await new Promise<void>((resolve, reject) => { const proc = spawn(python, args, { stdio: ["ignore", "pipe", "pipe"], signal: job.signal }); let tail = ""; const onData = (chunk: Buffer) => { tail = (tail + chunk.toString()).slice(-4000); void appendFile(logFile, chunk, { mode: 0o600 }).catch(() => {}); }; proc.stdout.on("data", onData); proc.stderr.on("data", onData); proc.once("error", (error) => reject(new Error(`Could not start ${python} (set MACHINE_PYTHON to a Python with tinker-cookbook): ${error.message}`))); proc.once("exit", (exitCode) => (exitCode === 0 ? resolve() : reject(new Error(`Training exited with code ${exitCode} (full output in ${logFile}): ${tail.trim().slice(-1500)}`)))); }); } finally { clearInterval(progress); } return readTrainingLog(job.logDir); };}
/** The final checkpoint and losses of a tinker-cookbook training log directory. */export async function readTrainingLog(logDir: string): Promise<TrainResult> { const jsonl = async (file: string) => (await readFile(path.join(logDir, file), "utf8").catch(() => "")) .split("\n") .filter((l) => l.trim()) .map((l) => JSON.parse(l) as Record<string, unknown>); const checkpoints = await jsonl("checkpoints.jsonl"); const final = checkpoints.findLast((c) => typeof c.sampler_path === "string"); if (!final) throw new Error(`Training finished without a sampler checkpoint in ${logDir}/checkpoints.jsonl.`); const metrics = (await jsonl("metrics.jsonl")).filter((m) => typeof m.train_mean_nll === "number"); const tokens = metrics.reduce((n, m) => n + (typeof m.num_tokens === "number" ? m.num_tokens : 0), 0); return { checkpoint: final.sampler_path as string, steps: metrics.length, ...(tokens ? { tokens } : {}), ...(metrics.length ? { firstLoss: metrics[0].train_mean_nll as number, lastLoss: metrics.at(-1)!.train_mean_nll as number } : {}), };}
export async function writeLines(file: string, lines: ExportLine[]): Promise<void> { await mkdir(path.dirname(file), { recursive: true, mode: 0o700 }); await writeFile(file, lines.map((l) => JSON.stringify(l)).join("\n") + (lines.length ? "\n" : ""), { mode: 0o600 });}
/** Named reasoning efforts as Tinker's OpenAI endpoint maps them to Inkling's thinking effort. */export const EFFORT_PRESETS: Record<string, number> = { none: 0, minimal: 0.1, low: 0.2, medium: 0.7, high: 0.9, xhigh: 0.99 };
/** The thinking effort runs sampled at: a preset or a number; unset means the endpoint's default, high. */export function effortScalar(effort: string | number | undefined): number { if (effort === undefined) return EFFORT_PRESETS.high; if (typeof effort === "number") return effort; const value = EFFORT_PRESETS[effort]; if (value === undefined) throw new Error(`Unknown reasoning effort "${effort}".`); return value;}