mlx-bun 0.4.0 → 1.0.0

package/src/dataset/index.ts

@@ -1,12 +0,0 @@
-// Public barrel for the native dataset-building engine: a TypeScript port of
-// optiq lab's 13 dataset templates. Wizards/server routes import from here.
-
-export { TEMPLATES, getTemplate, generate } from "./registry";
-export type { TemplateDef, TemplateField, GenerateResult } from "./registry";
-
-export { datasetRunner, registerDatasetRunner } from "./job";
-
-export { makeLlmClient } from "./llm";
-export type { LlmClient, ChatMessage, ChatOpts } from "./llm";
-
-export type { Row } from "./generators";

package/src/dataset/job.ts

@@ -1,81 +0,0 @@
-// The dataset JobRunner: wires a submitted dataset job to generate(). Speaks
-// the shared src/jobs/types.ts contract. The submit config carries the
-// template id, user inputs, output dir, the server's own loopback API base,
-// and an optional model name.
-
-import type { JobRunner } from "../jobs/types";
-import { makeLlmClient } from "./llm";
-import { generate } from "./registry";
-
-export const datasetRunner: JobRunner = async (emit, config) => {
-  const template_id = String(config.template_id ?? "");
-  const inputs = (config.inputs as Record<string, unknown>) ?? {};
-  const output_dir = String(config.output_dir ?? "");
-  const api_url = String(config.api_url ?? "");
-  const model_name = (config.model_name as string | undefined) ?? "local";
-
-  if (!template_id) throw new Error("dataset job: missing template_id");
-  if (!output_dir) throw new Error("dataset job: missing output_dir");
-
-  emit({
-    type: "stage",
-    stage: "starting",
-    progress: 0.05,
-    message: `Starting template ${template_id}…`,
-  });
-
-  const llm = api_url ? makeLlmClient(api_url, model_name) : undefined;
-  const r = await generate(template_id, inputs, output_dir, emit, llm);
-  return { outputPath: r.output_dir };
-};
-
-function applyRegister(mod: any): boolean {
-  const register =
-    mod?.registerRunner ?? mod?.register ?? mod?.default?.registerRunner ?? mod?.default?.register;
-  if (typeof register === "function") {
-    register("dataset", datasetRunner); // idempotent: registry is a Map by kind
-    return true;
-  }
-  return false;
-}
-
-/**
- * Best-effort self-registration with the jobs runner registry. Registers the
- * "dataset" kind so both submitInProcess (server path) and the subprocess
- * job-entry resolver can find this runner. Tries a synchronous require first
- * (so registration completes before any submit on the same tick — the
- * convention quantize/finetune use), then an async import as a fallback.
- * Wrapped in try/catch so importing job.ts never fails in tests or standalone
- * use where ../jobs/runner isn't present. Returns true if registration
- * succeeded synchronously; an async fallback may still complete shortly after.
- *
- * job-entry.ts's KIND_MODULES can map `dataset → ../dataset/job.ts` and rely on
- * the top-level call below; the server orchestrator can alternatively await
- * this explicitly at bootstrap.
- */
-export function registerDatasetRunner(): boolean {
-  // 1. Synchronous require — Bun supports import.meta.require for ESM. This is
-  //    the path that matters for in-process submit on the same tick.
-  try {
-    const req = (import.meta as any).require ?? (globalThis as any).require;
-    if (typeof req === "function") {
-      if (applyRegister(req("../jobs/runner"))) return true;
-    }
-  } catch {
-    // fall through to async
-  }
-  // 2. Async fallback — completes on a later microtask. The orchestrator that
-  //    needs guaranteed timing should call/await the runner module directly.
-  void (async () => {
-    try {
-      applyRegister(await import("../jobs/runner"));
-    } catch {
-      // runner module not present — orchestrator will wire registration.
-    }
-  })();
-  return false;
-}
-
-// Self-register on import so `kind: "dataset"` resolves wherever this module is
-// loaded (mirrors quantize/finetune's top-level registration on import).
-registerDatasetRunner();

