@simulatte/doppler 0.1.0

package/src/training/clip.js

@@ -0,0 +1,54 @@
+
+import { readBuffer } from '../memory/buffer-pool.js';
+import { runScale } from '../gpu/kernels/index.js';
+import { f16ToF32Array } from '../inference/kv-cache/types.js';
+
+async function readGradData(grad) {
+  const buffer = await readBuffer(grad.buffer);
+  if (grad.dtype === 'f16') {
+    return f16ToF32Array(new Uint16Array(buffer));
+  }
+  return new Float32Array(buffer);
+}
+
+export async function clipGradients(grads, config) {
+  const maxNorm = config?.training?.gradient?.maxNorm;
+  let sumSq = 0;
+  let totalParamCount = 0;
+
+  for (const grad of grads.values()) {
+    const data = await readGradData(grad);
+    totalParamCount += data.length;
+    for (let i = 0; i < data.length; i += 1) {
+      const value = data[i];
+      sumSq += value * value;
+    }
+  }
+
+  const globalNorm = Math.sqrt(sumSq);
+
+  if (!maxNorm || maxNorm <= 0 || !globalNorm || globalNorm <= maxNorm) {
+    return {
+      clippedGrads: grads,
+      gradient_norm_unclipped: globalNorm,
+      gradient_norm_clipped: globalNorm,
+      clipped_event_count: 0,
+      total_param_count: totalParamCount,
+    };
+  }
+
+  const scale = maxNorm / (globalNorm + 1e-6);
+  const clipped = new Map();
+  for (const [param, grad] of grads.entries()) {
+    const scaled = await runScale(grad, scale, { inplace: true });
+    clipped.set(param, scaled);
+  }
+
+  return {
+    clippedGrads: clipped,
+    gradient_norm_unclipped: globalNorm,
+    gradient_norm_clipped: maxNorm,
+    clipped_event_count: 1,
+    total_param_count: totalParamCount,
+  };
+}

package/src/training/dataloader.d.ts

@@ -0,0 +1,8 @@
+export declare class DataLoader<T> {
+  constructor(dataset: T[], batchSize: number, shuffle?: boolean);
+  dataset: T[];
+  batchSize: number;
+  shuffle: boolean;
+  collate(batch: T[]): T[] | unknown;
+  batches(): AsyncGenerator<T[] | unknown, void, void>;
+}

package/src/training/dataloader.js

@@ -0,0 +1,44 @@
+let fallbackRandomState = (Date.now() >>> 0) || 0x6d2b79f5;
+
+function unseededRandom() {
+  if (typeof crypto !== 'undefined' && typeof crypto.getRandomValues === 'function') {
+    const values = new Uint32Array(1);
+    crypto.getRandomValues(values);
+    return values[0] / 4294967296;
+  }
+  fallbackRandomState = (fallbackRandomState + 0x6d2b79f5) >>> 0;
+  return fallbackRandomState / 4294967296;
+}
+
+function shuffledIndices(length) {
+  const indices = Array.from({ length }, (_, i) => i);
+  for (let i = indices.length - 1; i > 0; i -= 1) {
+    const j = Math.floor(unseededRandom() * (i + 1));
+    const tmp = indices[i];
+    indices[i] = indices[j];
+    indices[j] = tmp;
+  }
+  return indices;
+}
+
+export class DataLoader {
+  constructor(dataset, batchSize, shuffle = true) {
+    this.dataset = dataset;
+    this.batchSize = batchSize;
+    this.shuffle = shuffle;
+  }
+
+  collate(batch) {
+    return batch;
+  }
+
+  async *batches() {
+    const length = this.dataset.length ?? 0;
+    const indices = this.shuffle ? shuffledIndices(length) : Array.from({ length }, (_, i) => i);
+    for (let i = 0; i < indices.length; i += this.batchSize) {
+      const batchIndices = indices.slice(i, i + this.batchSize);
+      const batch = batchIndices.map((idx) => this.dataset[idx]);
+      yield this.collate(batch);
+    }
+  }
+}

package/src/training/datasets/index.d.ts

