@genai-fi/nanogpt 0.20.0 → 0.20.2

package/dist/training/BasicTrainer.js

@@ -0,0 +1,265 @@
+import { _n as e, di as t, kt as n, ni as r, oi as i, qt as a } from "../dist-BewPQWjc.js";
+import { AdamWOptimizer as o } from "./AdamW.js";
+import { calculateAccuracy as s, calculateLoss as c } from "./loss.js";
+import l from "./Evaluator.js";
+import u from "../utilities/profile.js";
+import { createTensorStatistics as d } from "../checks/weights.js";
+//#region lib/training/BasicTrainer.ts
+var f = {
+	logInterval: 1,
+	maxEpochs: 100,
+	sftMode: "full",
+	batchSize: 32
+}, p = {
+	learningRate: 3e-4,
+	beta1: .9,
+	beta2: .99,
+	epsilon: 1e-8,
+	weightDecay: .01,
+	warmupSteps: 100,
+	decayEpochs: 100,
+	epochSteps: 1e4,
+	minLearningRate: 1e-5,
+	lossScaling: 1
+}, m = class {
+	tokenizer;
+	model;
+	optimizer;
+	running = !1;
+	lastState;
+	_gradientCheckpointing = !1;
+	_mixedPrecision = !1;
+	maskedLoss = !1;
+	optimizerConfig;
+	metrics = /* @__PURE__ */ new Set();
+	_labelSmoothing = 0;
+	_layerDrop = 0;
+	_dropout = 0;
+	constructor(e, t, n, r) {
+		this.tokenizer = t, this.model = e, this.optimizerConfig = {
+			...p,
+			...n,
+			lossScaling: e.lossScaling
+		};
+		let i = r || new o(this.optimizerConfig);
+		r && r.updateConfig(this.optimizerConfig), this.optimizer = i;
+	}
+	setLossMasking() {
+		this.maskedLoss = !0;
+	}
+	setGradientCheckpointing(e) {
+		this._gradientCheckpointing = e;
+	}
+	setMixedPrecision(e) {
+		this._mixedPrecision = e;
+	}
+	setLabelSmoothing(e) {
+		this._labelSmoothing = e;
+	}
+	setDropout(e) {
+		this._dropout = e;
+	}
+	setLayerDrop(e) {
+		this._layerDrop = e;
+	}
+	setLearningRate(e) {
+		this.optimizerConfig.learningRate = e, this.updateOptimizer();
+	}
+	setMetrics(e) {
+		this.metrics = new Set(e);
+	}
+	reset() {
+		this.lastState = void 0, this.running = !1;
+	}
+	stop() {
+		this.running = !1;
+	}
+	get isRunning() {
+		return this.running;
+	}
+	getOptimizer() {
+		return this.optimizer;
+	}
+	updateOptimizer(e) {
+		e && (this.optimizerConfig = {
+			...this.optimizerConfig,
+			...e
+		}), this.optimizer.updateConfig(this.optimizerConfig);
+	}
+	resumeFromLog(e) {
+		(!this.lastState || this.lastState.step === 0) && (this.lastState = {
+			losses: [],
+			validationLosses: [],
+			logStartTime: 0,
+			step: e.step,
+			lastLoss: e.trainingMetrics.loss,
+			totalSteps: e.step,
+			trainingDuration: e.duration
+		});
+	}
+	trainStep(n, o, l = !1, u = !1) {
+		return t(() => {
+			this.model.getProfiler()?.startMemory();
+			let { xs: t, ys: d } = o, { value: f, grads: p } = a(() => {
+				let r = this.model.forward({
+					training: !0,
+					checkpointing: this._gradientCheckpointing,
+					mixedPrecision: this._mixedPrecision,
+					dropout: this._dropout,
+					layerDrop: this._layerDrop,
+					ropePositionOffset: 0
+				}, t), a = c(r, d, this.maskedLoss, !1, this._labelSmoothing);
+				this.metrics.has("accuracy") && (n.accuracy = s(r, d), i(n.accuracy)), r.dispose();
+				let o = a.mul(e(this.optimizerConfig.lossScaling));
+				return a.dispose(), o;
+			});
+			if (l) this.model.getProfiler()?.endMemory("Training");
+			else {
+				let e = this.optimizer.applyGradients(p);
+				this.metrics.has("gradientNorm") ? (n.gradientNorm = e, i(e)) : (n.gradientNorm = void 0, e.dispose());
+				let t = Object.keys(p);
+				this.model.weightStore.touchVariables(t), this.model.getProfiler()?.endMemory("Training"), u ? (n.gradients = p, Object.values(p).forEach((e) => i(e))) : r(p);
+			}
+			return f.mul(e(1 / this.optimizerConfig.lossScaling));
+		});
+	}
+	async dummyPass() {
+		let e = n([1, this.model.config.blockSize], "int32"), t = n([1, this.model.config.blockSize], "int32");
+		try {
+			let n = this.trainStep({}, {
+				xs: e,
+				ys: t
+			}, !0);
+			await n.data(), n.dispose();
+		} catch (e) {
+			console.error("Error during dummy pass:", e);
+		} finally {
+			e.dispose(), t.dispose();
+		}
+	}
+	dispose() {
+		this.optimizer && this.optimizer.dispose();
+	}
+	createEmptyState() {
+		return {
+			step: 0,
+			lastLoss: 1e6,
+			totalSteps: 0,
+			losses: [],
+			validationLosses: [],
+			logStartTime: 0,
+			trainingDuration: 0,
+			...this.lastState || {}
+		};
+	}
+	async stepDataset(e, t, n) {
+		let { logInterval: i = 10 } = {
+			...f,
+			...t
+		};
+		t.metrics && this.setMetrics(t.metrics);
+		let a = Date.now(), o = this.createEmptyState();
+		this.lastState = o, await this.dummyPass(), this.metrics.has("memoryUsage") && (this.model.getProfiler() || this.model.setProfiler(new u())), this.running = !0, o.logStartTime = a;
+		let s = n ? new l(this.model, n, this.maskedLoss) : void 0, c = await e.iterator();
+		try {
+			for (; this.running;) {
+				let e = await c.next();
+				if (e.done) break;
+				let n = e.value, r = this.trainStep(o, n, !1);
+				if (t.debug) {
+					let e = (await r.data())[0];
+					if (isNaN(e) || !isFinite(e)) throw console.error("Invalid loss value:", e), console.error("Batch xs:", n.xs.toString()), console.error("Batch ys:", n.ys.toString()), console.error("State:", o), Error("Loss is NaN or Infinity");
+					console.log(`Step ${o.step}: Loss = ${e}`);
+				}
+				n.xs.dispose(), n.ys.dispose(), o.step++, o.totalSteps++, o.step % i === 0 ? await this.performLogging(r, n.xs.shape[0], t, s) : (o.gradientNorm &&= (o.gradientNorm.dispose(), void 0), o.accuracy &&= (o.accuracy.dispose(), void 0)), r.dispose();
+			}
+		} catch (e) {
+			throw console.error("Training error:", e), e;
+		}
+		throw this.model.trainingState = {
+			steps: o.totalSteps,
+			learningRate: this.optimizer.lr,
+			batchSize: t.batchSize || 32,
+			loss: o.lastLoss,
+			tokensProcessed: o.totalSteps * (t.batchSize || 32) * this.model.config.blockSize,
+			duration: o.trainingDuration
+		}, r(), this.running = !1, Error("No log returned before training stopped.");
+	}
+	async performLogging(e, t, n, r) {
+		let i = n?.onStep, a = this.metrics.has("gradientStatistics"), o = (await e.data())[0], s = this.lastState;
+		s.lastLoss = o, s.trainingDuration += Date.now() - s.logStartTime;
+		let c = s.totalSteps * t * this.model.config.blockSize, l = {
+			trainingMetrics: {
+				loss: s.lastLoss,
+				perplexity: this.metrics.has("perplexity") ? Math.exp(s.lastLoss) : void 0,
+				accuracy: s.accuracy ? (await s.accuracy.data())[0] : void 0
+			},
+			step: s.step,
+			time: Date.now() - s.logStartTime,
+			gradientNorm: s.gradientNorm ? (await s.gradientNorm.data())[1] : void 0,
+			batchSize: t,
+			learningRate: this.metrics.has("learningRate") ? this.optimizer.lr : void 0,
+			duration: s.trainingDuration,
+			totalTokens: c,
+			tokensPerSecond: c / (s.trainingDuration / 1e3),
+			memoryUsage: this.metrics.has("memoryUsage") ? this.model.getProfiler()?.getPeakMemory() || 0 : void 0
+		};
+		if (s.gradientNorm &&= (s.gradientNorm.dispose(), void 0), s.accuracy &&= (s.accuracy.dispose(), void 0), this.model.trainingState = {
+			steps: s.totalSteps,
+			learningRate: this.optimizer.lr,
+			batchSize: t,
+			loss: s.lastLoss,
+			tokensProcessed: c,
+			duration: s.trainingDuration
+		}, a && s.gradients) {
+			let e = /* @__PURE__ */ new Map();
+			for (let [t, n] of Object.entries(s.gradients)) e.set(t, await d(n)), n.dispose();
+			l.gradientMetrics = e;
+		}
+		if (r) try {
+			let e = await r.evaluate(5);
+			Array.isArray(e) ? l.validationMetrics = {
+				loss: e[0].loss,
+				accuracy: e[0].accuracy
+			} : (s.validationLosses.push(e.loss), l.validationMetrics = {
+				accuracy: e.accuracy,
+				loss: e.loss,
+				perplexity: this.metrics.has("perplexity") ? Math.exp(e.loss) : void 0
+			});
+		} catch (e) {
+			console.error("Validation error:", e);
+		}
+		i && await i(l), s.logStartTime = Date.now();
+	}
+	async trainOnDataset(e, t, n) {
+		let { logInterval: i = 10, maxEpochs: a = Infinity } = {
+			...f,
+			...t
+		}, o = a * (t?.epochSteps || 1e3);
+		t.metrics && this.setMetrics(t.metrics);
+		let s = Date.now(), c = this.createEmptyState();
+		this.lastState = c, await this.dummyPass(), t?.metrics?.includes("memoryUsage") && (this.model.getProfiler() || this.model.setProfiler(new u())), this.running = !0, c.logStartTime = s;
+		let d = n ? new l(this.model, n, this.maskedLoss) : void 0, p = await e.iterator();
+		try {
+			for (; this.running;) {
+				let e = await p.next();
+				if (e.done) break;
+				let n = e.value, r = c.step % i === 0, a = (t?.metrics?.includes("gradientStatistics") || !1) && r, s = this.trainStep(c, n, !1, a);
+				if (t.debug) {
+					let e = (await s.data())[0];
+					if (isNaN(e) || !isFinite(e)) throw console.error("Invalid loss value:", e), console.error("Batch xs:", await n.xs.array()), console.error("Batch ys:", await n.ys.array()), console.error("State:", c), Error("Loss is NaN or Infinity");
+					console.log(`Step ${c.step}: Loss = ${e}`);
+				}
+				n.xs.dispose(), n.ys.dispose(), c.step++, c.totalSteps++, r ? await this.performLogging(s, n.xs.shape[0], t, d) : (c.gradientNorm &&= (c.gradientNorm.dispose(), void 0), c.accuracy &&= (c.accuracy.dispose(), void 0)), s.dispose(), c.step >= o && this.stop();
+			}
+		} catch (e) {
+			throw console.error("Training error:", e), r(), e;
+		}
+		return r(), this.running = !1, {
+			losses: c.losses,
+			validationLosses: c.validationLosses
+		};
+	}
+};
+//#endregion
+export { m as default };