package/src/dataset/llm.ts

@@ -1,71 +0,0 @@
-// Loopback chat client for LLM-driven dataset templates. Posts to the
-// mlx-bun server's own /v1/chat/completions endpoint (OpenAI shape, with
-// tool calling). Ports the _llm_call / _llm_chat / _llm_chat_raw helpers
-// from optiq's dataset_templates.py.
-
-export interface ChatMessage {
-  role: string;
-  content: string;
-  // Tool-call traces carry extra fields (tool_calls, tool_call_id, name).
-  [k: string]: unknown;
-}
-
-export interface ChatOpts {
-  maxTokens?: number;
-  temperature?: number;
-  /** Defaults to false: dataset generators need `content`, not a hidden
-   *  `<think>` channel that eats the whole token budget. */
-  enableThinking?: boolean;
-  /** OpenAI tools array, forwarded for tool-use traces. */
-  tools?: unknown[];
-}
-
-export interface LlmClient {
-  /** One-shot/multi-turn chat; returns the assistant content string. */
-  chat(messages: ChatMessage[], opts?: ChatOpts): Promise<string>;
-  /** Non-streaming chat returning the full JSON body (for tool_calls). */
-  chatRaw(messages: ChatMessage[], opts?: ChatOpts): Promise<any>;
-}
-
-const LOOPBACK_TOKEN = "sk-mlx-bun-local";
-
-/**
- * Build a loopback chat client targeting `apiUrl` (the server's own base,
- * e.g. http://127.0.0.1:8080). `model` defaults to "local", the stable id
- * mlx-bun advertises for its served model.
- */
-export function makeLlmClient(apiUrl: string, model = "local"): LlmClient {
-  const base = apiUrl.replace(/\/+$/, "");
-
-  async function chatRaw(messages: ChatMessage[], opts: ChatOpts = {}): Promise<any> {
-    const body: Record<string, unknown> = {
-      model,
-      messages,
-      max_tokens: opts.maxTokens ?? 512,
-      temperature: opts.temperature ?? 0.7,
-      chat_template_kwargs: { enable_thinking: opts.enableThinking ?? false },
-    };
-    if (opts.tools && opts.tools.length) body.tools = opts.tools;
-
-    const resp = await fetch(`${base}/v1/chat/completions`, {
-      method: "POST",
-      headers: {
-        "Content-Type": "application/json",
-        Authorization: `Bearer ${LOOPBACK_TOKEN}`,
-      },
-      body: JSON.stringify(body),
-    });
-    if (!resp.ok) {
-      const detail = await resp.text().catch(() => "");
-      throw new Error(`chat/completions ${resp.status}: ${detail.slice(0, 400)}`);
-    }
-    return resp.json();
-  }
-
-  async function chat(messages: ChatMessage[], opts?: ChatOpts): Promise<string> {
-    const data = await chatRaw(messages, opts);
-    return data?.choices?.[0]?.message?.content ?? "";
-  }
-
-  return { chat, chatRaw };
-}

package/src/dataset/registry.ts

