@simulatte/doppler 0.1.5 → 0.1.7

package/src/training/operator-eval.js

@@ -0,0 +1,230 @@
+import { readFile } from 'node:fs/promises';
+import { resolve } from 'node:path';
+
+import { parseJsonl } from './datasets/jsonl.js';
+
+function asTokenSequence(text) {
+  return String(text ?? '')
+    .trim()
+    .split(/\s+/)
+    .filter(Boolean);
+}
+
+function extractCharacterNgrams(text, n) {
+  const normalized = Array.from(String(text ?? '').trim());
+  if (normalized.length < n) {
+    return new Map();
+  }
+  const grams = new Map();
+  for (let index = 0; index <= normalized.length - n; index += 1) {
+    const gram = normalized.slice(index, index + n).join('');
+    grams.set(gram, (grams.get(gram) || 0) + 1);
+  }
+  return grams;
+}
+
+function countOverlap(source, target) {
+  let overlap = 0;
+  for (const [key, sourceCount] of source.entries()) {
+    const targetCount = target.get(key) || 0;
+    overlap += Math.min(sourceCount, targetCount);
+  }
+  return overlap;
+}
+
+function computeBleuStats(hypotheses, references, maxOrder = 4) {
+  const matchesByOrder = new Array(maxOrder).fill(0);
+  const possibleByOrder = new Array(maxOrder).fill(0);
+  let hypothesisLength = 0;
+  let referenceLength = 0;
+
+  for (let index = 0; index < hypotheses.length; index += 1) {
+    const hypothesis = asTokenSequence(hypotheses[index]);
+    const reference = asTokenSequence(references[index]);
+    hypothesisLength += hypothesis.length;
+    referenceLength += reference.length;
+    for (let order = 1; order <= maxOrder; order += 1) {
+      const hypothesisCounts = new Map();
+      const referenceCounts = new Map();
+      for (let tokenIndex = 0; tokenIndex <= hypothesis.length - order; tokenIndex += 1) {
+        const ngram = hypothesis.slice(tokenIndex, tokenIndex + order).join('\u0001');
+        hypothesisCounts.set(ngram, (hypothesisCounts.get(ngram) || 0) + 1);
+      }
+      for (let tokenIndex = 0; tokenIndex <= reference.length - order; tokenIndex += 1) {
+        const ngram = reference.slice(tokenIndex, tokenIndex + order).join('\u0001');
+        referenceCounts.set(ngram, (referenceCounts.get(ngram) || 0) + 1);
+      }
+      matchesByOrder[order - 1] += countOverlap(hypothesisCounts, referenceCounts);
+      possibleByOrder[order - 1] += Math.max(0, hypothesis.length - order + 1);
+    }
+  }
+
+  return {
+    matchesByOrder,
+    possibleByOrder,
+    hypothesisLength,
+    referenceLength,
+  };
+}
+
+export function computeBleuScore(hypotheses, references, options = {}) {
+  const maxOrder = Number.isInteger(options.maxOrder) && options.maxOrder > 0
+    ? options.maxOrder
+    : 4;
+  if (!Array.isArray(hypotheses) || !Array.isArray(references) || hypotheses.length !== references.length) {
+    throw new Error('computeBleuScore requires equally sized hypothesis and reference arrays.');
+  }
+  if (hypotheses.length === 0) {
+    return {
+      score: 0,
+      brevityPenalty: 0,
+      precisions: new Array(maxOrder).fill(0),
+      hypothesisLength: 0,
+      referenceLength: 0,
+    };
+  }
+
+  const stats = computeBleuStats(hypotheses, references, maxOrder);
+  const precisions = [];
+  let precisionLogSum = 0;
+  for (let order = 0; order < maxOrder; order += 1) {
+    const matches = stats.matchesByOrder[order];
+    const possible = stats.possibleByOrder[order];
+    const precision = possible === 0
+      ? 0
+      : ((matches + 1) / (possible + 1));
+    precisions.push(precision);
+    precisionLogSum += Math.log(Math.max(precision, 1e-16));
+  }
+  const brevityPenalty = stats.hypothesisLength > stats.referenceLength
+    ? 1
+    : Math.exp(1 - (stats.referenceLength / Math.max(stats.hypothesisLength, 1)));
+  const score = brevityPenalty * Math.exp(precisionLogSum / maxOrder);
+  return {
+    score,
+    brevityPenalty,
+    precisions,
+    hypothesisLength: stats.hypothesisLength,
+    referenceLength: stats.referenceLength,
+  };
+}
+
+export function computeChrfScore(hypotheses, references, options = {}) {
+  const maxOrder = Number.isInteger(options.maxOrder) && options.maxOrder > 0
+    ? options.maxOrder
+    : 6;
+  const beta = Number.isFinite(options.beta) && options.beta > 0 ? options.beta : 2;
+  if (!Array.isArray(hypotheses) || !Array.isArray(references) || hypotheses.length !== references.length) {
+    throw new Error('computeChrfScore requires equally sized hypothesis and reference arrays.');
+  }
+  if (hypotheses.length === 0) {
+    return {
+      score: 0,
+      precision: 0,
+      recall: 0,
+    };
+  }
+
+  let precisionSum = 0;
+  let recallSum = 0;
+  for (let order = 1; order <= maxOrder; order += 1) {
+    let overlap = 0;
+    let hypothesisTotal = 0;
+    let referenceTotal = 0;
+    for (let index = 0; index < hypotheses.length; index += 1) {
+      const hypothesisCounts = extractCharacterNgrams(hypotheses[index], order);
+      const referenceCounts = extractCharacterNgrams(references[index], order);
+      overlap += countOverlap(hypothesisCounts, referenceCounts);
+      for (const value of hypothesisCounts.values()) {
+        hypothesisTotal += value;
+      }
+      for (const value of referenceCounts.values()) {
+        referenceTotal += value;
+      }
+    }
+    precisionSum += hypothesisTotal > 0 ? (overlap / hypothesisTotal) : 0;
+    recallSum += referenceTotal > 0 ? (overlap / referenceTotal) : 0;
+  }
+
+  const precision = precisionSum / maxOrder;
+  const recall = recallSum / maxOrder;
+  const betaSquared = beta * beta;
+  const score = (precision + recall) === 0
+    ? 0
+    : ((1 + betaSquared) * precision * recall) / ((betaSquared * precision) + recall);
+  return { score, precision, recall };
+}
+
+export function computeExactMatch(hypotheses, references) {
+  if (!Array.isArray(hypotheses) || !Array.isArray(references) || hypotheses.length !== references.length) {
+    throw new Error('computeExactMatch requires equally sized hypothesis and reference arrays.');
+  }
+  if (hypotheses.length === 0) {
+    return { score: 0, matches: 0, total: 0 };
+  }
+  let matches = 0;
+  for (let index = 0; index < hypotheses.length; index += 1) {
+    if (String(hypotheses[index] ?? '').trim() === String(references[index] ?? '').trim()) {
+      matches += 1;
+    }
+  }
+  return {
+    score: matches / hypotheses.length,
+    matches,
+    total: hypotheses.length,
+  };
+}
+
+export function computeAccuracy(labels, predictions) {
+  return computeExactMatch(predictions, labels);
+}
+
+export function computeEvalMetrics(evalKind, hypotheses, references, options = {}) {
+  const normalizedKind = String(evalKind || '').trim();
+  if (normalizedKind === 'translation') {
+    const bleu = computeBleuScore(hypotheses, references, options.bleu || {});
+    const chrf = computeChrfScore(hypotheses, references, options.chrf || {});
+    return {
+      bleu,
+      chrf,
+      primaryMetric: 'bleu',
+      primaryScore: bleu.score,
+    };
+  }
+  if (normalizedKind === 'text_generation') {
+    const exactMatch = computeExactMatch(hypotheses, references);
+    return {
+      exactMatch,
+      primaryMetric: 'exact_match',
+      primaryScore: exactMatch.score,
+    };
+  }
+  if (normalizedKind === 'classification') {
+    const accuracy = computeAccuracy(references, hypotheses);
+    return {
+      accuracy,
+      primaryMetric: 'accuracy',
+      primaryScore: accuracy.score,
+    };
+  }
+  if (normalizedKind === 'retrieval' || normalizedKind === 'custom') {
+    throw new Error(`Eval kind "${normalizedKind}" requires a custom evaluator and is not yet implemented.`);
+  }
+  throw new Error(`Unsupported eval kind "${normalizedKind}".`);
+}
+
+export async function loadEvalDataset(datasetPath) {
+  const absolutePath = resolve(String(datasetPath));
+  const raw = await readFile(absolutePath, 'utf8');
+  const rows = absolutePath.endsWith('.json')
+    ? JSON.parse(raw)
+    : parseJsonl(raw);
+  if (!Array.isArray(rows)) {
+    throw new Error(`Eval dataset "${absolutePath}" must be a JSON array or JSONL file.`);
+  }
+  return {
+    absolutePath,
+    rows,
+    raw,
+  };
+}