package/dist/training/DatasetBuilder.d.ts

@@ -0,0 +1,26 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+import { Conversation, ITokeniser } from '../tokeniser/type';
+import { Dataset } from '@tensorflow/tfjs-data';
+export declare const PAGE_FACTOR = 8;
+export declare function flattenTokens(textData: Conversation[][], tokenizer: ITokeniser): Uint16Array;
+export declare function flattenTokensWithMask(textData: Conversation[][], tokenizer: ITokeniser): {
+    tokens: Uint16Array;
+    mask: Uint8Array;
+};
+export declare function shuffle(array: Uint32Array): Uint32Array;
+export interface DatasetState {
+    shuffledIndexes: Uint32Array;
+    step: number;
+}
+export declare class DatasetBuilder {
+    tokenizer: ITokeniser;
+    blockSize: number;
+    constructor(tokenizer: ITokeniser, blockSize?: number);
+    createTextDataset(flatTokens: Uint16Array, batchSize?: number, indexes?: Uint32Array, mask?: Uint8Array, ignoreIndex?: number): Promise<{
+        dataset: Dataset<{
+            xs: Tensor;
+            ys: Tensor;
+        }>;
+        state: DatasetState;
+    }>;
+}

package/dist/training/DatasetBuilder.js

@@ -0,0 +1,2 @@
+import { a as e, i as t, n, r, t as i } from "../DatasetBuilder-DgURD85T.js";
+export { i as DatasetBuilder, n as PAGE_FACTOR, r as flattenTokens, t as flattenTokensWithMask, e as shuffle };