@@ -0,0 +1,12 @@
+export { parseJsonl, loadJsonl, mapJsonl } from './jsonl.js';
+export { tokenizeTextPairs, buildCausalPair, type TextPair, type TokenizedSample } from './text-pairs.js';
+export {
+  mapTranslationPairs,
+  tokenizeTranslationPairs,
+  type TranslationPair,
+  type TranslationTokenizedSample,
+  type MapTranslationPairsOptions,
+  type TokenizeTranslationPairsOptions,
+} from './translation-pairs.js';
+export { buildTokenBatch, createTokenBatchTensors, type TokenBatch } from './token-batch.js';
+export { reploidTracesToTextPairs } from './reploid.js';

package/src/training/datasets/index.js

@@ -0,0 +1,6 @@
+
+export { parseJsonl, loadJsonl, mapJsonl } from './jsonl.js';
+export { tokenizeTextPairs, buildCausalPair } from './text-pairs.js';
+export { mapTranslationPairs, tokenizeTranslationPairs } from './translation-pairs.js';
+export { buildTokenBatch, createTokenBatchTensors } from './token-batch.js';
+export { reploidTracesToTextPairs } from './reploid.js';

package/src/training/datasets/jsonl.d.ts

@@ -0,0 +1,11 @@
+export declare function parseJsonl(text: string): unknown[];
+
+export declare function loadJsonl(
+  source: string,
+  options?: { fetch?: (url: string) => Promise<string> }
+): Promise<unknown[]>;
+
+export declare function mapJsonl<T, U>(
+  records: T[],
+  mapper?: (record: T) => U | null | undefined
+): U[];

package/src/training/datasets/jsonl.js

@@ -0,0 +1,50 @@
+
+function parseLine(line, lineNumber) {
+  if (!line.trim()) {
+    return null;
+  }
+  try {
+    return JSON.parse(line);
+  } catch (err) {
+    const message = err instanceof Error ? err.message : String(err);
+    throw new Error(`Invalid JSONL at line ${lineNumber}: ${message}`);
+  }
+}
+
+export function parseJsonl(text) {
+  const records = [];
+  const lines = text.split(/\r?\n/);
+  for (let i = 0; i < lines.length; i += 1) {
+    const value = parseLine(lines[i], i + 1);
+    if (value !== null) {
+      records.push(value);
+    }
+  }
+  return records;
+}
+
+export async function loadJsonl(source, options = {}) {
+  if (typeof source !== 'string') {
+    throw new Error('loadJsonl expects a string source');
+  }
+
+  const looksLikeUrl = /^https?:\/\//i.test(source);
+  const text = looksLikeUrl
+    ? await (options.fetch ? options.fetch(source) : fetch(source).then((res) => res.text()))
+    : source;
+  return parseJsonl(text);
+}
+
+export function mapJsonl(records, mapper) {
+  if (!mapper) {
+    return records;
+  }
+  const mapped = [];
+  for (const record of records) {
+    const next = mapper(record);
+    if (next !== null && next !== undefined) {
+      mapped.push(next);
+    }
+  }
+  return mapped;
+}

package/src/training/datasets/reploid.d.ts

@@ -0,0 +1,3 @@
+import type { TextPair } from './text-pairs.js';
+
+export declare function reploidTracesToTextPairs(traces: Array<Record<string, unknown>>): TextPair[];

package/src/training/datasets/reploid.js

@@ -0,0 +1,36 @@
+
+function formatToolCall(tool, args, result) {
+  const argText = args ? JSON.stringify(args) : '';
+  const resultText = result ? JSON.stringify(result) : '';
+  return `tool:${tool}\nargs:${argText}\nresult:${resultText}`;
+}
+
+function extractPair(event) {
+  if (event.prompt && event.completion) {
+    return { prompt: event.prompt, completion: event.completion };
+  }
+  if (event.query && event.response) {
+    return { prompt: event.query, completion: event.response };
+  }
+  if (event.type === 'tool:execute' && event.payload) {
+    const tool = event.payload.tool || 'tool';
+    const args = event.payload.args || null;
+    const result = event.payload.result || event.payload.output || null;
+    return {
+      prompt: event.payload.prompt || event.payload.query || `call ${tool}`,
+      completion: formatToolCall(tool, args, result),
+    };
+  }
+  return null;
+}
+
+export function reploidTracesToTextPairs(traces) {
+  const pairs = [];
+  for (const event of traces) {
+    const pair = extractPair(event);
+    if (pair) {
+      pairs.push(pair);
+    }
+  }
+  return pairs;
+}

