@mlx-node/trl 0.0.0

package/dist/data/dataset.d.ts

@@ -0,0 +1,22 @@
+import type { DatasetExample, ChatMessage, DatasetSplit, PromptFormatterOptions, PromptTemplate, DatasetLoader } from '../types';
+import { type PathValidationOptions } from '../utils/path-security';
+export interface LocalDatasetOptions extends PromptFormatterOptions, PathValidationOptions {
+    basePath?: string;
+    promptTemplate?: PromptTemplate;
+    metadata?: Record<string, unknown>;
+}
+export declare const SYSTEM_PROMPT: string;
+export declare const XML_COT_FORMAT = "<reasoning>\n{reasoning}\n</reasoning>\n<answer>\n{answer}\n</answer>";
+export declare const defaultPromptTemplate: PromptTemplate;
+export declare function createDatasetExample(prompt: ChatMessage[], metadata?: Record<string, unknown>): DatasetExample;
+export declare function extractGsm8kAnswer(raw: string): string | null;
+export declare function validateDatasetExample(example: DatasetExample): void;
+export declare function loadLocalGsm8kDataset(split: DatasetSplit, options?: LocalDatasetOptions & {
+    limit?: number;
+}): Promise<DatasetExample[]>;
+export declare class LocalGsm8kDatasetLoader implements DatasetLoader {
+    private readonly options;
+    constructor(options?: LocalDatasetOptions);
+    load(split: DatasetSplit, limit?: number): Promise<DatasetExample[]>;
+}
+//# sourceMappingURL=dataset.d.ts.map

package/dist/data/dataset.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"dataset.d.ts","sourceRoot":"","sources":["../../src/data/dataset.ts"],"names":[],"mappings":"AAEA,OAAO,KAAK,EACV,cAAc,EACd,WAAW,EAEX,YAAY,EACZ,sBAAsB,EACtB,cAAc,EACd,aAAa,EACd,MAAM,UAAU,CAAC;AAElB,OAAO,EAA2C,KAAK,qBAAqB,EAAE,MAAM,wBAAwB,CAAC;AAE7G,MAAM,WAAW,mBAAoB,SAAQ,sBAAsB,EAAE,qBAAqB;IACxF,QAAQ,CAAC,EAAE,MAAM,CAAC;IAClB,cAAc,CAAC,EAAE,cAAc,CAAC;IAChC,QAAQ,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;CACpC;AAUD,eAAO,MAAM,aAAa,QASlB,CAAC;AAET,eAAO,MAAM,cAAc,0EAKjB,CAAC;AAWX,eAAO,MAAM,qBAAqB,EAAE,cAYnC,CAAC;AAEF,wBAAgB,oBAAoB,CAAC,MAAM,EAAE,WAAW,EAAE,EAAE,QAAQ,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,GAAG,cAAc,CAK9G;AAED,wBAAgB,kBAAkB,CAAC,GAAG,EAAE,MAAM,GAAG,MAAM,GAAG,IAAI,CAE7D;AAED,wBAAgB,sBAAsB,CAAC,OAAO,EAAE,cAAc,GAAG,IAAI,CAYpE;AAwDD,wBAAsB,qBAAqB,CACzC,KAAK,EAAE,YAAY,EACnB,OAAO,GAAE,mBAAmB,GAAG;IAAE,KAAK,CAAC,EAAE,MAAM,CAAA;CAAO,GACrD,OAAO,CAAC,cAAc,EAAE,CAAC,CA4B3B;AAED,qBAAa,uBAAwB,YAAW,aAAa;IAC3D,OAAO,CAAC,QAAQ,CAAC,OAAO,CAAsB;gBAElC,OAAO,GAAE,mBAAwB;IAIvC,IAAI,CAAC,KAAK,EAAE,YAAY,EAAE,KAAK,CAAC,EAAE,MAAM,GAAG,OAAO,CAAC,cAAc,EAAE,CAAC;CAG3E"}

package/dist/data/dataset.js