package/dist/training/Evaluator.d.ts

@@ -0,0 +1,19 @@
+import { Dataset } from '@tensorflow/tfjs-data';
+import { TensorContainer } from '@tensorflow/tfjs-core';
+import { default as Model, ModelForwardAttributes } from '../../models/model';
+interface Result {
+    loss: number;
+    accuracy: number;
+}
+export default class Evaluator {
+    private model;
+    private iterator?;
+    private xs?;
+    private ys?;
+    private masked;
+    constructor(model: Model<ModelForwardAttributes>, dataset: Dataset<TensorContainer>, masked?: boolean);
+    dispose(): void;
+    private calculateBatchLoss;
+    evaluate(maxBatches?: number): Promise<Result | Result[]>;
+}
+export {};

package/dist/training/Evaluator.js

@@ -0,0 +1,48 @@
+import { di as e } from "../dist-BewPQWjc.js";
+import { calculateAccuracy as t, calculateLoss as n } from "./loss.js";
+//#region lib/training/Evaluator.ts
+var r = class {
+	model;
+	iterator;
+	xs;
+	ys;
+	masked = !1;
+	constructor(e, t, n) {
+		this.model = e, this.masked = !!n, this.iterator = t.iterator();
+	}
+	dispose() {
+		this.xs && this.xs.dispose(), this.ys && this.ys.dispose();
+	}
+	async calculateBatchLoss(r, i, a, o) {
+		let [s, c] = e(() => {
+			let e = this.model.forward({ training: !1 }, r), s = n(e, i, o, a), c = t(e, i);
+			return e.dispose(), [s, c];
+		}), l = await s.array(), u = await c.array(), d = l, f = u;
+		return c.dispose(), s.dispose(), Array.isArray(d) ? d.map((e) => ({
+			loss: e,
+			accuracy: f
+		})) : {
+			loss: d,
+			accuracy: f
+		};
+	}
+	async evaluate(e = 100) {
+		let t = 0, n = 0, r = 0;
+		if (this.iterator) {
+			let i = await this.iterator;
+			for (let a = 0; a < e; a++) {
+				let e = await i.next();
+				if (e.done) break;
+				let { xs: a, ys: o } = e.value, s = await this.calculateBatchLoss(a, o, !1, this.masked);
+				a.dispose(), o.dispose(), t += s.loss, n += s.accuracy, r++;
+			}
+			return {
+				loss: t / r,
+				accuracy: n / r
+			};
+		} else if (this.xs && this.ys) return this.calculateBatchLoss(this.xs, this.ys, !0, !0);
+		throw Error("No data available for evaluation");
+	}
+};
+//#endregion
+export { r as default };

