create-walle 0.9.25 → 0.9.26

package/template/wall-e/training/tool-sft-dataset.js

@@ -0,0 +1,312 @@
+'use strict';
+
+const crypto = require('crypto');
+const fs = require('fs');
+const path = require('path');
+
+const VALID_SPLITS = new Set(['train', 'valid', 'calibration', 'ab_holdout', 'test']);
+
+function stableStringify(value) {
+  if (value === null || typeof value !== 'object') return JSON.stringify(value);
+  if (Array.isArray(value)) return '[' + value.map(stableStringify).join(',') + ']';
+  return '{' + Object.keys(value).sort().map((key) => (
+    JSON.stringify(key) + ':' + stableStringify(value[key])
+  )).join(',') + '}';
+}
+
+function sha256(value) {
+  return crypto.createHash('sha256').update(String(value ?? '')).digest('hex');
+}
+
+function cloneJson(value) {
+  return JSON.parse(JSON.stringify(value));
+}
+
+function toolName(tool) {
+  if (!tool || typeof tool !== 'object') return null;
+  if (tool.type === 'function' && tool.function?.name) return String(tool.function.name);
+  if (tool.name) return String(tool.name);
+  return null;
+}
+
+function toolNameSet(tools = []) {
+  return new Set((tools || []).map(toolName).filter(Boolean));
+}
+
+function messageToolCalls(message = {}) {
+  if (Array.isArray(message.tool_calls)) return message.tool_calls;
+  if (Array.isArray(message.toolCalls)) return message.toolCalls;
+  return [];
+}
+
+function toolCallName(call = {}) {
+  if (call.function?.name) return String(call.function.name);
+  if (call.name) return String(call.name);
+  return null;
+}
+
+function validateToolCall(call, knownTools) {
+  const issues = [];
+  if (!call || typeof call !== 'object') {
+    issues.push('tool_call_not_object');
+    return issues;
+  }
+  const name = toolCallName(call);
+  if (!name) issues.push('tool_call_missing_name');
+  if (name && knownTools.size && !knownTools.has(name)) issues.push(`unknown_tool:${name}`);
+  const args = call.function?.arguments ?? call.arguments ?? call.args;
+  if (args === undefined) {
+    issues.push('tool_call_missing_arguments');
+  } else if (typeof args === 'string') {
+    try { JSON.parse(args); } catch { issues.push('tool_call_arguments_not_json'); }
+  } else if (args === null || typeof args !== 'object' || Array.isArray(args)) {
+    issues.push('tool_call_arguments_invalid');
+  }
+  return issues;
+}
+
+function validateCanonicalExample(example) {
+  const issues = [];
+  if (!example || typeof example !== 'object') return { ok: false, issues: ['example_not_object'] };
+
+  if (!Array.isArray(example.messages) || example.messages.length === 0) {
+    issues.push('missing_messages');
+  }
+  if (!Array.isArray(example.tools) || example.tools.length === 0) {
+    issues.push('missing_tools');
+  }
+  if (!example.tool_catalog || typeof example.tool_catalog !== 'string') {
+    issues.push('missing_tool_catalog');
+  }
+  if (!VALID_SPLITS.has(example.split)) {
+    issues.push('invalid_split');
+  }
+  if (typeof example.quality !== 'number' || example.quality < 0 || example.quality > 1) {
+    issues.push('invalid_quality');
+  }
+  if (!example.leakage || typeof example.leakage !== 'object') {
+    issues.push('missing_leakage');
+  }
+
+  const knownTools = toolNameSet(example.tools || []);
+  for (const [idx, message] of (example.messages || []).entries()) {
+    if (!message || typeof message !== 'object') {
+      issues.push(`message_${idx}_not_object`);
+      continue;
+    }
+    if (!['system', 'user', 'assistant', 'tool'].includes(message.role)) {
+      issues.push(`message_${idx}_invalid_role`);
+    }
+    const calls = messageToolCalls(message);
+    if (message.role === 'assistant' && calls.length > 0) {
+      for (const call of calls) issues.push(...validateToolCall(call, knownTools));
+    }
+    if (message.role === 'tool' && !message.tool_call_id && !message.toolCallId) {
+      issues.push(`message_${idx}_missing_tool_call_id`);
+    }
+    if (message.role !== 'assistant' && calls.length > 0) {
+      issues.push(`message_${idx}_tool_calls_on_${message.role}`);
+    }
+  }
+
+  return { ok: issues.length === 0, issues: [...new Set(issues)] };
+}
+
+function normalizeTool(tool) {
+  const name = toolName(tool);
+  if (!name) return tool;
+  if (tool.type === 'function' && tool.function) {
+    return {
+      type: 'function',
+      function: {
+        name,
+        description: String(tool.function.description || ''),
+        parameters: tool.function.parameters || { type: 'object', properties: {} },
+      },
+    };
+  }
+  return {
+    type: 'function',
+    function: {
+      name,
+      description: String(tool.description || ''),
+      parameters: tool.parameters || tool.input_schema || { type: 'object', properties: {} },
+    },
+  };
+}
+
+function normalizeToolCall(call) {
+  const name = toolCallName(call);
+  const rawArgs = call.function?.arguments ?? call.arguments ?? call.args ?? {};
+  let args = rawArgs;
+  if (typeof args !== 'string') args = JSON.stringify(args || {});
+  return {
+    id: call.id || `call_${sha256(`${name}:${args}`).slice(0, 8)}`,
+    type: 'function',
+    function: { name, arguments: args },
+  };
+}
+
+function normalizeMessage(message) {
+  const out = { role: message.role };
+  if (message.content !== undefined) out.content = String(message.content ?? '');
+  const calls = messageToolCalls(message);
+  if (calls.length) out.tool_calls = calls.map(normalizeToolCall);
+  if (message.tool_call_id || message.toolCallId) out.tool_call_id = message.tool_call_id || message.toolCallId;
+  return out;
+}
+
+function normalizeCanonicalExample(example) {
+  const normalized = cloneJson(example || {});
+  normalized.id = normalized.id || `tooluse-${sha256(stableStringify(normalized)).slice(0, 12)}`;
+  normalized.dataset_version = normalized.dataset_version || normalized.datasetVersion || 'gemma-e4b-tooluse-local';
+  normalized.split = normalized.split || 'train';
+  normalized.source = normalized.source || 'unknown';
+  normalized.task_type = normalized.task_type || normalized.taskType || 'coding-agent';
+  normalized.quality = Number.isFinite(normalized.quality) ? normalized.quality : 0.5;
+  normalized.tool_catalog = normalized.tool_catalog || normalized.toolCatalog || 'gemma-compact-v1';
+  normalized.messages = (normalized.messages || []).map(normalizeMessage);
+  normalized.tools = (normalized.tools || []).map(normalizeTool);
+  normalized.labels = normalized.labels || {};
+  normalized.leakage = {
+    holdout: Boolean(normalized.leakage?.holdout || normalized.split === 'ab_holdout'),
+    source_session_id: normalized.leakage?.source_session_id || normalized.source_session_id || null,
+    source_task_id: normalized.leakage?.source_task_id || normalized.source_task_id || null,
+  };
+  return normalized;
+}
+
+function leakagePayload(example) {
+  const normalized = normalizeCanonicalExample(example);
+  return {
+    source: normalized.source,
+    sourceSessionId: normalized.leakage?.source_session_id || null,
+    sourceTaskId: normalized.leakage?.source_task_id || null,
+    toolCatalog: normalized.tool_catalog,
+    messages: normalized.messages.map((message) => ({
+      role: message.role,
+      content: message.content || '',
+      tool_calls: (message.tool_calls || []).map((call) => ({
+        name: call.function?.name || '',
+        arguments: call.function?.arguments || '',
+      })),
+      tool_call_id: message.tool_call_id || '',
+    })),
+  };
+}
+
+function hashExampleForLeakage(example) {
+  return sha256(stableStringify(leakagePayload(example)));
+}
+
+function groupKeyForExample(example) {
+  const normalized = normalizeCanonicalExample(example);
+  return normalized.leakage?.source_session_id ||
+    normalized.leakage?.source_task_id ||
+    normalized.source_session_id ||
+    normalized.source_task_id ||
+    normalized.id;
+}
+
+function splitForHash(hash, ratios = {}) {
+  const train = ratios.train ?? 0.8;
+  const valid = ratios.valid ?? 0.1;
+  const calibration = ratios.calibration ?? 0.1;
+  const n = parseInt(hash.slice(0, 8), 16) / 0xffffffff;
+  if (n < train) return 'train';
+  if (n < train + valid) return 'valid';
+  if (n < train + valid + calibration) return 'calibration';
+  return 'ab_holdout';
+}
+
+function splitExamplesBySource(examples = [], ratios = {}, seed = 'tooluse-v1') {
+  const out = { train: [], valid: [], calibration: [], ab_holdout: [] };
+  const groupToSplit = new Map();
+  for (const input of examples) {
+    const example = normalizeCanonicalExample(input);
+    const groupKey = groupKeyForExample(example);
+    if (!groupToSplit.has(groupKey)) {
+      groupToSplit.set(groupKey, splitForHash(sha256(`${seed}:${groupKey}`), ratios));
+    }
+    const split = groupToSplit.get(groupKey);
+    example.split = split;
+    example.leakage.holdout = split === 'ab_holdout';
+    out[split].push(example);
+  }
+  return out;
+}
+
+function asArray(value) {
+  if (!value) return [];
+  if (Array.isArray(value)) return value;
+  return Object.values(value).flat();
+}
+
+function checkHoldoutLeakage({ train = [], valid = [], calibration = [], holdout = [], ab_holdout = [] } = {}) {
+  const holdoutRows = [...asArray(holdout), ...asArray(ab_holdout)];
+  const holdoutHashes = new Set(holdoutRows.map(hashExampleForLeakage));
+  const leaked = [];
+  for (const [split, rows] of Object.entries({ train, valid, calibration })) {
+    for (const row of asArray(rows)) {
+      const hash = hashExampleForLeakage(row);
+      if (holdoutHashes.has(hash)) leaked.push({ split, id: row.id || null, hash });
+    }
+  }
+  return { ok: leaked.length === 0, leaked };
+}
+
+function buildDatasetManifest({ examples = [], splits = null, config = {} } = {}) {
+  const all = splits ? asArray(splits) : examples.map(normalizeCanonicalExample);
+  const bySplit = {};
+  const bySource = {};
+  for (const example of all) {
+    bySplit[example.split] = (bySplit[example.split] || 0) + 1;
+    bySource[example.source] = (bySource[example.source] || 0) + 1;
+  }
+  const datasetHash = sha256(stableStringify(all.map((example) => ({
+    id: example.id,
+    split: example.split,
+    leakageHash: hashExampleForLeakage(example),
+  }))));
+  return {
+    datasetVersion: config.datasetVersion || 'gemma-e4b-tooluse-v1',
+    datasetHash,
+    createdAt: config.createdAt || new Date().toISOString(),
+    total: all.length,
+    bySplit,
+    bySource,
+    config: { ...config, createdAt: undefined },
+  };
+}
+
+function writeJsonl(filePath, rows) {
+  const lines = rows.map((row) => JSON.stringify(row));
+  fs.writeFileSync(filePath, lines.join('\n') + (lines.length ? '\n' : ''), 'utf8');
+}
+
+function writeDatasetFiles({ outputDir, examples = [], splits = null, manifest = null } = {}) {
+  if (!outputDir) throw new Error('outputDir is required');
+  fs.mkdirSync(outputDir, { recursive: true });
+  const splitMap = splits || splitExamplesBySource(examples);
+  const all = asArray(splitMap).map(normalizeCanonicalExample);
+  writeJsonl(path.join(outputDir, 'canonical.jsonl'), all);
+  for (const [split, rows] of Object.entries(splitMap)) {
+    writeJsonl(path.join(outputDir, `${split}.jsonl`), rows.map(normalizeCanonicalExample));
+  }
+  const finalManifest = manifest || buildDatasetManifest({ splits: splitMap });
+  fs.writeFileSync(path.join(outputDir, 'manifest.json'), JSON.stringify(finalManifest, null, 2) + '\n', 'utf8');
+  return { outputDir, manifest: finalManifest, splits: splitMap };
+}
+
+module.exports = {
+  VALID_SPLITS,
+  stableStringify,
+  sha256,
+  validateCanonicalExample,
+  normalizeCanonicalExample,
+  hashExampleForLeakage,
+  splitExamplesBySource,
+  checkHoldoutLeakage,
+  buildDatasetManifest,
+  writeDatasetFiles,
+};