@@ -0,0 +1,142 @@
+import { readFileSync } from 'node:fs';
+import { resolve as resolvePath } from 'node:path';
+import { extractHashAnswer } from '../utils/xml-parser';
+import { validatePathContainment, getAllowedRoot } from '../utils/path-security';
+const DEFAULT_BASE_PATH = resolvePath(process.cwd(), 'data/gsm8k');
+const VALID_SPLITS = new Set(['train', 'test']);
+export const SYSTEM_PROMPT = `
+Respond in the following format:
+
+<reasoning>
+...
+</reasoning>
+<answer>
+...
+</answer>
+`.trim();
+export const XML_COT_FORMAT = `<reasoning>
+{reasoning}
+</reasoning>
+<answer>
+{answer}
+</answer>`;
+const SYSTEM_MESSAGE = {
+    role: 'system',
+    content: SYSTEM_PROMPT,
+};
+function createMessage(role, content) {
+    return { role, content };
+}
+export const defaultPromptTemplate = (question, options) => {
+    const messages = [SYSTEM_MESSAGE];
+    if (options?.includeOneShot && options.oneShotExample) {
+        const { question: exampleQuestion, reasoning, answer } = options.oneShotExample;
+        messages.push(createMessage('user', exampleQuestion), createMessage('assistant', XML_COT_FORMAT.replace('{reasoning}', reasoning).replace('{answer}', answer)));
+    }
+    messages.push(createMessage('user', question));
+    return messages;
+};
+export function createDatasetExample(prompt, metadata) {
+    return {
+        prompt: prompt.map((message) => ({ ...message })), // defensive copy
+        metadata: metadata ? { ...metadata } : undefined,
+    };
+}
+export function extractGsm8kAnswer(raw) {
+    return extractHashAnswer(raw);
+}
+export function validateDatasetExample(example) {
+    if (!Array.isArray(example.prompt) || example.prompt.length === 0) {
+        throw new Error('Dataset example must contain at least one prompt message.');
+    }
+    for (const message of example.prompt) {
+        if (!message || typeof message.content !== 'string' || message.content.trim() === '') {
+            throw new Error('Prompt messages must include non-empty textual content.');
+        }
+        if (message.role !== 'system' && message.role !== 'user' && message.role !== 'assistant') {
+            throw new Error(`Unsupported chat role: ${String(message.role)}`);
+        }
+    }
+}
+function resolveBasePath(optionPath, options) {
+    const allowedRoot = getAllowedRoot(options);
+    if (!optionPath) {
+        // Default path - validate it's within allowed root
+        validatePathContainment(DEFAULT_BASE_PATH, allowedRoot);
+        return DEFAULT_BASE_PATH;
+    }
+    // Resolve and validate user-provided path
+    const resolved = resolvePath(allowedRoot, optionPath);
+    validatePathContainment(resolved, allowedRoot);
+    return resolved;
+}
+function datasetFileForSplit(split) {
+    if (!VALID_SPLITS.has(split)) {
+        throw new Error(`Unsupported GSM8K split "${split}". Expected one of: ${Array.from(VALID_SPLITS).join(', ')}`);
+    }
+    return `${split}.jsonl`;
+}
+function readDatasetFile(filePath) {
+    try {
+        return readFileSync(filePath, 'utf8');
+    }
+    catch (error) {
+        const message = error instanceof Error ? error.message : String(error);
+        throw new Error(`Failed to read dataset file at ${filePath}: ${message}`);
+    }
+}
+function readJsonl(path, limit) {
+    const fileContents = readDatasetFile(path);
+    const lines = fileContents.split(/\r?\n/).filter((line) => line.trim().length > 0);
+    const records = [];
+    const max = typeof limit === 'number' && limit >= 0 ? limit : Number.POSITIVE_INFINITY;
+    for (let i = 0; i < lines.length && records.length < max; i += 1) {
+        const line = lines[i];
+        try {
+            const parsed = JSON.parse(line);
+            if (typeof parsed.question !== 'string' || typeof parsed.answer !== 'string') {
+                throw new Error('Record must include string "question" and "answer" fields.');
+            }
+            records.push({ question: parsed.question, answer: parsed.answer });
+        }
+        catch (error) {
+            const message = error instanceof Error ? error.message : String(error);
+            throw new Error(`Failed to parse JSONL record at ${path}:${i + 1} - ${message}`);
+        }
+    }
+    return records;
+}
+export async function loadLocalGsm8kDataset(split, options = {}) {
+    const basePath = resolveBasePath(options.basePath, options);
+    const fileName = datasetFileForSplit(split);
+    const filePath = resolvePath(basePath, fileName);
+    // Additional validation: ensure the final file path stays within the base path
+    // This protects against any edge cases where the filename could escape
+    validatePathContainment(filePath, basePath);
+    const promptTemplate = options.promptTemplate ?? defaultPromptTemplate;
+    const records = readJsonl(filePath, options.limit);
+    const examples = records.map((record, index) => {
+        const prompt = promptTemplate(record.question, {
+            includeOneShot: options.includeOneShot,
+            oneShotExample: options.oneShotExample,
+        });
+        const example = createDatasetExample(prompt, {
+            split,
+            index,
+            raw_answer: record.answer,
+            ...options.metadata,
+        });
+        validateDatasetExample(example);
+        return example;
+    });
+    return examples;
+}
+export class LocalGsm8kDatasetLoader {
+    options;
+    constructor(options = {}) {
+        this.options = { ...options };
+    }
+    async load(split, limit) {
+        return loadLocalGsm8kDataset(split, { ...this.options, limit });
+    }
+}