package/dist/training/LRScheduler.d.ts

@@ -0,0 +1,12 @@
+import { LRSchedulerConfig } from './types';
+export default class LRScheduler {
+    protected learningRate: number;
+    private config;
+    private step;
+    private startLearningRate;
+    constructor(learningRate: number, config: LRSchedulerConfig);
+    serializeConfig(): LRSchedulerConfig;
+    updateConfig(newConfig: Partial<LRSchedulerConfig>, learningRate?: number): void;
+    get lr(): number;
+    getNextLR(): number;
+}

package/dist/training/LRScheduler.js

@@ -0,0 +1,38 @@
+//#region lib/training/LRScheduler.ts
+var e = class {
+	learningRate;
+	config;
+	step = 0;
+	startLearningRate;
+	constructor(e, t) {
+		this.learningRate = e, this.config = t, this.startLearningRate = e, t.step !== void 0 && (this.step = t.step);
+	}
+	serializeConfig() {
+		return {
+			...this.config,
+			step: this.step
+		};
+	}
+	updateConfig(e, t) {
+		this.config = {
+			...this.config,
+			...e
+		}, t !== void 0 && (this.startLearningRate = t);
+	}
+	get lr() {
+		return this.learningRate;
+	}
+	getNextLR() {
+		let e = this.step;
+		if (this.config.warmupSteps > 0 && e < this.config.warmupSteps) {
+			let t = (e + 1) / this.config.warmupSteps, n = this.startLearningRate * t;
+			return this.learningRate = n, this.step++, n;
+		}
+		let t = this.config.epochSteps * this.config.decayEpochs;
+		if (e >= t || t <= this.config.warmupSteps) return this.learningRate = this.config.minLearningRate, this.step++, this.config.minLearningRate;
+		let n = (e - this.config.warmupSteps) / (t - this.config.warmupSteps), r = .5 * (1 + Math.cos(Math.PI * n)), i = this.config.minLearningRate + r * (this.startLearningRate - this.config.minLearningRate);
+		return this.learningRate = i, this.step++, i;
+	}
+};
+//#endregion
+export { e as default };

