@mlx-node/trl 0.0.0 → 0.0.1

package/dist/trainers/sft-config.d.ts

@@ -1,32 +0,0 @@
-export declare class SFTConfigError extends Error {
-    constructor(message: string);
-}
-export interface SFTTrainerConfig {
-    model_name: string;
-    output_dir: string;
-    run_name: string;
-    learning_rate: number;
-    batch_size: number;
-    gradient_accumulation_steps: number;
-    num_epochs: number;
-    max_train_samples: number;
-    max_grad_norm: number;
-    weight_decay: number;
-    max_seq_length: number;
-    completion_only: boolean;
-    label_smoothing: number;
-    logging_steps: number;
-    save_steps: number;
-    max_checkpoints: number;
-    log_jsonl: boolean;
-    tui_mode: boolean;
-    seed: number;
-    resume_from_checkpoint: string;
-}
-declare const DEFAULT_SFT_CONFIG: SFTTrainerConfig;
-export declare function getDefaultSFTConfig(): SFTTrainerConfig;
-export declare function mergeSFTConfig(base: SFTTrainerConfig, update: Partial<SFTTrainerConfig>): SFTTrainerConfig;
-export declare function loadSFTTomlConfig(filePath: string): SFTTrainerConfig;
-export declare function applySFTOverrides(config: SFTTrainerConfig, overrides: string[]): SFTTrainerConfig;
-export { DEFAULT_SFT_CONFIG };
-//# sourceMappingURL=sft-config.d.ts.map

package/dist/trainers/sft-config.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"sft-config.d.ts","sourceRoot":"","sources":["../../src/trainers/sft-config.ts"],"names":[],"mappings":"AAIA,qBAAa,cAAe,SAAQ,KAAK;gBAC3B,OAAO,EAAE,MAAM;CAI5B;AAED,MAAM,WAAW,gBAAgB;IAE/B,UAAU,EAAE,MAAM,CAAC;IACnB,UAAU,EAAE,MAAM,CAAC;IACnB,QAAQ,EAAE,MAAM,CAAC;IAGjB,aAAa,EAAE,MAAM,CAAC;IACtB,UAAU,EAAE,MAAM,CAAC;IACnB,2BAA2B,EAAE,MAAM,CAAC;IACpC,UAAU,EAAE,MAAM,CAAC;IACnB,iBAAiB,EAAE,MAAM,CAAC;IAC1B,aAAa,EAAE,MAAM,CAAC;IACtB,YAAY,EAAE,MAAM,CAAC;IAGrB,cAAc,EAAE,MAAM,CAAC;IACvB,eAAe,EAAE,OAAO,CAAC;IACzB,eAAe,EAAE,MAAM,CAAC;IAGxB,aAAa,EAAE,MAAM,CAAC;IACtB,UAAU,EAAE,MAAM,CAAC;IACnB,eAAe,EAAE,MAAM,CAAC;IACxB,SAAS,EAAE,OAAO,CAAC;IACnB,QAAQ,EAAE,OAAO,CAAC;IAGlB,IAAI,EAAE,MAAM,CAAC;IACb,sBAAsB,EAAE,MAAM,CAAC;CAChC;AAED,QAAA,MAAM,kBAAkB,EAAE,gBAyBxB,CAAC;AA0GH,wBAAgB,mBAAmB,IAAI,gBAAgB,CAEtD;AAED,wBAAgB,cAAc,CAAC,IAAI,EAAE,gBAAgB,EAAE,MAAM,EAAE,OAAO,CAAC,gBAAgB,CAAC,GAAG,gBAAgB,CAa1G;AAED,wBAAgB,iBAAiB,CAAC,QAAQ,EAAE,MAAM,GAAG,gBAAgB,CAsBpE;AAED,wBAAgB,iBAAiB,CAAC,MAAM,EAAE,gBAAgB,EAAE,SAAS,EAAE,MAAM,EAAE,GAAG,gBAAgB,CAkBjG;AAED,OAAO,EAAE,kBAAkB,EAAE,CAAC"}