@@ -1,528 +0,0 @@
-// Dataset-template registry: the 13 TemplateDefs ported field-for-field from
-// optiq lab's dataset_templates.py. This file is pure metadata (the UI renders
-// forms from it) plus the generate() dispatcher that does the 90/10 split and
-// JSONL writing exactly like the Python reference.
-
-import type { Emit } from "../jobs/types";
-import type { LlmClient } from "./llm";
-import {
-  genCodeCompletion,
-  genCotSynthesis,
-  genDpoPairs,
-  genFormatConversion,
-  genHfDatasetImport,
-  genMultiTurnChat,
-  genPromptReconstruction,
-  genRagQa,
-  genSelfInstruct,
-  genSftQa,
-  genStyleTransfer,
-  genToolUseTraces,
-  genVerifiedCode,
-  type Row,
-} from "./generators";
-
-/** A single user-facing form field. UI forms are built from these. */
-export interface TemplateField {
-  name: string;
-  type: "text" | "textarea" | "number";
-  label: string;
-  hint?: string;
-  required?: boolean;
-  default?: string | number;
-}
-
-/** Metadata for one dataset-building template. */
-export interface TemplateDef {
-  id: string;
-  label: string;
-  description: string;
-  output_format: "messages" | "prompt_completion" | "dpo" | "text";
-  needs_llm: boolean;
-  fields: TemplateField[];
-}
-
-export const TEMPLATES: TemplateDef[] = [
-  // ---- non-LLM templates (fast, local) ----
-  {
-    id: "sft_qa_pairs",
-    label: "SFT from QA pairs",
-    description:
-      "Paste Q/A pairs or upload a .txt with `Q: …\\nA: …` blocks. Produces chat-format JSONL.",
-    output_format: "messages",
-    needs_llm: false,
-    fields: [
-      {
-        name: "pairs_text",
-        type: "textarea",
-        label: "Pairs (Q: / A: blocks)",
-        hint: "Q: What is OptIQ?\\nA: A mixed-precision quantizer …",
-        required: true,
-      },
-    ],
-  },
-  {
-    id: "dpo_pref_pairs",
-    label: "DPO from preference pairs",
-    description: "Upload CSV with columns `prompt,chosen,rejected`. Emits DPO JSONL.",
-    output_format: "dpo",
-    needs_llm: false,
-    fields: [
-      {
-        name: "csv_text",
-        type: "textarea",
-        label: "CSV (header `prompt,chosen,rejected`)",
-        required: true,
-      },
-    ],
-  },
-  {
-    id: "code_completion",
-    label: "Code completion",
-    description:
-      "Walks a directory's .py files, splits each function at a random midpoint into prompt/completion pairs.",
-    output_format: "prompt_completion",
-    needs_llm: false,
-    fields: [
-      {
-        name: "src_dir",
-        type: "text",
-        label: "Source directory (containing .py files)",
-        required: true,
-      },
-      {
-        name: "max_pairs",
-        type: "number",
-        label: "Max pairs to produce",
-        default: 500,
-      },
-    ],
-  },
-  {
-    id: "hf_dataset_import",
-    label: "Hugging Face dataset import",
-    description:
-      "Pull a public dataset from the Hugging Face Hub by id, optionally " +
-      "filter rows by a column value, slice to a row cap, and emit the " +
-      "chosen output format. Stand-alone (no LLM call) and idempotent — " +
-      "use it as the first step of any pipeline that starts from a " +
-      "published corpus (EditLens, no_robots, dolly, your own dataset).",
-    output_format: "text",
-    needs_llm: false,
-    fields: [
-      {
-        name: "hf_id",
-        type: "text",
-        label: "Dataset id",
-        hint: "e.g. pangram/editlens_iclr or HuggingFaceH4/no_robots",
-        required: true,
-      },
-      {
-        name: "config",
-        type: "text",
-        label: "Config / subset (optional)",
-        hint: "Some datasets have multiple configs (e.g. wikitext-2-raw-v1).",
-      },
-      {
-        name: "split",
-        type: "text",
-        label: "Split",
-        default: "train",
-      },
-      {
-        name: "text_column",
-        type: "text",
-        label: "Text column",
-        hint: "Field on each row that holds the body text.",
-        default: "text",
-      },
-      {
-        name: "label_column",
-        type: "text",
-        label: "Filter column (optional)",
-        hint: "e.g. text_type for EditLens. Leave blank for no filter.",
-      },
-      {
-        name: "label_filter",
-        type: "text",
-        label: "Filter value (optional)",
-        hint: "Keep only rows where filter-column == this value. e.g. human_written.",
-      },
-      {
-        name: "max_rows",
-        type: "number",
-        label: "Row cap (0 = no cap)",
-        default: 1000,
-      },
-      {
-        name: "min_chars",
-        type: "number",
-        label: "Drop rows shorter than this many chars",
-        default: 200,
-      },
-      {
-        name: "output_format",
-        type: "text",
-        label: "Output format",
-        hint:
-          "One of: text | messages_user_only | prompt_completion. " +
-          "'text' writes {\"text\": ...}; 'messages_user_only' writes a " +
-          "messages row with the text as the user turn (downstream " +
-          "templates can read either).",
-        default: "text",
-      },
-    ],
-  },
-  {
-    id: "format_conversion",
-    label: "Format conversion",
-    description:
-      "Upload existing JSONL in any shape + a key mapping. Emits OptIQ-expected JSONL.",
-    output_format: "messages",
-    needs_llm: false,
-    fields: [
-      {
-        name: "input_jsonl",
-        type: "textarea",
-        label: "Existing JSONL",
-        required: true,
-      },
-      { name: "user_key", type: "text", label: "user field", default: "input" },
-      { name: "assistant_key", type: "text", label: "assistant field", default: "output" },
-    ],
-  },
-  // ---- LLM-driven templates (need a served model) ----
-  {
-    id: "style_transfer",
-    label: "Style transfer",
-    description:
-      "Provide reference samples and raw text. Uses the served model to rewrite in the reference style.",
-    output_format: "prompt_completion",
-    needs_llm: true,
-    fields: [
-      {
-        name: "reference_samples",
-        type: "textarea",
-        label: "Reference style samples (separated by ---)",
-        required: true,
-      },
-      {
-        name: "raw_text",
-        type: "textarea",
-        label: "Raw text to rewrite (one paragraph per row)",
-        required: true,
-      },
-    ],
-  },
-  {
-    id: "self_instruct",
-    label: "Self-instruct expansion",
-    description:
-      "Upload seed instructions; uses the served model (via data-designer when available) to generate K variants per seed.",
-    output_format: "messages",
-    needs_llm: true,
-    fields: [
-      {
-        name: "seeds",
-        type: "textarea",
-        label: "Seed instructions (one per line)",
-        required: true,
-      },
-      {
-        name: "variants_per_seed",
-        type: "number",
-        label: "Variants per seed",
-        default: 5,
-      },
-    ],
-  },
-  {
-    id: "prompt_reconstruction",
-    label: "Prompt reconstruction",
-    description:
-      "Build (AI draft → target) training pairs by working backwards: " +
-      "for each target paragraph the served model infers a likely prompt " +
-      "and writes a generic AI draft. The assistant target is the original " +
-      "paragraph verbatim, so facts and formatting are preserved by construction.",
-    output_format: "messages",
-    needs_llm: true,
-    fields: [
-      {
-        name: "target_text",
-        type: "textarea",
-        label: "Target paragraphs (separated by blank lines)",
-        hint:
-          "Paste the text you want the model to learn to produce. " +
-          "Posts, memos, brand-voice samples, edited drafts, etc.",
-        required: true,
-      },
-      {
-        name: "style",
-        type: "text",
-        label: "Style label",
-        hint: "Free text, surfaced in the system prompt.",
-        default: "direct technical blog",
-      },
-      {
-        name: "tone",
-        type: "text",
-        label: "Tone",
-        default: "analytical, clear, non-corporate",
-      },
-      {
-        name: "preserve",
-        type: "text",
-        label: "Things to preserve (comma-separated)",
-        default: "facts, names, numbers, URLs, citations, code blocks, quotes",
-      },
-      {
-        name: "avoid",
-        type: "text",
-        label: "Things to avoid (comma-separated)",
-        default: "em dashes, generic transitions, marketing language",
-      },
-    ],
-  },
-  {
-    id: "multi_turn_chat",
-    label: "Multi-turn chat synthesis",
-    description:
-      "Take seed user prompts and a persona; expand each into an " +
-      "N-turn synthetic user/assistant conversation by alternating " +
-      "model-as-assistant and model-as-followup-user calls.",
-    output_format: "messages",
-    needs_llm: true,
-    fields: [
-      {
-        name: "seeds",
-        type: "textarea",
-        label: "Seed user prompts (one per line)",
-        required: true,
-      },
-      {
-        name: "turns",
-        type: "number",
-        label: "Total turns per conversation (user + assistant pairs)",
-        default: 4,
-      },
-      {
-        name: "persona",
-        type: "text",
-        label: "Assistant persona (system prompt)",
-        default: "You are a helpful, concise assistant.",
-      },
-      {
-        name: "user_persona",
-        type: "text",
-        label: "Followup-user persona",
-        hint: "How follow-up questions should sound.",
-        default: "a curious developer probing for specifics",
-      },
-    ],
-  },
-  {
-    id: "tool_use_traces",
-    label: "Tool-use traces",
-    description:
-      "Generate (user, tool_call, tool_result, final) training traces " +
-      "in OpenAI tool-call shape. Provide a list of tool schemas plus " +
-      "scenario prompts; the model picks a tool, you supply a mocked " +
-      "result, the model writes the final answer.",
-    output_format: "messages",
-    needs_llm: true,
-    fields: [
-      {
-        name: "tools_json",
-        type: "textarea",
-        label: "Tool schemas (OpenAI tools format, JSON array)",
-        hint: "Paste a JSON array of {type, function:{name,description,parameters}}.",
-        required: true,
-      },
-      {
-        name: "scenarios",
-        type: "textarea",
-        label: "User scenarios (one per line)",
-        required: true,
-      },
-      {
-        name: "mock_results",
-        type: "textarea",
-        label: "Mocked tool results (one per line, matches scenarios order)",
-        hint: "Free text; gets passed back as the tool_result content.",
-        required: true,
-      },
-    ],
-  },
-  {
-    id: "rag_qa",
-    label: "RAG Q/A from documents",
-    description:
-      "Chunk pasted text into passages; for each passage the model " +
-      "writes a question whose answer is grounded in that passage and " +
-      "then writes the grounded answer. Outputs messages-format rows " +
-      "where the passage is in a system message and the question/answer " +
-      "in user/assistant turns.",
-    output_format: "messages",
-    needs_llm: true,
-    fields: [
-      {
-        name: "documents",
-        type: "textarea",
-        label: "Source documents (paragraphs separated by blank lines)",
-        required: true,
-      },
-      {
-        name: "questions_per_chunk",
-        type: "number",
-        label: "Questions per chunk",
-        default: 2,
-      },
-      {
-        name: "min_chunk_chars",
-        type: "number",
-        label: "Min chunk size (chars)",
-        default: 200,
-      },
-    ],
-  },
-  {
-    id: "cot_synthesis",
-    label: "Reasoning trace (CoT) synthesis",
-    description:
-      "For each question, the served model emits a step-by-step " +
-      "<think> trace followed by the final answer. Output is messages-" +
-      "format with the reasoning preserved inside the assistant content " +
-      "as a `<think>...</think>` block, matching the Qwen3 / GPT-OSS " +
-      "reasoning convention.",
-    output_format: "messages",
-    needs_llm: true,
-    fields: [
-      {
-        name: "questions",
-        type: "textarea",
-        label: "Questions (one per line)",
-        required: true,
-      },
-      {
-        name: "category",
-        type: "text",
-        label: "Domain hint",
-        default: "math, logic, planning, technical analysis",
-      },
-    ],
-  },
-  {
-    id: "verified_code",
-    label: "Verified code generation",
-    description:
-      "For each natural-language spec, the served model writes Python " +
-      "code AND a set of `assert` checks. We run the assertions in the " +
-      "sandbox and keep only the pairs that pass. Output is messages-" +
-      "format; failed pairs are dropped before writing the JSONL.",
-    output_format: "messages",
-    needs_llm: true,
-    fields: [
-      {
-        name: "specs",
-        type: "textarea",
-        label: "Natural-language specs (one per line)",
-        hint: "e.g. `write a function that returns the n-th Fibonacci number`",
-        required: true,
-      },
-      {
-        name: "language",
-        type: "text",
-        label: "Language",
-        default: "python",
-      },
-    ],
-  },
-];
-
-export function getTemplate(id: string): TemplateDef | undefined {
-  return TEMPLATES.find((t) => t.id === id);
-}
-
-/** Result of a successful generate() run. */
-export interface GenerateResult {
-  n_train: number;
-  n_valid: number;
-  output_dir: string;
-}
-
-type Inputs = Record<string, unknown>;
-
-/** Dispatch table: template id → generator function. */
-const GENERATORS: Record<
-  string,
-  (inputs: Inputs, emit: Emit, llm?: LlmClient) => Promise<Row[]>
-> = {
-  sft_qa_pairs: genSftQa,
-  dpo_pref_pairs: genDpoPairs,
-  code_completion: genCodeCompletion,
-  format_conversion: genFormatConversion,
-  hf_dataset_import: genHfDatasetImport,
-  style_transfer: genStyleTransfer,
-  self_instruct: genSelfInstruct,
-  prompt_reconstruction: genPromptReconstruction,
-  multi_turn_chat: genMultiTurnChat,
-  tool_use_traces: genToolUseTraces,
-  rag_qa: genRagQa,
-  cot_synthesis: genCotSynthesis,
-  verified_code: genVerifiedCode,
-};
-
-/**
- * Run the named template with `inputs`, write `train.jsonl` + `valid.jsonl`
- * into `outputDir`, and emit progress events. Mirrors the Python reference:
- * 90/10 split with `split_idx = max(1, floor(n * 0.9))`, valid = remainder
- * (or the last row if the remainder is empty).
- */
-export async function generate(
-  id: string,
-  inputs: Inputs,
-  outputDir: string,
-  emit: Emit,
-  llm?: LlmClient,
-): Promise<GenerateResult> {
-  const gen = GENERATORS[id];
-  if (!gen) throw new Error(`unknown template ${JSON.stringify(id)}`);
-
-  // mkdir -p
-  const { mkdir } = await import("node:fs/promises");
-  await mkdir(outputDir, { recursive: true });
-
-  const rows = await gen(inputs, emit, llm);
-  const n = rows.length;
-  if (n === 0) {
-    throw new Error("generator produced 0 rows — check your inputs");
-  }
-
-  // 90/10 split (matches dataset_templates.py exactly).
-  const splitIdx = Math.max(1, Math.floor(n * 0.9));
-  const trainRows = rows.slice(0, splitIdx);
-  let validRows = rows.slice(splitIdx);
-  if (validRows.length === 0) validRows = rows.slice(-1); // last row fallback
-
-  const sep = outputDir.endsWith("/") ? "" : "/";
-  const trainPath = `${outputDir}${sep}train.jsonl`;
-  const validPath = `${outputDir}${sep}valid.jsonl`;
-
-  await Bun.write(trainPath, trainRows.map((r) => JSON.stringify(r)).join("\n") + "\n");
-  await Bun.write(validPath, validRows.map((r) => JSON.stringify(r)).join("\n") + "\n");
-
-  const nTrain = splitIdx;
-  const nValid = Math.max(1, n - splitIdx);
-
-  emit({
-    type: "stage",
-    stage: "done",
-    progress: 1.0,
-    message: `Wrote ${n} rows.`,
-    output_dir: outputDir,
-    n_train: nTrain,
-    n_valid: nValid,
-  });
-
-  return { n_train: nTrain, n_valid: nValid, output_dir: outputDir };
-}