package/src/training/datasets/text-pairs.d.ts

@@ -0,0 +1,21 @@
+export interface TextPair {
+  prompt: string;
+  completion: string;
+}
+
+export interface TokenizedSample {
+  inputIds: number[];
+  targetIds: number[];
+  text?: string;
+}
+
+export declare function buildCausalPair(tokens: number[]): {
+  inputIds: number[];
+  targetIds: number[];
+};
+
+export declare function tokenizeTextPairs(
+  tokenizer: { encode: (text: string) => number[] },
+  pairs: TextPair[],
+  options?: { maxLength?: number | null; joinWith?: string }
+): Promise<TokenizedSample[]>;

package/src/training/datasets/text-pairs.js

@@ -0,0 +1,42 @@
+
+function concatTokens(tokens, maxLength) {
+  if (!maxLength || tokens.length <= maxLength) {
+    return tokens;
+  }
+  return tokens.slice(0, maxLength);
+}
+
+export function buildCausalPair(tokens) {
+  if (tokens.length < 2) {
+    return { inputIds: [], targetIds: [] };
+  }
+  return {
+    inputIds: tokens.slice(0, tokens.length - 1),
+    targetIds: tokens.slice(1),
+  };
+}
+
+export async function tokenizeTextPairs(tokenizer, pairs, options = {}) {
+  if (!tokenizer || typeof tokenizer.encode !== 'function') {
+    throw new Error('tokenizeTextPairs requires a tokenizer with encode()');
+  }
+
+  const {
+    maxLength = null,
+    joinWith = '',
+  } = options;
+
+  const samples = [];
+  for (const pair of pairs) {
+    const prompt = pair.prompt ?? '';
+    const completion = pair.completion ?? '';
+    const fullText = `${prompt}${joinWith}${completion}`;
+    const tokens = tokenizer.encode(fullText);
+    const clipped = concatTokens(tokens, maxLength);
+    const { inputIds, targetIds } = buildCausalPair(clipped);
+    if (inputIds.length > 0) {
+      samples.push({ inputIds, targetIds, text: fullText });
+    }
+  }
+  return samples;
+}

package/src/training/datasets/token-batch.d.ts

@@ -0,0 +1,21 @@
+import type { Tensor } from '../../gpu/tensor.js';
+
+export interface TokenizedSample {
+  inputIds: number[];
+  targetIds: number[];
+  text?: string;
+}
+
+export interface TokenBatch {
+  inputFlat: Uint32Array;
+  targetFlat: Uint32Array;
+  offsets: number[];
+}
+
+export declare function buildTokenBatch(samples: TokenizedSample[]): TokenBatch;
+
+export declare function createTokenBatchTensors(batch: TokenBatch): {
+  input: Tensor;
+  targets: Tensor;
+  offsets: number[];
+};

package/src/training/datasets/token-batch.js

@@ -0,0 +1,40 @@
+
+import { acquireBuffer, uploadData } from '../../memory/buffer-pool.js';
+import { createTensor } from '../../gpu/tensor.js';
+
+function flattenTokenBatch(samples, key) {
+  const offsets = [];
+  let total = 0;
+  for (const sample of samples) {
+    offsets.push(total);
+    total += sample[key].length;
+  }
+  const flat = new Uint32Array(total);
+  for (let i = 0; i < samples.length; i += 1) {
+    const ids = samples[i][key];
+    flat.set(ids, offsets[i]);
+  }
+  return { flat, offsets };
+}
+
+export function buildTokenBatch(samples) {
+  if (!samples.length) {
+    throw new Error('buildTokenBatch requires at least one sample');
+  }
+  const { flat: inputFlat, offsets } = flattenTokenBatch(samples, 'inputIds');
+  const { flat: targetFlat } = flattenTokenBatch(samples, 'targetIds');
+  return { inputFlat, targetFlat, offsets };
+}
+
+export function createTokenBatchTensors(batch) {
+  const inputBuf = acquireBuffer(batch.inputFlat.byteLength, undefined, 'train_input_tokens');
+  uploadData(inputBuf, batch.inputFlat);
+
+  const targetBuf = acquireBuffer(batch.targetFlat.byteLength, undefined, 'train_target_tokens');
+  uploadData(targetBuf, batch.targetFlat);
+
+  const input = createTensor(inputBuf, 'f32', [batch.inputFlat.length], 'train_input_tokens');
+  const targets = createTensor(targetBuf, 'f32', [batch.targetFlat.length], 'train_target_tokens');
+
+  return { input, targets, offsets: batch.offsets };
+}