package/dist/training/PreTrainer.d.ts

@@ -0,0 +1,11 @@
+import { default as Model, ModelForwardAttributes } from '../../models/model';
+import { default as BasicTrainer } from './BasicTrainer';
+import { ITokeniser } from '../../tokeniser/type';
+import { DatasetBuilder } from './DatasetBuilder';
+import { AdamWOptimizer } from './AdamW';
+import { AdamWOptimizerConfig } from './types';
+export default class PreTrainer extends BasicTrainer {
+    tokenizer: ITokeniser;
+    datasetBuilder: DatasetBuilder;
+    constructor(model: Model<ModelForwardAttributes>, tokenizer: ITokeniser, optConfig?: Partial<AdamWOptimizerConfig>, optimizer?: AdamWOptimizer);
+}

package/dist/training/PreTrainer.js

@@ -0,0 +1,22 @@
+import { t as e } from "../DatasetBuilder-DgURD85T.js";
+import t from "./BasicTrainer.js";
+//#region lib/training/PreTrainer.ts
+var n = {
+	decayEpochs: 100,
+	epochSteps: 1e4,
+	warmupSteps: 1e3,
+	minLearningRate: 3e-5,
+	weightDecay: .1,
+	learningRate: 3e-4
+}, r = class extends t {
+	tokenizer;
+	datasetBuilder;
+	constructor(t, r, i, a) {
+		super(t, r, {
+			...n,
+			...i
+		}, a), this.tokenizer = r, this.optimizerConfig.minLearningRate = i?.minLearningRate ?? this.optimizerConfig.learningRate / 20, this.updateOptimizer(), this.datasetBuilder = new e(r, t.config.blockSize);
+	}
+};
+//#endregion
+export { r as default };

package/dist/training/SFTTrainer.d.ts

@@ -0,0 +1,12 @@
+import { default as Model, ModelForwardAttributes } from '../../models/model';
+import { default as BasicTrainer } from './BasicTrainer';
+import { ITokeniser } from '../../tokeniser/type';
+import { AdamWOptimizer } from './AdamW';
+import { AdamWOptimizerConfig } from './types';
+import { DatasetBuilder } from './DatasetBuilder';
+export default class SFTTrainer extends BasicTrainer {
+    tokenizer: ITokeniser;
+    datasetBuilder: DatasetBuilder;
+    loraName?: string;
+    constructor(model: Model<ModelForwardAttributes>, tokenizer: ITokeniser, optConfig?: Partial<AdamWOptimizerConfig>, optimizer?: AdamWOptimizer);
+}

package/dist/training/SFTTrainer.js

@@ -0,0 +1,24 @@
+import { t as e } from "../DatasetBuilder-DgURD85T.js";
+import t from "./BasicTrainer.js";
+//#region lib/training/SFTTrainer.ts
+var n = {
+	decayEpochs: 100,
+	epochSteps: 1e4,
+	warmupSteps: 100,
+	minLearningRate: 1e-5,
+	weightDecay: .1,
+	beta2: .95,
+	learningRate: 3e-4
+}, r = class extends t {
+	tokenizer;
+	datasetBuilder;
+	loraName;
+	constructor(t, r, i, a) {
+		super(t, r, {
+			...n,
+			...i
+		}, a), this.tokenizer = r, this.optimizerConfig.minLearningRate = i?.minLearningRate ?? this.optimizerConfig.learningRate / 20, this.updateOptimizer(), this.datasetBuilder = new e(r, t.config.blockSize), this.maskedLoss = !0;
+	}
+};
+//#endregion
+export { r as default };

