@genai-fi/nanogpt 0.23.0 → 1.0.0

package/dist/training/DatasetBuilder.d.ts

@@ -7,11 +7,11 @@ export declare function flattenTokensWithMask(textData: Conversation[][], tokeni
     tokens: Uint16Array;
     mask: Uint8Array;
 };
-export declare function shuffle(array: Uint32Array): Uint32Array;
+export declare function shuffle(array: Uint32Array | Uint16Array): Uint32Array | Uint16Array;
 export interface DatasetState {
-    shuffledShards: Uint32Array;
-    shuffledIndexes: Uint32Array;
-    lastShardIndexes: Uint32Array;
+    shuffledShards: Uint16Array;
+    shuffledIndexes: Uint16Array;
+    lastShardIndexes: Uint16Array;
     currentShard: Uint16Array | null;
     nextShard: Uint16Array | null;
     currentMask: Uint8Array | null;

package/dist/training/DatasetBuilder.js

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

package/dist/training/PreTrainer.js

@@ -1,4 +1,4 @@
-import { t as e } from "../DatasetBuilder-C0iJT29K.js";
+import { t as e } from "../DatasetBuilder-DU1G1OKX.js";
 import t from "./BasicTrainer.js";
 //#region lib/training/PreTrainer.ts
 var n = {

package/dist/training/SFTTrainer.js

@@ -1,4 +1,4 @@
-import { t as e } from "../DatasetBuilder-C0iJT29K.js";
+import { t as e } from "../DatasetBuilder-DU1G1OKX.js";
 import t from "./BasicTrainer.js";
 //#region lib/training/SFTTrainer.ts
 var n = {

package/dist/training/configure.d.ts

@@ -0,0 +1,3 @@
+import { TrainingOptions } from './types';
+import { default as Model, ModelForwardAttributes } from '../../models/model';
+export default function configureModel(model: Model<ModelForwardAttributes>, options?: TrainingOptions): void;

package/dist/training/configure.js

@@ -0,0 +1,32 @@
+import { t as e } from "../v4-BK7K-jy_.js";
+//#region lib/training/configure.ts
+function t(t, n) {
+	let r = n?.method.supervised || "full", i = n?.method.type || "pretraining";
+	if (i === "pretraining" && (t.hasLoRA() && t.detachLoRA(), t.weightStore.setTrainable(["*"]), n && (t.metaData.pretrainingSettings = n)), i === "supervised") {
+		if (r === "lora") {
+			if (n?.loraName) {
+				if (!t.hasLoRA(n.loraName)) if (n.loraConfig) t.createLoRA(n.loraName, n.loraConfig), t.attachLoRA(n.loraName);
+				else throw Error(`LoRA configuration must be provided to create LoRA with name ${n.loraName}`);
+				else if (t.attachLoRA(n.loraName), n.loraConfig) {
+					let e = t.lora;
+					(e.alpha !== n.loraConfig.alpha || e.rank !== n.loraConfig.rank) && (t.detachLoRA(), t.deleteLoRA(n.loraName), t.createLoRA(n.loraName, n.loraConfig), t.attachLoRA(n.loraName), console.warn("Resetting LoRA with new configuration."));
+				}
+			} else if (n?.loraConfig) if (t.hasLoRA()) {
+				let r = t.lora;
+				if (r.alpha !== n.loraConfig.alpha || r.rank !== n.loraConfig.rank) {
+					t.detachLoRA();
+					let r = n.loraName || e();
+					t.createLoRA(r, n.loraConfig), t.attachLoRA(r);
+				}
+			} else {
+				let r = n.loraName || e();
+				t.createLoRA(r, n.loraConfig), t.attachLoRA(r);
+			}
+			else if (!t.hasLoRA()) throw Error("LoRA configuration must be provided for lora SFT mode");
+		} else t.hasLoRA() && t.detachLoRA();
+		r === "last-layer" ? t.weightStore.setTrainable([`block_${t.config.nLayer - 1}_*`, "token_embedding"]) : r === "full" && t.weightStore.setTrainable(["*"]);
+	}
+	n?.trainableWeights && t.weightStore.setTrainable(n.trainableWeights);
+}
+//#endregion
+export { t as default };

package/dist/training/factory.d.ts

@@ -0,0 +1,6 @@
+import { ModelForwardAttributes, default as Model } from '../../models/model';
+import { TrainingOptions } from '../../training/types';
+import { default as BasicTrainer } from '../../training/BasicTrainer';
+import { ITokeniser } from '../../tokeniser/type';
+import { AdamWOptimizer } from './AdamW';
+export default function createTrainer(model: Model<ModelForwardAttributes>, tokenizer: ITokeniser, options?: TrainingOptions, optimizer?: AdamWOptimizer): BasicTrainer;

package/dist/training/factory.js

@@ -0,0 +1,8 @@
+import e from "./SFTTrainer.js";
+import t from "./PreTrainer.js";
+//#region lib/training/factory.ts
+function n(n, r, i, a) {
+	return i?.method.type === "supervised" ? new e(n, r, i, a) : new t(n, r, i, a);
+}
+//#endregion
+export { n as default };

package/dist/training/prepareData.d.ts

@@ -0,0 +1,22 @@
+import { default as Model, ModelForwardAttributes } from '../../models/model';
+import { TrainingOptions } from './types';
+import { TokenStore } from './tasks/TokenStore';
+import { DatasetMetadata } from '../../loader/types';
+import { ConversationStream } from '../../data/stream';
+import { ITokeniser } from '../../tokeniser/type';
+import { Dataset } from '@tensorflow/tfjs-data';
+import { Tensor } from '@tensorflow/tfjs-core';
+interface PrepareDataResult {
+    trainDataset: Dataset<{
+        xs: Tensor;
+        ys: Tensor;
+    }>;
+    validationDataset?: Dataset<{
+        xs: Tensor;
+        ys: Tensor;
+    }>;
+    totalTokens: number;
+}
+/** Take our training options, model, tokeniser, and tasks, and prepare the training and validation datasets in Tensorflow format. */
+export default function prepareData(options: TrainingOptions, model: Model<ModelForwardAttributes>, tokeniser: ITokeniser, tasks: ConversationStream[] | Uint16Array[] | TokenStore, validation?: Uint16Array[] | TokenStore, datasets?: DatasetMetadata[]): Promise<PrepareDataResult>;
+export {};

package/dist/training/prepareData.js

@@ -0,0 +1,49 @@
+import { TokenStore as e } from "./tasks/TokenStore.js";
+import { tokensFromStreams as t } from "./tasks/tokenStream.js";
+import { t as n } from "../DatasetBuilder-DU1G1OKX.js";
+import { createTrainValidationDatasets as r, storeFromArray as i } from "./validation.js";
+//#region lib/training/prepareData.ts
+async function a(a, o, s, c, l, u) {
+	let d = a.loraName || a.loraConfig;
+	if (u && d) throw Error("Cannot specify datasets when using LoRA fine-tuning");
+	if (!u && !d) throw Error("Must specify datasets for non-LoRA training");
+	if (u) {
+		let e = o.metaData.pretrainingData || [], t = [...e], n = !1;
+		for (let r of u) e.some((e) => e.id === r.id) || t.push({
+			id: r.id,
+			name: r.name,
+			conversational: r.conversational
+		}), r.conversational && (n = !0);
+		o.metaData.pretrainingData = t, n ? o.metaData.mode = "conversational" : o.metaData.mode !== "conversational" && (o.metaData.mode = "completion");
+	} else o.metaData.mode !== "conversational" && (o.metaData.mode = "completion");
+	let f = a.maskedLoss ?? a.method.type === "supervised", p, m = l;
+	if (Array.isArray(c)) if (c[0] instanceof Uint16Array) p = c;
+	else {
+		let e = await t(c, s, {
+			masking: f,
+			validationSplit: a.validationSplit
+		});
+		p = e.trainingTokens, l || (m = e.validationTokens);
+	}
+	else p = c;
+	let h = p instanceof e ? p.getTokenCount() : p.reduce((e, t) => e + t.length, 0);
+	a.epochSteps = Math.ceil(h / ((a?.batchSize || 32) * o.config.blockSize));
+	let g = new n(s, o.config.blockSize);
+	if (m) {
+		let { trainDataset: e, validationDataset: t } = await r(p, m, s, g, a?.batchSize || 32);
+		return {
+			trainDataset: e,
+			validationDataset: t,
+			totalTokens: h
+		};
+	} else {
+		let t = p instanceof e ? p : await i(p, s);
+		return {
+			trainDataset: (await g.createTextDataset(t, a)).dataset,
+			validationDataset: void 0,
+			totalTokens: h
+		};
+	}
+}
+//#endregion
+export { a as default };

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

@@ -31,7 +31,8 @@ export declare class TokenStore {
     getMask(index: number): Promise<Uint8Array | undefined>;
     getShardCount(): number;
     getTokenCount(): number;
-    appendShard(shard: Uint16Array, mask?: Uint8Array): Promise<void>;
+    finish(): Promise<void>;
+    appendShard(shard: Uint16Array, mask?: Uint8Array): void;
     dispose(): Promise<void>;
     clear(): Promise<void>;
 }

package/dist/training/tasks/TokenStore.js

@@ -7,12 +7,12 @@ var e = "llm-tokenstore", t = class {
 	masks;
 	shardCount = 0;
 	lastShardLength = -1;
-	_shardSize = 4e3 * 1024;
+	_shardSize = 8e3 * 1024;
 	dirHandle = null;
 	opfsAvailable = !1;
 	lru = /* @__PURE__ */ new Map();
 	maxCachedShards;
-	constructor(t, n, r, i = 8, a = 4e3 * 1024) {
+	constructor(t, n, r, i = 2, a = 8e3 * 1024) {
 		this.tokeniserId = t, this.datasetId = n, this.maxCachedShards = Math.max(1, Math.trunc(i)), this._shardSize = Math.max(1, Math.trunc(a)), this.name = r ?? e;
 	}
 	get shardSize() {
@@ -148,12 +148,15 @@ var e = "llm-tokenstore", t = class {
 	getTokenCount() {
 		return this.shardCount === 0 ? 0 : (this.shardCount - 1) * this._shardSize + this.lastShardLength;
 	}
-	async appendShard(e, t) {
+	async finish() {
+		await this.writeManifest();
+	}
+	appendShard(e, t) {
 		if (this.lastShardLength >= 0 && this.lastShardLength < this._shardSize) throw Error("Previous shard was not full");
 		let n = this.shards.length;
 		this.shards.push(e), !this.masks && t && (this.masks = Array(n).fill(null)), this.masks && this.masks.push(t ?? null), this.shardCount = this.shards.length, this.lastShardLength = e.length, this.touchLRU(n);
 		try {
-			this.opfsAvailable && this.dirHandle && (await this.writeShardToOPFS(n, e), t && await this.writeMaskToOPFS(n, t), await this.writeManifest());
+			this.opfsAvailable && this.dirHandle && (this.writeShardToOPFS(n, e), t && this.writeMaskToOPFS(n, t));
 		} catch (e) {
 			console.error(e), this.maxCachedShards = 2 ** 53 - 1;
 		}
@@ -208,7 +211,7 @@ async function a(e, r, i, a) {
 	let o = n.get(e);
 	if (o && o.tokeniserId === r && o.datasetId === i) return o;
 	o && (await o.dispose(), n.delete(e));
-	let s = new t(r, i, e, a?.maxCachedShards ?? 8, a?.shardSize ?? 4e3 * 1024);
+	let s = new t(r, i, e, a?.maxCachedShards, a?.shardSize);
 	return a?.noOPFS || await s.init(), n.set(e, s), s;
 }
 //#endregion

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

@@ -0,0 +1,17 @@
+import { ITokeniser } from '../../../tokeniser/type';
+import { ConversationStream } from '../../../data/stream';
+import { TokenStore } from './TokenStore';
+interface TokensFromTasksOptions {
+    masking?: boolean;
+    maxCachedShards?: number;
+    noOPFS?: boolean;
+    shardSize?: number;
+    validationSplit?: number;
+    validationSeed?: string | number;
+    cb?: (tokens: number) => void;
+}
+export declare function tokensFromStreams(tasks: ConversationStream[], tokenizer: ITokeniser, options?: TokensFromTasksOptions): Promise<{
+    trainingTokens: TokenStore;
+    validationTokens?: TokenStore;
+}>;
+export {};

package/dist/training/tasks/tokenStream.js

@@ -0,0 +1,46 @@
+import { createTokenStore as e, deleteTokenStore as t } from "./TokenStore.js";
+import { seededRng as n } from "../../utilities/random.js";
+//#region lib/training/tasks/tokenStream.ts
+function r(e, t, n, r, i, a) {
+	let o = n.encodeConversation(e, !1, !!a);
+	if (o) {
+		let e = Array.isArray(o) ? o : o.tokens;
+		r.total += e.length;
+		let n = t[t.length - 1], s = a ? a[a.length - 1] : null;
+		if (r.offset + e.length > n.length) {
+			let c = n.length - r.offset;
+			n.set(e.slice(0, c), r.offset);
+			let l = e.length - c;
+			if (l > i) throw Error(`Estimated tokens (${i}) is too small for the next batch of tokens (${l}).`);
+			let u = new Uint16Array(i);
+			if (u.set(e.slice(c), 0), t.push(u), a && s && !Array.isArray(o)) {
+				s.set(o.mask.slice(0, c).map((e) => +!!e), r.offset);
+				let e = new Uint8Array(u.length);
+				e.set(o.mask.slice(c).map((e) => +!!e), 0), a.push(e);
+			}
+			r.offset = e.length - c;
+		} else n.set(e, r.offset), s && !Array.isArray(o) && s.set(o.mask.map((e) => +!!e), r.offset), r.offset += e.length;
+	}
+}
+async function i(i, a, o) {
+	await t("training-tokens");
+	let s = await e("training-tokens", a.id, a.datasetID ?? "", o);
+	await t("validation-tokens");
+	let c = o?.validationSplit && o.validationSplit > 0 ? await e("validation-tokens", a.id, a.datasetID ?? "", o) : void 0, l = [new Uint16Array(s.shardSize)], u = o?.masking ? [new Uint8Array(s.shardSize)] : null, d = {
+		offset: 0,
+		total: 0
+	}, f = o?.validationSplit && o.validationSplit > 0 ? [new Uint16Array(c.shardSize)] : void 0, p = o?.masking && f ? [new Uint8Array(c.shardSize)] : null, m = {
+		offset: 0,
+		total: 0
+	}, h = 0, g = o?.cb, _ = o?.validationSeed === void 0 ? Math.random : n(o.validationSeed);
+	for (; h < i.length;) await i[h++].begin((e) => {
+		let t = o?.validationSplit && o.validationSplit > 0 && _() < o.validationSplit, n = t ? f : l, i = t ? m : d, h = t ? p : u, g = t ? c : s;
+		r(e, n, a, i, g.shardSize, h || void 0), n.length > 1 && (g.appendShard(n[0], h ? h[0] : void 0), n.shift(), h && h.shift());
+	}, g ? () => g(d.total) : void 0);
+	return l.length === 1 && (l[0] = l[0].subarray(0, d.offset), s.appendShard(l[0], u ? u[0].subarray(0, d.offset) : void 0)), f && f.length === 1 && (f[0] = f[0].subarray(0, m.offset), c.appendShard(f[0], p ? p[0].subarray(0, m.offset) : void 0)), await s.finish(), c && await c.finish(), {
+		trainingTokens: s,
+		validationTokens: f ? c : void 0
+	};
+}
+//#endregion
+export { i as tokensFromStreams };

package/dist/training/types.d.ts

@@ -1,6 +1,9 @@
 import { TensorStatistics } from '../../checks/weights';
 import { LoRAConfig } from '../../models/config';
 import { NamedTensorMap, Tensor } from '@tensorflow/tfjs-core';
+import { TokenStore } from './tasks/TokenStore';
+import { ConversationStream } from '../../data/stream';
+import { DatasetMetadata } from '../../loader/types';
 export interface Metrics {
     accuracy?: number;
     perplexity?: number;
@@ -43,7 +46,13 @@ export interface AdamWOptimizerConfig extends LRSchedulerConfig {
     orthoGrad?: boolean;
 }
 export type TrainingMetrics = 'accuracy' | 'perplexity' | 'gradientNorm' | 'gradientStatistics' | 'weightNorm' | 'weightStatistics' | 'memoryUsage' | 'tokensPerSecond' | 'learningRate';
+export interface TrainingMethod {
+    type: 'pretraining' | 'supervised';
+    supervised?: 'full' | 'lora' | 'last-layer';
+}
 export interface TrainingOptions extends Partial<AdamWOptimizerConfig> {
+    previous_job_id?: string;
+    method: TrainingMethod;
     batchSize: number;
     maxEpochs?: number;
     logInterval?: number;
@@ -54,7 +63,6 @@ export interface TrainingOptions extends Partial<AdamWOptimizerConfig> {
     trainableWeights?: string[];
     loraConfig?: LoRAConfig;
     loraName?: string;
-    sftMode: 'full' | 'lora' | 'last-layer';
     maskedLoss?: boolean;
     metrics?: TrainingMetrics[];
     contextScaling?: number;
@@ -62,6 +70,11 @@ export interface TrainingOptions extends Partial<AdamWOptimizerConfig> {
     dropout?: number;
     layerDrop?: number;
     debug?: boolean;
+}
+export interface TrainingContext {
+    training_data: ConversationStream[] | Uint16Array[] | TokenStore;
+    validation_data?: Uint16Array[] | TokenStore;
+    datasets?: DatasetMetadata[];
     onStep?: (log: TrainingLogEntry) => void;
 }
 export interface TrainingState {

package/dist/training/validateOptions.d.ts

@@ -0,0 +1,2 @@
+import { TrainingOptions } from './types';
+export default function validateOptions(options: TrainingOptions, training: boolean, oldOptions?: TrainingOptions): TrainingOptions;

package/dist/training/validateOptions.js

@@ -0,0 +1,19 @@
+//#region lib/training/validateOptions.ts
+function e(e, t, n) {
+	let r = new Set(Object.keys(e).filter((t) => e[t] !== n?.[t]));
+	if (t) {
+		if (r.has("batchSize")) throw Error("Cannot change batch size during training");
+		if (r.has("sftMode")) throw Error("Cannot change SFT mode during training");
+		if (r.has("loraConfig")) throw Error("Cannot change LoRA configuration during training");
+		if (r.has("validationSplit")) throw Error("Cannot change validation split during training");
+		if (r.has("trainableWeights")) throw Error("Cannot change trainable weights during training");
+		if (r.has("mixedPrecision")) throw Error("Cannot change mixed precision setting during training");
+		if (r.has("gradientCheckpointing")) throw Error("Cannot change gradient checkpointing setting during training");
+	}
+	return {
+		...n,
+		...e
+	};
+}
+//#endregion
+export { e as default };

package/dist/training/validation.js

@@ -1,8 +1,10 @@
 import { TokenStore as e, createTokenStore as t } from "./tasks/TokenStore.js";
 //#region lib/training/validation.ts
 async function n(e, n) {
-	let r = await t("training-tokens", n.id, n.datasetID ?? ""), i = e.map((e) => r.appendShard(e));
-	return await Promise.all(i), r;
+	let r = await t("training-tokens", n.id, n.datasetID ?? "");
+	return e.forEach((e) => {
+		r.appendShard(e);
+	}), await r.finish(), r;
 }
 async function r(t, r, i, a, o) {
 	let s = t instanceof e ? t : await n(t, i), c = r instanceof e ? r : await n(r, i), l = s.getTokenCount(), { dataset: u, state: d } = await a.createTextDataset(s, { batchSize: o }), { dataset: f, state: p } = await a.createTextDataset(c, {

package/dist/utilities/arrayShape.d.ts

@@ -0,0 +1 @@
+export default function arrayShape(arr: unknown[]): number[];

package/dist/utilities/arrayShape.js

@@ -0,0 +1,8 @@
+//#region lib/utilities/arrayShape.ts
+function e(e) {
+	let t = [], n = e;
+	for (; Array.isArray(n);) t.push(n.length), n = n[n.length - 1];
+	return t;
+}
+//#endregion
+export { e as default };

package/dist/utilities/random.d.ts

@@ -0,0 +1 @@
+export declare function seededRng(seed: string | number): () => number;

package/dist/utilities/random.js

@@ -0,0 +1,19 @@
+//#region lib/utilities/random.ts
+function e(e) {
+	let t = 1779033703 ^ e.length;
+	for (let n = 0; n < e.length; n++) t = Math.imul(t ^ e.charCodeAt(n), 3432918353), t = t << 13 | t >>> 19;
+	return () => (t = Math.imul(t ^ t >>> 16, 2246822507), t = Math.imul(t ^ t >>> 13, 3266489909), (t ^= t >>> 16) >>> 0);
+}
+function t(e) {
+	let t = e >>> 0;
+	return () => {
+		t += 1831565813;
+		let e = Math.imul(t ^ t >>> 15, t | 1);
+		return e ^= e + Math.imul(e ^ e >>> 7, e | 61), ((e ^ e >>> 14) >>> 0) / 4294967296;
+	};
+}
+function n(n) {
+	return t(typeof n == "number" ? n >>> 0 : e(n)());
+}
+//#endregion
+export { n as seededRng };

package/dist/utilities/waitForModel.d.ts

@@ -1,2 +1,2 @@
-import { TeachableLLM } from '../../main';
+import { default as TeachableLLM } from '../../TeachableLLM';
 export default function waitForModel(model: TeachableLLM): Promise<void>;

package/dist/v4-BK7K-jy_.js

@@ -0,0 +1,30 @@
+//#region node_modules/uuid/dist/stringify.js
+var e = [];
+for (let t = 0; t < 256; ++t) e.push((t + 256).toString(16).slice(1));
+function t(t, n = 0) {
+	return (e[t[n + 0]] + e[t[n + 1]] + e[t[n + 2]] + e[t[n + 3]] + "-" + e[t[n + 4]] + e[t[n + 5]] + "-" + e[t[n + 6]] + e[t[n + 7]] + "-" + e[t[n + 8]] + e[t[n + 9]] + "-" + e[t[n + 10]] + e[t[n + 11]] + e[t[n + 12]] + e[t[n + 13]] + e[t[n + 14]] + e[t[n + 15]]).toLowerCase();
+}
+//#endregion
+//#region node_modules/uuid/dist/rng.js
+var n = new Uint8Array(16);
+function r() {
+	return crypto.getRandomValues(n);
+}
+//#endregion
+//#region node_modules/uuid/dist/v4.js
+function i(e, t, n) {
+	return !t && !e && crypto.randomUUID ? crypto.randomUUID() : a(e, t, n);
+}
+function a(e, n, i) {
+	e ||= {};
+	let a = e.random ?? e.rng?.() ?? r();
+	if (a.length < 16) throw Error("Random bytes length must be >= 16");
+	if (a[6] = a[6] & 15 | 64, a[8] = a[8] & 63 | 128, n) {
+		if (i ||= 0, i < 0 || i + 16 > n.length) throw RangeError(`UUID byte range ${i}:${i + 15} is out of buffer bounds`);
+		for (let e = 0; e < 16; ++e) n[i + e] = a[e];
+		return n;
+	}
+	return t(a);
+}
+//#endregion
+export { i as t };

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@genai-fi/nanogpt",
-    "version": "0.23.0",
+    "version": "1.0.0",
     "type": "module",
     "main": "dist/main.js",
     "types": "dist/main.d.ts",
@@ -29,7 +29,7 @@
     "devDependencies": {
         "@eslint/js": "^10.0.1",
         "@playwright/test": "^1.55.0",
-        "@types/node": "^22.13.14",
+        "@types/node": "^26.2.0",
         "@types/papaparse": "^5.3.16",
         "@types/picomatch": "^4.0.2",
         "@types/yargs": "^17.0.33",
@@ -60,5 +60,11 @@
         "pdfjs-dist": "^5.4.149",
         "picomatch": "^4.0.3",
         "uuid": "^14.0.0"
+    },
+    "allowScripts": {
+        "webgpu@0.4.0": true,
+        "esbuild@0.28.1": true,
+        "core-js@3.29.1": true,
+        "patch-package@8.0.1": true
     }
 }

package/dist/Generator.js

@@ -1,2 +0,0 @@
-import { a as e, o as t } from "./main-BSaDGH7I.js";
-export { e as default, t as isConversation };