package/dist/trainers/sft-config.js

@@ -1,186 +0,0 @@
-import { parse as parseToml } from '@std/toml';
-import { readFileSync } from 'node:fs';
-import { resolve as resolvePath } from 'node:path';
-export class SFTConfigError extends Error {
-    constructor(message) {
-        super(message);
-        this.name = 'SFTConfigError';
-    }
-}
-const DEFAULT_SFT_CONFIG = Object.freeze({
-    model_name: 'Qwen/Qwen3-0.6B',
-    output_dir: 'outputs/sft',
-    run_name: 'sft-run',
-    learning_rate: 2e-5,
-    batch_size: 4,
-    gradient_accumulation_steps: 1,
-    num_epochs: 3,
-    max_train_samples: 0,
-    max_grad_norm: 1.0,
-    weight_decay: 0.01,
-    max_seq_length: 2048,
-    completion_only: false, // Changed to false for TRL parity
-    label_smoothing: 0.0,
-    logging_steps: 10,
-    save_steps: 100,
-    max_checkpoints: 3,
-    log_jsonl: true,
-    tui_mode: false,
-    seed: 42,
-    resume_from_checkpoint: '',
-});
-const SFT_CONFIG_VALUE_TYPES = {
-    model_name: 'string',
-    output_dir: 'string',
-    run_name: 'string',
-    learning_rate: 'number',
-    batch_size: 'number',
-    gradient_accumulation_steps: 'number',
-    num_epochs: 'number',
-    max_train_samples: 'number',
-    max_grad_norm: 'number',
-    weight_decay: 'number',
-    max_seq_length: 'number',
-    completion_only: 'boolean',
-    label_smoothing: 'number',
-    logging_steps: 'number',
-    save_steps: 'number',
-    max_checkpoints: 'number',
-    log_jsonl: 'boolean',
-    tui_mode: 'boolean',
-    seed: 'number',
-    resume_from_checkpoint: 'string',
-};
-const SFT_INTEGER_KEYS = new Set([
-    'batch_size',
-    'gradient_accumulation_steps',
-    'num_epochs',
-    'max_train_samples',
-    'max_seq_length',
-    'logging_steps',
-    'save_steps',
-    'max_checkpoints',
-    'seed',
-]);
-function cloneDefaults() {
-    return { ...DEFAULT_SFT_CONFIG };
-}
-function isConfigKey(value) {
-    return Object.prototype.hasOwnProperty.call(SFT_CONFIG_VALUE_TYPES, value);
-}
-function coerceBoolean(value, key) {
-    if (typeof value === 'boolean')
-        return value;
-    if (typeof value === 'string') {
-        const normalized = value.trim().toLowerCase();
-        if (['true', '1', 'yes', 'on'].includes(normalized))
-            return true;
-        if (['false', '0', 'no', 'off'].includes(normalized))
-            return false;
-    }
-    throw new SFTConfigError(`Invalid boolean for ${key}: ${String(value)}`);
-}
-function coerceNumber(value, key) {
-    if (typeof value === 'string' && value.trim() === '') {
-        throw new SFTConfigError(`Invalid number for ${key}: empty string`);
-    }
-    const parsed = typeof value === 'number' ? value : Number(value);
-    if (!Number.isFinite(parsed)) {
-        throw new SFTConfigError(`Invalid number for ${key}: ${String(value)}`);
-    }
-    if (SFT_INTEGER_KEYS.has(key) && !Number.isInteger(parsed)) {
-        throw new SFTConfigError(`Expected integer for ${key}, received ${parsed}`);
-    }
-    return parsed;
-}
-function coerceString(value, key) {
-    if (typeof value === 'string')
-        return value;
-    throw new SFTConfigError(`Invalid string for ${key}: ${String(value)}`);
-}
-function coerceValue(key, value) {
-    const expected = SFT_CONFIG_VALUE_TYPES[key];
-    if (expected === 'boolean') {
-        return coerceBoolean(value, key);
-    }
-    if (expected === 'number') {
-        return coerceNumber(value, key);
-    }
-    return coerceString(value, key);
-}
-function setConfigValue(config, key, value) {
-    config[key] = value;
-}
-function normalizeTomlRecord(record) {
-    const normalized = {};
-    for (const [rawKey, rawValue] of Object.entries(record)) {
-        if (!isConfigKey(rawKey)) {
-            continue;
-        }
-        setConfigValue(normalized, rawKey, coerceValue(rawKey, rawValue));
-    }
-    return normalized;
-}
-export function getDefaultSFTConfig() {
-    return cloneDefaults();
-}
-export function mergeSFTConfig(base, update) {
-    if (!update) {
-        return { ...base };
-    }
-    const result = { ...base };
-    for (const [key, value] of Object.entries(update)) {
-        if (value === undefined)
-            continue;
-        if (!isConfigKey(key)) {
-            throw new SFTConfigError(`Unknown configuration key: ${key}`);
-        }
-        setConfigValue(result, key, value);
-    }
-    return result;
-}
-export function loadSFTTomlConfig(filePath) {
-    const absolutePath = resolvePath(filePath);
-    let fileContents;
-    try {
-        fileContents = readFileSync(absolutePath, 'utf8');
-    }
-    catch (error) {
-        const message = error instanceof Error ? error.message : String(error);
-        throw new SFTConfigError(`Failed to read config at ${absolutePath}: ${message}`);
-    }
-    let parsedRaw;
-    try {
-        parsedRaw = parseToml(fileContents);
-    }
-    catch (error) {
-        const message = error instanceof Error ? error.message : String(error);
-        throw new SFTConfigError(`Failed to parse TOML at ${absolutePath}: ${message}`);
-    }
-    if (parsedRaw === null || typeof parsedRaw !== 'object' || Array.isArray(parsedRaw)) {
-        throw new SFTConfigError(`Expected table at ${absolutePath}`);
-    }
-    const parsed = parsedRaw;
-    const normalized = normalizeTomlRecord(parsed);
-    return mergeSFTConfig(getDefaultSFTConfig(), normalized);
-}
-export function applySFTOverrides(config, overrides) {
-    if (!overrides.length) {
-        return { ...config };
-    }
-    const accumulated = {};
-    for (const entry of overrides) {
-        const idx = entry.indexOf('=');
-        if (idx === -1) {
-            throw new SFTConfigError(`Invalid override "${entry}", expected key=value format`);
-        }
-        const key = entry.slice(0, idx).trim();
-        const rawValue = entry.slice(idx + 1).trim();
-        if (!isConfigKey(key)) {
-            throw new SFTConfigError(`Unknown configuration key in override: ${key}`);
-        }
-        setConfigValue(accumulated, key, coerceValue(key, rawValue));
-    }
-    return mergeSFTConfig(config, accumulated);
-}
-export { DEFAULT_SFT_CONFIG };