package/src/training/datasets/translation-pairs.d.ts

@@ -0,0 +1,34 @@
+export interface TranslationPair {
+  source: string;
+  target: string;
+}
+
+export interface TranslationTokenizedSample {
+  inputIds: number[];
+  targetIds: number[];
+  source: string;
+  target: string;
+  text: string;
+}
+
+export interface MapTranslationPairsOptions {
+  sourceKey?: string;
+  targetKey?: string;
+}
+
+export interface TokenizeTranslationPairsOptions {
+  maxLength?: number | null;
+  promptPrefix?: string;
+  separator?: string;
+}
+
+export declare function mapTranslationPairs(
+  records: unknown[],
+  options?: MapTranslationPairsOptions
+): TranslationPair[];
+
+export declare function tokenizeTranslationPairs(
+  tokenizer: { encode: (text: string) => number[] },
+  pairs: TranslationPair[],
+  options?: TokenizeTranslationPairsOptions
+): Promise<TranslationTokenizedSample[]>;

package/src/training/datasets/translation-pairs.js

@@ -0,0 +1,49 @@
+function normalizePair(record, sourceKey, targetKey) {
+  if (!record || typeof record !== 'object') return null;
+  const source = String(record[sourceKey] ?? '').trim();
+  const target = String(record[targetKey] ?? '').trim();
+  if (!source || !target) return null;
+  return { source, target };
+}
+
+export function mapTranslationPairs(records, options = {}) {
+  const sourceKey = String(options.sourceKey || 'source');
+  const targetKey = String(options.targetKey || 'target');
+  if (!Array.isArray(records)) return [];
+  const pairs = [];
+  for (const record of records) {
+    const pair = normalizePair(record, sourceKey, targetKey);
+    if (pair) pairs.push(pair);
+  }
+  return pairs;
+}
+
+export async function tokenizeTranslationPairs(tokenizer, pairs, options = {}) {
+  if (!tokenizer || typeof tokenizer.encode !== 'function') {
+    throw new Error('tokenizeTranslationPairs requires a tokenizer with encode()');
+  }
+  const maxLength = Number.isFinite(options.maxLength)
+    ? Math.max(2, Math.floor(options.maxLength))
+    : null;
+  const promptPrefix = String(options.promptPrefix || 'Translate to target: ');
+  const separator = String(options.separator || '\n');
+
+  const samples = [];
+  for (const pair of pairs) {
+    const source = String(pair?.source || '').trim();
+    const target = String(pair?.target || '').trim();
+    if (!source || !target) continue;
+    const text = `${promptPrefix}${source}${separator}${target}`;
+    const tokens = tokenizer.encode(text);
+    const clipped = maxLength ? tokens.slice(0, maxLength) : tokens;
+    if (clipped.length < 2) continue;
+    samples.push({
+      inputIds: clipped.slice(0, clipped.length - 1),
+      targetIds: clipped.slice(1),
+      source,
+      target,
+      text,
+    });
+  }
+  return samples;
+}

package/src/training/export.d.ts