package/dist/data/sft-dataset.d.ts

@@ -0,0 +1,156 @@
+/**
+ * SFT Dataset handling for Supervised Fine-Tuning
+ *
+ * Supports two data formats (auto-detected):
+ * 1. Prompt-Completion: { prompt: ChatMessage[], completion: ChatMessage }
+ * 2. Full Conversation: { messages: ChatMessage[] }
+ *
+ * Both formats produce tokenized batches with labels masked appropriately.
+ */
+import type { Qwen3Tokenizer } from '@mlx-node/core';
+import type { ChatMessage } from '../types';
+import { type PathValidationOptions } from '../utils/path-security';
+/**
+ * Special token IDs for SFT label masking
+ *
+ * These are used to detect assistant message boundaries in tokenized conversations.
+ * The IDs can be derived from the tokenizer or provided explicitly.
+ */
+export interface SpecialTokenIds {
+    /** Token ID for <|im_start|> */
+    imStart: number;
+    /** Token ID for <|im_end|> */
+    imEnd: number;
+    /** Token IDs that represent newlines (for detecting end of role header) */
+    newlineTokens: number[];
+}
+/**
+ * Prompt-Completion format for tool-use training
+ */
+export interface SFTPromptCompletionExample {
+    prompt: ChatMessage[];
+    completion: ChatMessage;
+}
+/**
+ * Full conversation format for multi-turn dialogue
+ */
+export interface SFTConversationExample {
+    messages: ChatMessage[];
+}
+/**
+ * Union type for SFT examples
+ */
+export type SFTExample = SFTPromptCompletionExample | SFTConversationExample;
+/**
+ * A tokenized batch ready for SFT training
+ */
+export interface SFTBatch {
+    inputIds: Int32Array;
+    labels: Int32Array;
+    shape: [number, number];
+}
+/**
+ * Configuration for SFT dataset
+ */
+export interface SFTDatasetConfig {
+    maxSeqLength?: number;
+    completionOnly?: boolean;
+    enableThinking?: boolean;
+    seed?: number;
+    /**
+     * Special token IDs for label masking.
+     *
+     * If not provided, these are automatically derived from the tokenizer.
+     * This option allows explicit overriding for custom tokenizers or
+     * non-standard vocabularies.
+     */
+    specialTokenIds?: Partial<SpecialTokenIds>;
+}
+/**
+ * SFT Dataset class for handling SFT training data
+ */
+export declare class SFTDataset {
+    private examples;
+    private tokenizer;
+    private config;
+    private format;
+    private shuffledIndices;
+    private rng;
+    /** Cached special token IDs for label masking */
+    private specialTokenIds;
+    constructor(examples: SFTExample[], tokenizer: Qwen3Tokenizer, config?: SFTDatasetConfig);
+    /**
+     * Get the number of examples in the dataset
+     */
+    get length(): number;
+    /**
+     * Shuffle dataset for a specific epoch using epoch-based seeding.
+     * This ensures reproducible shuffles across training resumes.
+     * Each epoch gets a deterministic shuffle based on (baseSeed + epoch).
+     *
+     * @param epoch - The epoch number (used as seed offset)
+     */
+    shuffleForEpoch(epoch: number): void;
+    /**
+     * Create a seeded pseudo-random number generator (Linear Congruential Generator)
+     */
+    private createSeededRandom;
+    /**
+     * Find length of common prefix between two token arrays
+     * Handles chat template quirks where prompt tokens may not be exact prefix of full tokens
+     */
+    private findCommonPrefixLength;
+    /**
+     * Tokenize a prompt-completion example
+     */
+    private tokenizePromptCompletion;
+    /**
+     * Check if a token ID is a newline token
+     */
+    private isNewlineToken;
+    /**
+     * Tokenize a conversation example
+     *
+     * For conversations, we train on all assistant turns.
+     * Non-assistant tokens (system, user) are masked with -100.
+     *
+     * Uses single-pass tokenization with token-based boundary detection.
+     * Token IDs are derived from the tokenizer for portability across models.
+     */
+    private tokenizeConversation;
+    /**
+     * Tokenize a single example based on its format
+     */
+    private tokenizeExample;
+    /**
+     * Collate multiple examples into a padded batch
+     */
+    collateBatch(indices: number[]): Promise<SFTBatch>;
+    /**
+     * Generate batches for training
+     */
+    batches(batchSize: number): AsyncGenerator<SFTBatch>;
+    /**
+     * Get total number of batches for a given batch size
+     */
+    numBatches(batchSize: number): number;
+}
+/**
+ * Load SFT dataset from a JSONL file
+ *
+ * Supports two formats:
+ * 1. Prompt-Completion: {"prompt": [...], "completion": {...}}
+ * 2. Conversation: {"messages": [...]}
+ *
+ * @param path - Path to the JSONL file (relative to cwd or allowedRoot)
+ * @param tokenizer - Qwen3 tokenizer instance
+ * @param config - Optional configuration including path validation options
+ */
+export declare function loadSFTDataset(path: string, tokenizer: Qwen3Tokenizer, config?: SFTDatasetConfig & {
+    limit?: number;
+} & PathValidationOptions): Promise<SFTDataset>;
+/**
+ * Create SFT dataset from examples directly
+ */
+export declare function createSFTDataset(examples: SFTExample[], tokenizer: Qwen3Tokenizer, config?: SFTDatasetConfig): SFTDataset;
+//# sourceMappingURL=sft-dataset.d.ts.map