package/dist/trainers/sft-trainer.d.ts

@@ -1,141 +0,0 @@
-/**
- * SFT (Supervised Fine-Tuning) Trainer
- *
- * This module provides a Rust-native SFT training engine for training
- * models on fixed prompt-completion pairs using cross-entropy loss.
- *
- * ## Key Features
- * - Training loop runs in Rust (eliminates FFI overhead)
- * - Cross-entropy loss with completion masking (ignore_index=-100)
- * - Label smoothing support
- * - Gradient accumulation and clipping
- * - High-level train() method for full training runs
- * - Low-level trainStep() for custom training loops
- *
- * ## Usage
- * ```typescript
- * const trainer = await SFTTrainer.create({
- *   modelPath: './model',
- *   learningRate: 2e-5,
- *   numEpochs: 3,
- * });
- * await trainer.train(dataset);
- * ```
- */
-import { Qwen3Model, Qwen3Tokenizer, type SftStepMetrics } from '@mlx-node/core';
-import type { SFTTrainerConfig } from './sft-config';
-import { SFTDataset, type SFTBatch } from '../data/sft-dataset';
-import { type TrainingLogger } from './training-logger';
-export { SftTrainingEngine } from '@mlx-node/core';
-export type { SftEngineConfig, SftStepMetrics, SftEpochMetrics } from '@mlx-node/core';
-/**
- * Training state saved with checkpoints for resumption
- */
-export interface SFTTrainingState {
-    step: number;
-    epoch: number;
-    timestamp: string;
-    trainerType: 'sft';
-}
-/**
- * Training step result
- */
-export interface SFTTrainStepResult {
-    /** Step metrics */
-    metrics: SftStepMetrics;
-    /** Current epoch */
-    epoch: number;
-}
-/**
- * SFT Trainer - Rust-Native Training Engine
- *
- * Provides a TypeScript-friendly interface to the Rust SFT training engine.
- */
-export declare class SFTTrainer {
-    private engine;
-    private model;
-    private tokenizer;
-    private config;
-    private currentEpoch;
-    private currentStep;
-    /** Original model path (for tokenizer files when saving checkpoints) */
-    private originalModelPath?;
-    private paused;
-    private stopRequested;
-    private stdinInterface?;
-    private logger;
-    private sampleDisplayMode;
-    /**
-     * Create a new SFT trainer from a model
-     *
-     * @param model - Pre-loaded Qwen3 model
-     * @param tokenizer - Pre-loaded tokenizer
-     * @param config - Training configuration
-     * @param logger - Optional custom logger
-     */
-    constructor(model: Qwen3Model, tokenizer: Qwen3Tokenizer, config?: Partial<SFTTrainerConfig>, logger?: TrainingLogger);
-    /**
-     * Setup stdin handler for TUI control commands
-     */
-    private setupStdinHandler;
-    /**
-     * Handle a command received from stdin
-     */
-    private handleStdinCommand;
-    /**
-     * Wait for resume if paused
-     */
-    private waitForResume;
-    /**
-     * Create a trainer by loading a model from disk
-     *
-     * @param config - Configuration including modelPath
-     * @returns Promise<SFTTrainer>
-     */
-    static create(config: Partial<SFTTrainerConfig>): Promise<SFTTrainer>;
-    /**
-     * Find the latest checkpoint in the output directory
-     */
-    static findLatestCheckpoint(outputDir?: string): string | null;
-    /**
-     * Run a single training step
-     *
-     * @param batch - Tokenized batch with input_ids and labels
-     * @returns Training step metrics
-     */
-    trainStep(batch: SFTBatch): Promise<SFTTrainStepResult>;
-    /**
-     * Run a full training loop over a dataset
-     *
-     * @param dataset - SFT dataset or path to JSONL file
-     */
-    train(dataset: SFTDataset | string): Promise<void>;
-    /**
-     * Save a checkpoint with model weights and training state
-     *
-     * @param name - Checkpoint name (default: "checkpoint-{step}")
-     * @returns Path to saved checkpoint
-     */
-    saveCheckpoint(name?: string): Promise<string>;
-    /**
-     * Remove old checkpoints, keeping only the most recent ones
-     */
-    private cleanupOldCheckpoints;
-    /**
-     * Get current training step
-     */
-    get step(): number;
-    /**
-     * Get current epoch
-     */
-    get epoch(): number;
-    /**
-     * Get the underlying model for inference
-     */
-    getModel(): Qwen3Model;
-    /**
-     * Get the tokenizer
-     */
-    getTokenizer(): Qwen3Tokenizer;
-}
-//# sourceMappingURL=sft-trainer.d.ts.map