@@ -0,0 +1,32 @@
+import type { AdapterManifest } from '../adapters/adapter-manifest.js';
+import type { LoRAModuleName } from '../inference/pipelines/text/lora-types.js';
+import type { Tensor } from '../gpu/tensor.js';
+
+export interface LoRAExportTensor {
+  name: string;
+  tensor: Tensor | Float32Array | GPUBuffer;
+  shape?: [number, number];
+  dtype?: 'f16' | 'f32';
+}
+
+export interface LoRAExportOptions {
+  id: string;
+  name: string;
+  baseModel: string;
+  rank: number;
+  alpha: number;
+  targetModules: LoRAModuleName[];
+  version?: string;
+  description?: string;
+  metadata?: Record<string, unknown>;
+  tensors: LoRAExportTensor[];
+  format?: 'base64' | 'array';
+  pretty?: boolean;
+}
+
+export interface LoRAExportResult {
+  manifest: AdapterManifest;
+  json: string;
+}
+
+export declare function exportLoRAAdapter(options: LoRAExportOptions): Promise<LoRAExportResult>;

package/src/training/export.js

@@ -0,0 +1,112 @@
+
+import { readBuffer } from '../memory/buffer-pool.js';
+import { f16ToF32Array } from '../inference/kv-cache/types.js';
+import { createManifest, serializeManifest } from '../adapters/adapter-manifest.js';
+
+function encodeBase64(bytes) {
+  if (typeof btoa === 'function') {
+    let binary = '';
+    const view = new Uint8Array(bytes);
+    for (let i = 0; i < view.length; i += 1) {
+      binary += String.fromCharCode(view[i]);
+    }
+    return btoa(binary);
+  }
+  if (typeof Buffer !== 'undefined') {
+    return Buffer.from(bytes).toString('base64');
+  }
+  throw new Error('Base64 encoding not supported in this environment');
+}
+
+async function resolveTensorData(entry) {
+  const dtype = entry.dtype ?? entry.tensor?.dtype ?? 'f32';
+  let shape = entry.shape ?? entry.tensor?.shape;
+  if (!shape) {
+    throw new Error(`Missing shape for tensor ${entry.name}`);
+  }
+
+  const hasGPUBuffer = typeof GPUBuffer !== 'undefined';
+  let data;
+  if (entry.tensor instanceof Float32Array) {
+    data = entry.tensor;
+  } else if (hasGPUBuffer && entry.tensor?.buffer instanceof GPUBuffer) {
+    const raw = await readBuffer(entry.tensor.buffer);
+    data = dtype === 'f16'
+      ? f16ToF32Array(new Uint16Array(raw))
+      : new Float32Array(raw);
+  } else if (hasGPUBuffer && entry.tensor instanceof GPUBuffer) {
+    const raw = await readBuffer(entry.tensor);
+    data = dtype === 'f16'
+      ? f16ToF32Array(new Uint16Array(raw))
+      : new Float32Array(raw);
+  } else {
+    throw new Error(`Unsupported tensor type for ${entry.name}`);
+  }
+
+  return { dtype: 'f32', shape: [...shape], data };
+}
+
+export async function exportLoRAAdapter(options) {
+  const {
+    id,
+    name,
+    baseModel,
+    rank,
+    alpha,
+    targetModules,
+    version,
+    description,
+    metadata,
+    tensors,
+    format = 'base64',
+    pretty = false,
+  } = options;
+
+  if (!Array.isArray(tensors) || tensors.length === 0) {
+    throw new Error('exportLoRAAdapter requires tensors');
+  }
+
+  const manifest = createManifest({
+    id,
+    name,
+    baseModel,
+    rank,
+    alpha,
+    targetModules,
+    version,
+    description,
+    metadata,
+    weightsFormat: 'json',
+  });
+
+  const serialized = [];
+  let totalBytes = 0;
+
+  for (const entry of tensors) {
+    const resolved = await resolveTensorData(entry);
+    const byteLength = resolved.data.byteLength;
+    totalBytes += byteLength;
+    const tensorSpec = {
+      name: entry.name,
+      shape: resolved.shape,
+      dtype: resolved.dtype,
+    };
+
+    if (format === 'array') {
+      tensorSpec.data = Array.from(resolved.data);
+    } else {
+      const slice = resolved.data.buffer.slice(
+        resolved.data.byteOffset,
+        resolved.data.byteOffset + resolved.data.byteLength
+      );
+      tensorSpec.base64 = encodeBase64(slice);
+    }
+
+    serialized.push(tensorSpec);
+  }
+
+  manifest.tensors = serialized;
+  manifest.weightsSize = totalBytes;
+
+  return { manifest, json: serializeManifest(manifest, pretty) };
+}