package/dist/data/sft-dataset.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"sft-dataset.d.ts","sourceRoot":"","sources":["../../src/data/sft-dataset.ts"],"names":[],"mappings":"AAAA;;;;;;;;GAQG;AAIH,OAAO,KAAK,EAAE,cAAc,EAAE,MAAM,gBAAgB,CAAC;AACrD,OAAO,KAAK,EAAE,WAAW,EAAE,MAAM,UAAU,CAAC;AAC5C,OAAO,EAA2C,KAAK,qBAAqB,EAAE,MAAM,wBAAwB,CAAC;AAK7G;;;;;GAKG;AACH,MAAM,WAAW,eAAe;IAC9B,gCAAgC;IAChC,OAAO,EAAE,MAAM,CAAC;IAChB,8BAA8B;IAC9B,KAAK,EAAE,MAAM,CAAC;IACd,2EAA2E;IAC3E,aAAa,EAAE,MAAM,EAAE,CAAC;CACzB;AAgDD;;GAEG;AACH,MAAM,WAAW,0BAA0B;IACzC,MAAM,EAAE,WAAW,EAAE,CAAC;IACtB,UAAU,EAAE,WAAW,CAAC;CACzB;AAED;;GAEG;AACH,MAAM,WAAW,sBAAsB;IACrC,QAAQ,EAAE,WAAW,EAAE,CAAC;CACzB;AAED;;GAEG;AACH,MAAM,MAAM,UAAU,GAAG,0BAA0B,GAAG,sBAAsB,CAAC;AAE7E;;GAEG;AACH,MAAM,WAAW,QAAQ;IACvB,QAAQ,EAAE,UAAU,CAAC;IACrB,MAAM,EAAE,UAAU,CAAC;IACnB,KAAK,EAAE,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;CACzB;AAED;;GAEG;AACH,MAAM,WAAW,gBAAgB;IAC/B,YAAY,CAAC,EAAE,MAAM,CAAC;IACtB,cAAc,CAAC,EAAE,OAAO,CAAC;IACzB,cAAc,CAAC,EAAE,OAAO,CAAC;IACzB,IAAI,CAAC,EAAE,MAAM,CAAC;IAEd;;;;;;OAMG;IACH,eAAe,CAAC,EAAE,OAAO,CAAC,eAAe,CAAC,CAAC;CAC5C;AAeD;;GAEG;AACH,qBAAa,UAAU;IACrB,OAAO,CAAC,QAAQ,CAAe;IAC/B,OAAO,CAAC,SAAS,CAAiB;IAClC,OAAO,CAAC,MAAM,CAAkF;IAChG,OAAO,CAAC,MAAM,CAAuC;IACrD,OAAO,CAAC,eAAe,CAAW;IAClC,OAAO,CAAC,GAAG,CAAe;IAC1B,iDAAiD;IACjD,OAAO,CAAC,eAAe,CAAkB;gBAE7B,QAAQ,EAAE,UAAU,EAAE,EAAE,SAAS,EAAE,cAAc,EAAE,MAAM,GAAE,gBAAqB;IAsC5F;;OAEG;IACH,IAAI,MAAM,IAAI,MAAM,CAEnB;IAED;;;;;;OAMG;IACH,eAAe,CAAC,KAAK,EAAE,MAAM,GAAG,IAAI;IAYpC;;OAEG;IACH,OAAO,CAAC,kBAAkB;IAQ1B;;;OAGG;IACH,OAAO,CAAC,sBAAsB;IAS9B;;OAEG;YACW,wBAAwB;IA8CtC;;OAEG;IACH,OAAO,CAAC,cAAc;IAItB;;;;;;;;OAQG;YACW,oBAAoB;IA2DlC;;OAEG;YACW,eAAe;IAQ7B;;OAEG;IACG,YAAY,CAAC,OAAO,EAAE,MAAM,EAAE,GAAG,OAAO,CAAC,QAAQ,CAAC;IA8CxD;;OAEG;IACI,OAAO,CAAC,SAAS,EAAE,MAAM,GAAG,cAAc,CAAC,QAAQ,CAAC;IAQ3D;;OAEG;IACH,UAAU,CAAC,SAAS,EAAE,MAAM,GAAG,MAAM;CAGtC;AAoFD;;;;;;;;;;GAUG;AACH,wBAAsB,cAAc,CAClC,IAAI,EAAE,MAAM,EACZ,SAAS,EAAE,cAAc,EACzB,MAAM,CAAC,EAAE,gBAAgB,GAAG;IAAE,KAAK,CAAC,EAAE,MAAM,CAAA;CAAE,GAAG,qBAAqB,GACrE,OAAO,CAAC,UAAU,CAAC,CAarB;AAED;;GAEG;AACH,wBAAgB,gBAAgB,CAC9B,QAAQ,EAAE,UAAU,EAAE,EACtB,SAAS,EAAE,cAAc,EACzB,MAAM,CAAC,EAAE,gBAAgB,GACxB,UAAU,CAEZ"}