package/src/training/operator-scoreboard.d.ts

@@ -0,0 +1,5 @@
+export declare function appendScoreboardRow(
+  scoreboardDir: string,
+  row: Record<string, unknown>,
+  options?: { selectionMetric?: string | null; selectionGoal?: string | null }
+): Promise<{ rowsPath: string; summaryPath: string }>;

package/src/training/operator-scoreboard.js

@@ -0,0 +1,44 @@
+import { join } from 'node:path';
+
+import { writeJsonArtifact, writeNdjsonRow } from './operator-artifacts.js';
+
+function resolveComparableMetric(row, metric) {
+  if (!row || typeof row !== 'object') return null;
+  const direct = row[metric];
+  if (typeof direct === 'number' && Number.isFinite(direct)) {
+    return direct;
+  }
+  const metrics = row.metrics && typeof row.metrics === 'object' ? row.metrics : null;
+  const nested = metrics?.[metric];
+  if (typeof nested === 'number' && Number.isFinite(nested)) {
+    return nested;
+  }
+  return null;
+}
+
+export async function appendScoreboardRow(scoreboardDir, row, options = {}) {
+  const rowsPath = join(scoreboardDir, 'scoreboard.ndjson');
+  await writeNdjsonRow(rowsPath, row);
+  const metric = String(options.selectionMetric || row.selectionMetric || row.primaryMetric || '').trim();
+  const goal = String(options.selectionGoal || row.selectionGoal || 'max').trim();
+  const comparable = resolveComparableMetric(row, metric);
+  const summary = {
+    artifactType: 'training_scoreboard',
+    schemaVersion: 1,
+    generatedAt: new Date().toISOString(),
+    selectionMetric: metric || null,
+    selectionGoal: goal,
+    latest: row,
+    best: comparable === null
+      ? row
+      : {
+        ...row,
+        selectionMetricValue: comparable,
+      },
+  };
+  const summaryResult = await writeJsonArtifact(join(scoreboardDir, 'latest.json'), summary);
+  return {
+    rowsPath,
+    summaryPath: summaryResult.path,
+  };
+}