package/dist/training/loss.d.ts

@@ -0,0 +1,3 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+export declare function calculateLoss(logits: Tensor, targets: Tensor, masked?: boolean, keepBatch?: boolean, labelSmoothing?: number): Tensor;
+export declare function calculateAccuracy(logits: Tensor, targets: Tensor): Tensor;

package/dist/training/loss.js

@@ -0,0 +1,19 @@
+import { createSoftmaxCrossEntropyWithGrad as e } from "./sparseCrossEntropy.js";
+//#region lib/training/loss.ts
+function t(t, n, r, i, a) {
+	try {
+		return e(r, i, a && a > 0 ? a : void 0)(t, n);
+	} catch (e) {
+		throw console.error("Error computing loss:", e), Error(`Loss computation failed: ${e}`);
+	}
+}
+function n(e, t) {
+	try {
+		let n = e.argMax(-1), r = n.equal(t).cast("float32"), i = r.mean();
+		return n.dispose(), r.dispose(), i;
+	} catch (e) {
+		throw console.error("Error computing accuracy:", e), Error(`Accuracy computation failed: ${e}`);
+	}
+}
+//#endregion
+export { n as calculateAccuracy, t as calculateLoss };

package/dist/training/orthoGrad.d.ts

@@ -0,0 +1,2 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+export declare function orthogonalizeGradient(weight: Tensor, gradient: Tensor, epsilon: number): Tensor;

package/dist/training/orthoGrad.js

@@ -0,0 +1,10 @@
+import { di as e } from "../dist-BewPQWjc.js";
+//#region lib/training/orthoGrad.ts
+function t(t, n, r) {
+	return e(() => {
+		let e = t.reshape([-1]), i = n.reshape([-1]), a = e.mul(e).sum().add(r), o = e.mul(i).sum().div(a), s = i.sub(e.mul(o)), c = i.norm(), l = s.norm().add(r);
+		return s.mul(c.div(l)).reshape(n.shape);
+	});
+}
+//#endregion
+export { t as orthogonalizeGradient };

package/dist/training/sparseCrossEntropy.d.ts

@@ -0,0 +1,7 @@
+import * as tf from '@tensorflow/tfjs-core';
+/**
+ * Numerically stable sparse cross-entropy with gradient support
+ * This version handles potential numerical issues better
+ */
+export declare function sparseSoftmaxCrossEntropy(logits: tf.Tensor, labels: tf.Tensor, validMask?: tf.Tensor, keepBatch?: boolean, originalBatchShape?: number[], labelSmoothing?: number): tf.Tensor;
+export declare function createSoftmaxCrossEntropyWithGrad(masked?: boolean, keepBatch?: boolean, labelSmoothing?: number): (...args: tf.Tensor[]) => tf.Tensor<tf.Rank>;

package/dist/training/sparseCrossEntropy.js

@@ -0,0 +1,47 @@
+import { At as e, Bt as t, Gr as n, Nn as r, Pn as i, Rt as a, St as o, Wr as s, Wt as c, Y as l, _n as u, bn as d, di as f, mn as p, qr as m } from "../dist-BewPQWjc.js";
+import { gatherSub as h } from "../ops/gatherSub.js";
+import { scatterSub as g } from "../ops/scatterSub.js";
+//#region lib/training/sparseCrossEntropy.ts
+function _(r, i, o, c, l, u = 0) {
+	return f(() => {
+		let f = r.shape[r.shape.length - 1], g = l || r.shape.slice(0, -1), _ = g.reduce((e, t) => e * t, 1), v = r.shape.length > 2 ? r.reshape([_, f]) : r, y = i.shape.length > 1 ? i.reshape([_]).cast("int32") : i.cast("int32"), b = t(v, d(v, -1, !0)), x = a(b, -1), S = h(x, y, b), C;
+		if (u > 0) {
+			let n = t(x, e(b, -1));
+			C = m(s(S, 1 - u), s(n, u));
+		} else C = S;
+		if (o) if (C = s(C, o), c) {
+			let e = p(o.reshape(g), -1);
+			C = n(p(C.reshape(g), -1), e);
+		} else {
+			let e = p(o);
+			C = n(p(C), e);
+		}
+		else C = c ? e(C.reshape(g), -1) : e(C);
+		return C;
+	});
+}
+function v(e, n, a = 0) {
+	return c((c, d, m) => {
+		let h = c.shape[c.shape.length - 1], v = c.shape.slice(0, -1), y = v.reduce((e, t) => e * t, 1), b = c.reshape([y, h]), x = d.reshape([y]).cast("int32"), S, C = null;
+		if (e) {
+			let e = u(65535, "int32"), t = o(x, e);
+			C = t.cast("float32"), S = i(t, x, r(x)), e.dispose(), t.dispose();
+		} else S = x;
+		let w = _(b, S, C || void 0, n, v, a);
+		return m(C ? [
+			b,
+			S,
+			C
+		] : [b, S]), b.dispose(), x.dispose(), {
+			value: w,
+			gradFunc: (n, i) => f(() => {
+				let o = i[0], f = i[1], m = e ? i[2] : void 0, _ = l(o), v = m ? p(m) : u(o.shape[0], "float32"), y = n.div(v).broadcastTo([o.shape[0]]), b = m && e ? s(y, m) : y, x;
+				x = a > 0 ? g(t(_, a / h), f, s(b, 1 - a)) : g(_, f, b);
+				let S = r(d);
+				return [x.reshape(c.shape), S];
+			})
+		};
+	});
+}
+//#endregion
+export { v as createSoftmaxCrossEntropyWithGrad, _ as sparseSoftmaxCrossEntropy };