package/template/wall-e/training/tool-sft-renderers.js

@@ -0,0 +1,144 @@
+'use strict';
+
+const fs = require('fs');
+const path = require('path');
+const { normalizeCanonicalExample } = require('./tool-sft-dataset');
+const { prepareMessagesForMlxProvider } = require('../llm/mlx');
+
+const RENDER_FORMATS = new Set(['mlx-tools', 'mlx-provider', 'mlx-provider-bare', 'mlx-provider-state-bare', 'trl-tools', 'text']);
+
+function cloneJson(value) {
+  return JSON.parse(JSON.stringify(value));
+}
+
+function compactTool(tool) {
+  const out = cloneJson(tool);
+  if (out.function) {
+    out.function.description = firstSentence(out.function.description || '');
+  } else if (out.description) {
+    out.description = firstSentence(out.description);
+  }
+  return out;
+}
+
+function firstSentence(text) {
+  const normalized = String(text || '').replace(/\s+/g, ' ').trim();
+  if (!normalized) return '';
+  const idx = normalized.search(/[.!?](\s|$)/);
+  return idx >= 0 ? normalized.slice(0, idx + 1) : normalized;
+}
+
+function normalizeArgumentsForBackend(messages, { objectArguments = false } = {}) {
+  return (messages || []).map((message) => {
+    const out = cloneJson(message);
+    if (!Array.isArray(out.tool_calls)) return out;
+    out.tool_calls = out.tool_calls.map((call) => {
+      const next = cloneJson(call);
+      const raw = next.function?.arguments;
+      if (objectArguments && typeof raw === 'string') {
+        try { next.function.arguments = JSON.parse(raw); } catch { /* leave string */ }
+      } else if (!objectArguments && raw !== undefined && typeof raw !== 'string') {
+        next.function.arguments = JSON.stringify(raw || {});
+      }
+      return next;
+    });
+    return out;
+  });
+}
+
+function renderToolResultAsMlxProviderMessage(message = {}) {
+  if (message.role !== 'tool') return message;
+  const id = message.tool_call_id || message.toolCallId || 'tool';
+  return {
+    role: 'user',
+    content: `[Tool result for ${id}]: ${message.content || ''}`,
+  };
+}
+
+function renderMlxToolsExample(example, opts = {}) {
+  const normalized = normalizeCanonicalExample(example);
+  return {
+    messages: normalizeArgumentsForBackend(
+      normalized.messages.map(renderToolResultAsMlxProviderMessage),
+      { objectArguments: !!opts.objectArguments }
+    ),
+    tools: opts.compactTools ? normalized.tools.map(compactTool) : cloneJson(normalized.tools),
+  };
+}
+
+function renderMlxProviderExample(example, opts = {}) {
+  const normalized = normalizeCanonicalExample(example);
+  const tools = opts.compactTools ? normalized.tools.map(compactTool) : normalized.tools;
+  return {
+    messages: prepareMessagesForMlxProvider({
+      messages: normalized.messages,
+      tools,
+      compactDuplicates: true,
+      toolCallFormat: opts.toolCallFormat || 'fenced',
+    }),
+  };
+}
+
+function renderTrlToolsExample(example, opts = {}) {
+  const normalized = normalizeCanonicalExample(example);
+  return {
+    messages: normalizeArgumentsForBackend(normalized.messages, { objectArguments: !!opts.objectArguments }),
+    tools: opts.compactTools ? normalized.tools.map(compactTool) : cloneJson(normalized.tools),
+  };
+}
+
+function renderTextFallbackExample(example) {
+  const normalized = normalizeCanonicalExample(example);
+  const lines = [];
+  for (const message of normalized.messages) {
+    if (message.role === 'assistant' && Array.isArray(message.tool_calls) && message.tool_calls.length) {
+      lines.push('assistant tool_calls:');
+      for (const call of message.tool_calls) {
+        lines.push(`- ${call.function?.name || ''} ${call.function?.arguments || '{}'}`);
+      }
+      continue;
+    }
+    if (message.role === 'tool') {
+      lines.push(`tool ${message.tool_call_id || ''}: ${message.content || ''}`);
+      continue;
+    }
+    lines.push(`${message.role}: ${message.content || ''}`);
+  }
+  return { text: lines.join('\n') };
+}
+
+function renderExample(example, format, opts = {}) {
+  if (format === 'mlx-tools') return renderMlxToolsExample(example, opts);
+  if (format === 'mlx-provider') return renderMlxProviderExample(example, opts);
+  if (format === 'mlx-provider-bare') return renderMlxProviderExample(example, { ...opts, toolCallFormat: 'bare-json' });
+  if (format === 'mlx-provider-state-bare') return renderMlxProviderExample(example, { ...opts, toolCallFormat: 'bare-json' });
+  if (format === 'trl-tools') return renderTrlToolsExample(example, opts);
+  if (format === 'text') return renderTextFallbackExample(example, opts);
+  throw new Error(`Unsupported render format: ${format}`);
+}
+
+function writeJsonl(filePath, rows) {
+  const lines = rows.map((row) => JSON.stringify(row));
+  fs.writeFileSync(filePath, lines.join('\n') + (lines.length ? '\n' : ''), 'utf8');
+}
+
+function writeRenderedSplit({ examples = [], outputDir, format = 'mlx-tools', splitName = 'train', compactTools = true } = {}) {
+  if (!RENDER_FORMATS.has(format)) throw new Error(`Unsupported render format: ${format}`);
+  if (!outputDir) throw new Error('outputDir is required');
+  fs.mkdirSync(outputDir, { recursive: true });
+  const filePath = path.join(outputDir, `${splitName}.jsonl`);
+  const rows = examples.map((example) => renderExample(example, format, { compactTools }));
+  writeJsonl(filePath, rows);
+  return { filePath, count: rows.length };
+}
+
+module.exports = {
+  RENDER_FORMATS,
+  renderMlxToolsExample,
+  renderMlxProviderExample,
+  renderTrlToolsExample,
+  renderExample,
+  renderTextFallbackExample,
+  writeRenderedSplit,
+  _private: { compactTool, normalizeArgumentsForBackend, renderToolResultAsMlxProviderMessage },
+};