package/src/training/optimizer.js

@@ -1,4 +1,4 @@
-import { acquireBuffer, BufferUsage } from '../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer, BufferUsage } from '../memory/buffer-pool.js';
 import { createTensor, tensorBytes } from '../gpu/tensor.js';
 import { runAdam } from '../gpu/kernels/backward/adam.js';
 
@@ -72,12 +72,24 @@ export class AdamOptimizer {
     let entry = this.state.get(param);
     if (!entry) {
       const bytes = tensorBytes(param.shape, param.dtype);
-      const mBuf = acquireBuffer(bytes, BufferUsage.STORAGE, 'adam_m');
-      const vBuf = acquireBuffer(bytes, BufferUsage.STORAGE, 'adam_v');
-      entry = {
-        m: createTensor(mBuf, param.dtype, [...param.shape], 'adam_m'),
-        v: createTensor(vBuf, param.dtype, [...param.shape], 'adam_v'),
-      };
+      let mBuf = null;
+      let vBuf = null;
+      try {
+        mBuf = acquireBuffer(bytes, BufferUsage.STORAGE, 'adam_m');
+        vBuf = acquireBuffer(bytes, BufferUsage.STORAGE, 'adam_v');
+        entry = {
+          m: createTensor(mBuf, param.dtype, [...param.shape], 'adam_m'),
+          v: createTensor(vBuf, param.dtype, [...param.shape], 'adam_v'),
+        };
+      } catch (error) {
+        if (mBuf) {
+          releaseBuffer(mBuf);
+        }
+        if (vBuf) {
+          releaseBuffer(vBuf);
+        }
+        throw error;
+      }
       this.state.set(param, entry);
     }
     return entry;

package/src/training/runner.d.ts

@@ -90,6 +90,16 @@ export interface TrainingStepMetricsEntry {
 export interface TrainingRunnerCallbacks {
   onStep?: (entry: TrainingStepMetricsEntry) => Promise<void> | void;
   onEpoch?: (entry: { epoch: number; steps: number; loss: number }) => Promise<void> | void;
+  onCheckpoint?: (entry: {
+    key: string;
+    defaultCheckpointKey: string | null;
+    path: string | null;
+    metadata: Record<string, unknown> | null;
+    payload: unknown;
+    step: number;
+    epoch: number;
+    batch: number;
+  }) => Promise<void> | void;
 }
 
 export interface TrainingRunnerOptions extends TrainingRunnerCallbacks {
@@ -106,6 +116,12 @@ export interface TrainingRunnerOptions extends TrainingRunnerCallbacks {
   ) => Promise<ClipMetrics>;
   lossScaler?: DynamicLossScaler;
   trainingObjective?: TrainingObjective;
+  resolveCheckpointKey?: (entry: {
+    defaultCheckpointKey: string | null;
+    step: number;
+    epoch: number;
+    batch: number;
+  }) => Promise<string> | string;
 }
 
 export interface TrainingRunOptions {
@@ -159,6 +175,9 @@ export declare class TrainingRunner {
   lastArtifact: UlArtifactFinalizeResult | DistillArtifactFinalizeResult | null;
   lastCheckpoint: {
     key: string;
+    defaultKey?: string | null;
+    path?: string | null;
+    metadata?: Record<string, unknown> | null;
     step: number;
     epoch: number;
     batch: number;
@@ -194,3 +213,36 @@ export declare function runTraining(
   config: TrainingConfigSchema,
   options?: TrainingRunOptions & TrainingRunnerOptions
 ): Promise<TrainingStepMetricsEntry[]>;
+
+export declare function createTrainingCheckpointPayload(
+  model: {
+    loraParams?: () => Tensor[];
+    paramGroups?: () => Record<string, Tensor[]>;
+  },
+  optimizer: unknown,
+  context: {
+    step: number;
+    epoch: number;
+    batch: number;
+    config: TrainingConfigSchema;
+  }
+): Promise<unknown>;
+
+export declare function restoreTrainingCheckpointState(
+  model: {
+    loraParams?: () => Tensor[];
+    paramGroups?: () => Record<string, Tensor[]>;
+  },
+  optimizer: unknown,
+  checkpointRecord: unknown,
+  config: TrainingConfigSchema
+): Promise<{
+  step: number;
+  epoch: number;
+  batch: number;
+  checkpointHash: string | null;
+  previousCheckpointHash: string | null;
+  checkpointKey: string | null;
+  resumeAudits: Array<Record<string, unknown>>;
+  resumeAuditCount: number;
+} | null>;

package/src/training/runner.js

@@ -617,7 +617,6 @@ function buildExpectedCheckpointMetadata(metadata) {
     'configHash',
     'datasetHash',
     'tokenizerHash',
-    'optimizerHash',
     'runtimePresetId',
     'kernelPathId',
   ]) {
@@ -713,7 +712,7 @@ function looksLikeTrainingCheckpointRecord(value) {
   return Number.isInteger(progress.step) && progress.step >= 0;
 }
 
-async function createTrainingCheckpointPayload(model, optimizer, context) {
+export async function createTrainingCheckpointPayload(model, optimizer, context) {
   const freezeMap = context.config?.training?.ul?.freeze
     ?? context.config?.training?.distill?.freeze
     ?? {};
@@ -747,7 +746,7 @@ async function createTrainingCheckpointPayload(model, optimizer, context) {
   };
 }
 
-async function restoreTrainingCheckpointState(model, optimizer, checkpointRecord, config) {
+export async function restoreTrainingCheckpointState(model, optimizer, checkpointRecord, config) {
   if (!looksLikeTrainingCheckpointRecord(checkpointRecord)) {
     return null;
   }
@@ -837,12 +836,16 @@ export class TrainingRunner {
     this.lossScaler = options.lossScaler || new DynamicLossScaler(config.training.lossScaling);
     this.onStep = options.onStep || null;
     this.onEpoch = options.onEpoch || null;
+    this.onCheckpoint = options.onCheckpoint || null;
+    this.resolveCheckpointKey = options.resolveCheckpointKey || null;
     this.lastArtifact = null;
     this.lastCheckpoint = null;
     this.resumeState = null;
   }
 
   async run(model, dataset, options = {}) {
+    this.lastCheckpoint = null;
+    this.lastArtifact = null;
     const {
       epochs = 1,
       batchSize = 1,
@@ -911,16 +914,39 @@ export class TrainingRunner {
         batch: checkpointContext.batch,
         config: this.config,
       });
-      await saveCheckpoint(checkpointKey, payload, {
+      const resolvedCheckpointKey = this.resolveCheckpointKey
+        ? await this.resolveCheckpointKey({
+          defaultCheckpointKey: checkpointKey,
+          step: checkpointContext.step,
+          epoch: checkpointContext.epoch,
+          batch: checkpointContext.batch,
+        })
+        : checkpointKey;
+      const saveResult = await saveCheckpoint(resolvedCheckpointKey, payload, {
         ...checkpointMetadata,
         optimizerHash: hashStableJson(payload?.trainingState?.optimizerSlots || {}),
       });
       this.lastCheckpoint = {
-        key: checkpointKey,
+        key: resolvedCheckpointKey,
+        defaultKey: checkpointKey,
+        path: saveResult?.path || null,
+        metadata: saveResult?.metadata || null,
         step: checkpointContext.step,
         epoch: checkpointContext.epoch,
         batch: checkpointContext.batch,
       };
+      if (this.onCheckpoint) {
+        await this.onCheckpoint({
+          key: resolvedCheckpointKey,
+          defaultCheckpointKey: checkpointKey,
+          path: saveResult?.path || null,
+          metadata: saveResult?.metadata || null,
+          payload,
+          step: checkpointContext.step,
+          epoch: checkpointContext.epoch,
+          batch: checkpointContext.batch,
+        });
+      }
     };
 
     const artifactSession = distillContract.enabled

package/src/training/suite.d.ts

@@ -176,6 +176,66 @@ export interface RunTrainingSuiteOptions {
   timestamp?: string | Date;
 }
 
+export interface DistillDataScope {
+  sourceLangs: string[] | null;
+  targetLangs: string[] | null;
+  pairAllowlist: string[] | null;
+  sourceLangSet: Set<string> | null;
+  targetLangSet: Set<string> | null;
+  pairAllowlistSet: Set<string> | null;
+  strictPairContract: boolean;
+}
+
+export interface DistillDatasetReport {
+  absolutePath: string;
+  rowCount: number;
+  sampleCount: number;
+  directionCounts: Record<string, number>;
+  dataScope: {
+    sourceLangs: string[] | null;
+    targetLangs: string[] | null;
+    pairAllowlist: string[] | null;
+    strictPairContract: boolean;
+  } | null;
+  shardCount?: number;
+  shardPaths?: string[];
+  createDataset(options?: Record<string, unknown>): {
+    batches(): AsyncGenerator<Record<string, unknown>, void, unknown>;
+  };
+}
+
+export interface DistillRuntimeContext {
+  stage: 'stage_a' | 'stage_b';
+  teacherPipeline: Record<string, unknown>;
+  studentPipeline: Record<string, unknown>;
+  teacherModelId: string;
+  studentModelId: string;
+  teacherModelUrl: string | null;
+  studentModelUrl: string | null;
+  topK: number;
+  temperature: number;
+  alphaKd: number;
+  alphaCe: number;
+  tripletMargin: number;
+  studentGraphMode: string;
+  targetTokenMode: string;
+  cleanup(): Promise<void>;
+}
+
+export interface DistillStudentFixture {
+  config: Record<string, unknown>;
+  model: {
+    forward: (input: unknown, tape: unknown) => Promise<unknown>;
+    forwardDistill?: (batch: unknown, tape: unknown, options?: Record<string, unknown>) => Promise<{ logits: unknown }>;
+    cleanupDistillStep?: () => void;
+    loraParams?: () => unknown[];
+    paramGroups?: () => Record<string, unknown[]>;
+  };
+  outputDim?: number;
+  embeddingDim?: number;
+  cleanup(): void;
+}
+
 export declare const trainingHarness: TrainingHarness;
 
 export declare function runTrainingSuite(
@@ -185,3 +245,55 @@ export declare function runTrainingSuite(
 export declare function runTrainingBenchSuite(
   options?: RunTrainingSuiteOptions
 ): Promise<TrainingBenchSuiteResult>;
+
+export declare function resolveDistillDataScope(
+  options?: RunTrainingSuiteOptions,
+  trainingConfig?: Record<string, unknown> | null
+): DistillDataScope;
+
+export declare function buildDistillPrompt(sample: Record<string, unknown>): string;
+
+export declare function normalizeDistillStudentGraphMode(value: unknown): string;
+
+export declare function loadDistillDatasetFromJsonl(
+  datasetPath: string,
+  scopeOptions?: DistillDataScope | null
+): Promise<DistillDatasetReport | null>;
+
+export declare function loadDistillModelHandle(
+  modelRef: string,
+  role: string,
+  loadOptions?: Record<string, unknown>
+): Promise<{
+  modelRef: string;
+  modelUrl: string | null;
+  manifest: Record<string, unknown>;
+  pipeline: Record<string, unknown>;
+}>;
+
+export declare function createDistillRuntimeContext(
+  options?: RunTrainingSuiteOptions,
+  trainingConfig?: Record<string, unknown> | null
+): Promise<DistillRuntimeContext>;
+
+export declare function createToyModelFixture(
+  overrides?: Record<string, unknown>
+): {
+  config: Record<string, unknown>;
+  model: {
+    forward: (input: unknown, tape: unknown) => Promise<unknown>;
+    loraParams(): unknown[];
+    paramGroups(): Record<string, unknown[]>;
+  };
+  batch: Record<string, unknown>;
+  cleanup(): void;
+};
+
+export declare function createDistillStudentRuntimeModelFixture(
+  overrides?: Record<string, unknown>,
+  options?: Record<string, unknown>
+): Promise<DistillStudentFixture>;
+
+export declare function buildDistillTrainingOverrides(
+  options?: RunTrainingSuiteOptions
+): Record<string, unknown> | null;