package/src/training/index.d.ts

@@ -0,0 +1,52 @@
+export { AutogradTape, OpType } from './autograd.js';
+export { buildAttentionSoftmaxCache } from './attention-backward.js';
+export { recordAttentionForward } from './attention-forward.js';
+export { LoraAdapter } from './lora.js';
+export { AdamOptimizer } from './optimizer.js';
+export { trainStep } from './trainer.js';
+export { crossEntropyLoss } from './loss.js';
+export { clipGradients } from './clip.js';
+export { exportLoRAAdapter } from './export.js';
+export { DynamicLossScaler, detectOverflow } from './loss-scaling.js';
+export { TrainingRunner, runTraining } from './runner.js';
+export { runTrainingSuite, runTrainingBenchSuite, trainingHarness } from './suite.js';
+export type {
+  TrainingObjective,
+  TrainingObjectiveContext,
+  TrainingObjectiveLossResult,
+} from './objectives/base.js';
+export {
+  createTrainingObjective,
+  isTrainingObjective,
+  createCrossEntropyObjective,
+  CROSS_ENTROPY_OBJECTIVE,
+  createDistillKdObjective,
+  createDistillTripletObjective,
+  createUlStage1JointObjective,
+  createUlStage2BaseObjective,
+} from './objectives/index.js';
+export type {
+  UlArtifactSession,
+  UlArtifactFinalizeResult,
+  DistillArtifactSession,
+  DistillArtifactFinalizeResult,
+} from './artifacts.js';
+export {
+  createDistillArtifactSession,
+  resolveDistillTrainingContract,
+  resolveStageAArtifactContext,
+  createUlArtifactSession,
+  resolveUlTrainingContract,
+  resolveStage1ArtifactContext,
+} from './artifacts.js';
+export {
+  resolveUlNoiseScale,
+  resolveUlScheduledLambda,
+  buildNoisyLatentsFromInputTensor,
+  applyUlStage1Batch,
+  cleanupUlPreparedBatch,
+  computeLatentBitrateProxy,
+} from './ul_dataset.js';
+export * as datasets from './datasets/index.js';
+export { DataLoader } from './dataloader.js';
+export { saveCheckpoint, loadCheckpoint } from './checkpoint.js';

package/src/training/index.js

@@ -0,0 +1,41 @@
+export { AutogradTape, OpType } from './autograd.js';
+export { buildAttentionSoftmaxCache } from './attention-backward.js';
+export { recordAttentionForward } from './attention-forward.js';
+export { LoraAdapter } from './lora.js';
+export { AdamOptimizer } from './optimizer.js';
+export { trainStep } from './trainer.js';
+export { crossEntropyLoss } from './loss.js';
+export { clipGradients } from './clip.js';
+export { exportLoRAAdapter } from './export.js';
+export { DynamicLossScaler, detectOverflow } from './loss-scaling.js';
+export { TrainingRunner, runTraining } from './runner.js';
+export { runTrainingSuite, runTrainingBenchSuite, trainingHarness } from './suite.js';
+export {
+  createTrainingObjective,
+  isTrainingObjective,
+  createCrossEntropyObjective,
+  CROSS_ENTROPY_OBJECTIVE,
+  createDistillKdObjective,
+  createDistillTripletObjective,
+  createUlStage1JointObjective,
+  createUlStage2BaseObjective,
+} from './objectives/index.js';
+export {
+  createDistillArtifactSession,
+  resolveDistillTrainingContract,
+  resolveStageAArtifactContext,
+  createUlArtifactSession,
+  resolveUlTrainingContract,
+  resolveStage1ArtifactContext,
+} from './artifacts.js';
+export {
+  resolveUlNoiseScale,
+  resolveUlScheduledLambda,
+  buildNoisyLatentsFromInputTensor,
+  applyUlStage1Batch,
+  cleanupUlPreparedBatch,
+  computeLatentBitrateProxy,
+} from './ul_dataset.js';
+export * as datasets from './datasets/index.js';
+export { DataLoader } from './dataloader.js';
+export { saveCheckpoint, loadCheckpoint } from './checkpoint.js';