package/dist/training/tasks/ConversationTask.d.ts

@@ -0,0 +1,18 @@
+import { Conversation, ITokeniser } from '../../../main';
+import { Task } from './Task';
+export default class ConversationTask extends Task {
+    private rawConvo;
+    private shuffledIndices;
+    private index;
+    get length(): number;
+    constructor(conversations: Conversation[][]);
+    hasMoreConversations(): boolean;
+    nextConversation(): Conversation[] | null;
+    nextTokens(tokeniser: ITokeniser): number[] | null;
+    nextTokens(tokeniser: ITokeniser, masking: boolean): {
+        tokens: number[];
+        mask: boolean[];
+    } | null;
+    shuffle(): void;
+    estimateTokens(tokeniser: ITokeniser): Promise<number>;
+}

package/dist/training/tasks/ConversationTask.js

@@ -0,0 +1,38 @@
+import { a as e } from "../../DatasetBuilder-DgURD85T.js";
+import { Task as t } from "./Task.js";
+//#region lib/training/tasks/ConversationTask.ts
+var n = class extends t {
+	rawConvo;
+	shuffledIndices = null;
+	index = 0;
+	get length() {
+		return this.rawConvo.length;
+	}
+	constructor(e) {
+		super(), this.rawConvo = e;
+	}
+	hasMoreConversations() {
+		return this.index < this.rawConvo.length;
+	}
+	nextConversation() {
+		if (this.index >= this.rawConvo.length) return null;
+		let e = this.rawConvo[this.shuffledIndices ? this.shuffledIndices[this.index] : this.index];
+		return this.index++, e;
+	}
+	nextTokens(e, t) {
+		let n = this.nextConversation();
+		return n ? e.encodeConversation(n, !1, t) : null;
+	}
+	shuffle() {
+		if (!this.shuffledIndices) {
+			this.shuffledIndices = new Uint32Array(this.rawConvo.length);
+			for (let e = 0; e < this.rawConvo.length; e++) this.shuffledIndices[e] = e;
+		}
+		e(this.shuffledIndices), this.index = 0;
+	}
+	async estimateTokens(e) {
+		return e.encodeConversation(this.rawConvo[0]).length * this.length;
+	}
+};
+//#endregion
+export { n as default };

package/dist/training/tasks/PretrainingTask.d.ts

@@ -0,0 +1,17 @@
+import { Conversation, ITokeniser } from '../../../main';
+import { Task } from './Task';
+export default class PretrainingTask extends Task {
+    private rawText;
+    private index;
+    get length(): number;
+    constructor(texts: string[]);
+    hasMoreConversations(): boolean;
+    nextConversation(): Conversation[] | null;
+    nextTokens(tokeniser: ITokeniser): number[] | null;
+    nextTokens(tokeniser: ITokeniser, masking: boolean): {
+        tokens: number[];
+        mask: boolean[];
+    } | null;
+    shuffle(): void;
+    estimateTokens(tokeniser: ITokeniser): Promise<number>;
+}