package/dist/trainers/sft-trainer.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"sft-trainer.d.ts","sourceRoot":"","sources":["../../src/trainers/sft-trainer.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;;;GAuBG;AAMH,OAAO,EAEL,UAAU,EACV,cAAc,EAGd,KAAK,cAAc,EAEpB,MAAM,gBAAgB,CAAC;AAExB,OAAO,KAAK,EAAE,gBAAgB,EAAE,MAAM,cAAc,CAAC;AAErD,OAAO,EAAE,UAAU,EAAkB,KAAK,QAAQ,EAAE,MAAM,qBAAqB,CAAC;AAChF,OAAO,EAAwB,KAAK,cAAc,EAAE,MAAM,mBAAmB,CAAC;AAG9E,OAAO,EAAE,iBAAiB,EAAE,MAAM,gBAAgB,CAAC;AACnD,YAAY,EAAE,eAAe,EAAE,cAAc,EAAE,eAAe,EAAE,MAAM,gBAAgB,CAAC;AAEvF;;GAEG;AACH,MAAM,WAAW,gBAAgB;IAC/B,IAAI,EAAE,MAAM,CAAC;IACb,KAAK,EAAE,MAAM,CAAC;IACd,SAAS,EAAE,MAAM,CAAC;IAClB,WAAW,EAAE,KAAK,CAAC;CACpB;AAED;;GAEG;AACH,MAAM,WAAW,kBAAkB;IACjC,mBAAmB;IACnB,OAAO,EAAE,cAAc,CAAC;IACxB,oBAAoB;IACpB,KAAK,EAAE,MAAM,CAAC;CACf;AAED;;;;GAIG;AACH,qBAAa,UAAU;IACrB,OAAO,CAAC,MAAM,CAAoB;IAClC,OAAO,CAAC,KAAK,CAAa;IAC1B,OAAO,CAAC,SAAS,CAAiB;IAClC,OAAO,CAAC,MAAM,CAAmB;IACjC,OAAO,CAAC,YAAY,CAAa;IACjC,OAAO,CAAC,WAAW,CAAa;IAChC,wEAAwE;IACxE,OAAO,CAAC,iBAAiB,CAAC,CAAS;IAGnC,OAAO,CAAC,MAAM,CAAkB;IAChC,OAAO,CAAC,aAAa,CAAkB;IACvC,OAAO,CAAC,cAAc,CAAC,CAAqB;IAC5C,OAAO,CAAC,MAAM,CAAiB;IAC/B,OAAO,CAAC,iBAAiB,CAA0C;IAEnE;;;;;;;OAOG;gBAED,KAAK,EAAE,UAAU,EACjB,SAAS,EAAE,cAAc,EACzB,MAAM,GAAE,OAAO,CAAC,gBAAgB,CAAM,EACtC,MAAM,CAAC,EAAE,cAAc;IAwCzB;;OAEG;IACH,OAAO,CAAC,iBAAiB;IAezB;;OAEG;IACH,OAAO,CAAC,kBAAkB;IAmC1B;;OAEG;YACW,aAAa;IAM3B;;;;;OAKG;WACU,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,gBAAgB,CAAC,GAAG,OAAO,CAAC,UAAU,CAAC;IA+D3E;;OAEG;IACH,MAAM,CAAC,oBAAoB,CAAC,SAAS,CAAC,EAAE,MAAM,GAAG,MAAM,GAAG,IAAI;IAmB9D;;;;;OAKG;IACG,SAAS,CAAC,KAAK,EAAE,QAAQ,GAAG,OAAO,CAAC,kBAAkB,CAAC;IAqB7D;;;;OAIG;IACG,KAAK,CAAC,OAAO,EAAE,UAAU,GAAG,MAAM,GAAG,OAAO,CAAC,IAAI,CAAC;IA8KxD;;;;;OAKG;IACG,cAAc,CAAC,IAAI,CAAC,EAAE,MAAM,GAAG,OAAO,CAAC,MAAM,CAAC;IAgDpD;;OAEG;IACH,OAAO,CAAC,qBAAqB;IAiC7B;;OAEG;IACH,IAAI,IAAI,IAAI,MAAM,CAEjB;IAED;;OAEG;IACH,IAAI,KAAK,IAAI,MAAM,CAElB;IAED;;OAEG;IACH,QAAQ,IAAI,UAAU;IAItB;;OAEG;IACH,YAAY,IAAI,cAAc;CAG/B"}