@genai-fi/nanogpt 0.23.0 → 1.0.0

package/dist/{matMulGelu-BAIgQaRx.js → matMulGelu-CBoqTZM7.js}

@@ -1,6 +1,6 @@
 import { Ci as e, Ii as t, In as n, Ps as r, di as i, gr as a, ii as o, oc as s } from "./dist-Da20xy8E.js";
-import { a as c } from "./gpgpu_math-DBYEAAdI.js";
-import { t as l } from "./Reshape-dnm9bO3B.js";
+import { a as c } from "./gpgpu_math-20tPK8LM.js";
+import { t as l } from "./Reshape-BlylqwWy.js";
 //#region node_modules/@tensorflow/tfjs-backend-webgl/dist/mulmat_packed_gpu.js
 var u = class {
 	constructor(e, t, n, r = !1, i = !1, a = !1, o = null, s = !1, l = !1) {

package/dist/models/NanoGPTV1.js

@@ -1,2 +1,95 @@
-import { _ as e } from "../main-BSaDGH7I.js";
-export { e as default };
+import { di as e, oi as t } from "../dist-Da20xy8E.js";
+import { packingSupported as n } from "../utilities/packed.js";
+import { r, t as i } from "../pack16-BhuXNUS7.js";
+import a from "../layers/RMSNorm.js";
+import o from "../layers/TransformerBlock.js";
+import s from "../layers/TiedEmbedding.js";
+import c from "../layers/RoPECache.js";
+import l from "./model.js";
+import u from "../layers/PositionEmbedding.js";
+//#region lib/models/NanoGPTV1.ts
+var d = {
+	modelType: "GenAI_NanoGPT_v1",
+	vocabSize: 2e3,
+	blockSize: 128,
+	nLayer: 6,
+	nHead: 4,
+	nEmbed: 256,
+	mlpFactor: 4,
+	useRope: !0
+}, f = class extends l {
+	wte;
+	wpe;
+	blocks;
+	lnF;
+	ropeCache;
+	constructor(e = {}) {
+		super({
+			...d,
+			...e
+		});
+		let t = {
+			activation: "gelu",
+			hiddenFactor: this.config.mlpFactor,
+			useGamma: !0,
+			useQKNorm: !1
+		};
+		this.wte = new s(this.config, "token_embedding", this), this.config.useRope === !1 ? this.wpe = new u(this.config, "positional_embedding", this) : this.ropeCache = new c(this.config), this.blocks = [];
+		for (let e = 0; e < this.config.nLayer; e++) this.blocks.push(new o(e, this.config, t, this));
+		this.lnF = new a(this.config, t, "final_rms_norm", this);
+	}
+	getClassName() {
+		return "GenAI_NanoGPT_v1";
+	}
+	inputPhase(t, n) {
+		return e(() => {
+			let e = this.wte.embed(t);
+			if (this.config.useRope === !1) {
+				let t = this.wpe.call(n, e);
+				if (Array.isArray(t)) throw Error("PositionEmbedding output should not be an array");
+				return t;
+			}
+			return e;
+		});
+	}
+	forward(a, o) {
+		return this.validateInput(o), a.ropeCache = this.ropeCache, a.outputEmbeddings && (a.embeddings = []), e(() => {
+			this.startMemory();
+			let e = this.inputPhase(o, a);
+			if (a.cache && a.cache.length !== this.blocks.length) throw console.error("Cache", a.cache), Error(`Cache length ${a.cache.length} does not match number of blocks ${this.blocks.length}`);
+			let s = a.mixedPrecision === !0 && n(), c = s ? i(e) : e;
+			s && e !== c && e.dispose();
+			for (let e = 0; e < this.blocks.length; e++) {
+				let n = this.blocks[e], r = Math.random() * 1e9, i = {
+					...a,
+					seed: r,
+					pastKV: a.cache ? a.cache[e] : void 0,
+					mixedPrecision: s
+				}, o = a.checkpointing && a.training ? n.callCheckpoint(i, c) : n.call(i, c);
+				a.outputEmbeddings ? (t(c), a.embeddings.push({
+					name: `block_output_${e}`,
+					tensor: c
+				})) : c.dispose(), c = o;
+			}
+			if (e = this.lnF.call({
+				...a,
+				mixedPrecision: s
+			}, c), c.dispose(), a.skipLogits) return this.endMemory("Forward"), e;
+			let l = this.wte.project(e);
+			a.outputEmbeddings ? (t(e), a.embeddings.push({
+				name: "final_norm_output",
+				tensor: e
+			})) : e.dispose();
+			let u = s ? r(l) : l;
+			return s && l !== u && l.dispose(), u;
+		});
+	}
+	project(t) {
+		return e(() => this.wte.project(t));
+	}
+	dispose() {
+		this.weightStore.dispose(), this.wte.dispose(), this.wpe && this.wpe.dispose(), this.blocks.forEach((e) => e.dispose()), this.lnF.dispose();
+	}
+};
+//#endregion
+export { f as default };

package/dist/models/NanoGPTV2.js

@@ -1,2 +1,86 @@
-import { g as e } from "../main-BSaDGH7I.js";
-export { e as default };
+import { di as e, oi as t } from "../dist-Da20xy8E.js";
+import { packingSupported as n } from "../utilities/packed.js";
+import { r, t as i } from "../pack16-BhuXNUS7.js";
+import a from "../layers/RMSNorm.js";
+import o from "../layers/TransformerBlock.js";
+import s from "../layers/TiedEmbedding.js";
+import c from "../layers/RoPECache.js";
+import l from "./model.js";
+//#region lib/models/NanoGPTV2.ts
+var u = {
+	modelType: "GenAI_NanoGPT_v2",
+	vocabSize: 2e3,
+	blockSize: 128,
+	nLayer: 6,
+	nHead: 4,
+	nEmbed: 256,
+	mlpFactor: 4
+}, d = class extends l {
+	wte;
+	wpe;
+	blocks;
+	lnF;
+	ropeCache;
+	constructor(e = {}) {
+		super({
+			...u,
+			...e
+		});
+		let t = {
+			activation: "relu2",
+			hiddenFactor: this.config.mlpFactor,
+			useGamma: !1,
+			useQKNorm: !0
+		};
+		this.wte = new s(this.config, "token_embedding", this), this.ropeCache = new c(this.config), this.blocks = [];
+		for (let e = 0; e < this.config.nLayer; e++) this.blocks.push(new o(e, this.config, t, this));
+		this.lnF = new a(this.config, t, "final_rms_norm", this);
+	}
+	getClassName() {
+		return "GenAI_NanoGPT_v2";
+	}
+	inputPhase(t) {
+		return e(() => this.wte.embed(t));
+	}
+	forward(a, o) {
+		return this.validateInput(o), a.ropeCache = this.ropeCache, a.outputEmbeddings && (a.embeddings = []), e(() => {
+			this.startMemory();
+			let e = this.inputPhase(o);
+			if (a.cache && a.cache.length !== this.blocks.length) throw console.error("Cache", a.cache), Error(`Cache length ${a.cache.length} does not match number of blocks ${this.blocks.length}`);
+			let s = a.mixedPrecision === !0 && n(), c = s ? i(e) : e;
+			s && e !== c && e.dispose();
+			for (let e = 0; e < this.blocks.length; e++) {
+				if (a.layerDrop && Math.random() < a.layerDrop * (e / this.blocks.length)) continue;
+				let n = this.blocks[e], r = Math.random() * 1e9, i = {
+					...a,
+					seed: r,
+					pastKV: a.cache ? a.cache[e] : void 0,
+					mixedPrecision: s
+				}, o = a.checkpointing && a.training ? n.callCheckpoint(i, c) : n.call(i, c);
+				a.outputEmbeddings ? (t(c), a.embeddings.push({
+					name: `block_output_${e}`,
+					tensor: c
+				})) : c.dispose(), c = o;
+			}
+			if (e = this.lnF.call({
+				...a,
+				mixedPrecision: s
+			}, c), c.dispose(), a.skipLogits) return this.endMemory("Forward"), e;
+			let l = this.wte.project(e);
+			a.outputEmbeddings ? (t(e), a.embeddings.push({
+				name: "final_norm_output",
+				tensor: e
+			})) : e.dispose();
+			let u = s ? r(l) : l;
+			return s && l !== u && l.dispose(), u;
+		});
+	}
+	project(t) {
+		return e(() => this.wte.project(t));
+	}
+	dispose() {
+		this.weightStore.dispose(), this.wte.dispose(), this.wpe && this.wpe.dispose(), this.blocks.forEach((e) => e.dispose()), this.lnF.dispose();
+	}
+};
+//#endregion
+export { d as default };

package/dist/models/factory.js

@@ -1,2 +1,13 @@
-import { h as e } from "../main-BSaDGH7I.js";
-export { e as default };
+import e from "./NanoGPTV1.js";
+import t from "./NanoGPTV2.js";
+//#region lib/models/factory.ts
+function n(n) {
+	let r = n.modelType || "GenAI_NanoGPT_v1";
+	switch (r) {
+		case "GenAI_NanoGPT_v1": return new e(n);
+		case "GenAI_NanoGPT_v2": return new t(n);
+		default: throw Error(`Unsupported model type: ${r}`);
+	}
+}
+//#endregion
+export { n as default };

package/dist/models/model.js

@@ -1,2 +1,76 @@
-import { v as e } from "../main-BSaDGH7I.js";
-export { e as default };
+import e from "../layers/BaseLayer.js";
+import { estimateParameterCount as t } from "../utilities/parameters.js";
+import n from "../layers/LoRA.js";
+//#region lib/models/model.ts
+var r = class extends e {
+	lossScaling = 128;
+	trainingState = null;
+	metaData = {
+		version: 2,
+		application: "@genai-fi/nanogpt"
+	};
+	loraLayer;
+	loraMap = /* @__PURE__ */ new Map();
+	constructor(e) {
+		super(e), e.loraConfig && e.loraConfig.forEach((e, t) => {
+			this.createLoRA(t, e);
+		});
+	}
+	createLoRA(e, t) {
+		if (this.loraMap.has(e)) return;
+		let r = new n(e, this.weightStore, t.alpha, t.rank, t.variables);
+		this.loraMap.set(e, r), this.config.loraConfig = this.config.loraConfig || /* @__PURE__ */ new Map(), this.config.loraConfig.set(e, t);
+	}
+	deleteLoRA(e) {
+		let t = this.loraMap.get(e);
+		if (!t) throw Error(`No LoRA with name ${e} exists.`);
+		this.loraLayer === t && this.detachLoRA(), t.dispose(), this.loraMap.delete(e), this.config.loraConfig && this.config.loraConfig.delete(e);
+	}
+	renameLoRA(e, t) {
+		if (!this.loraMap.has(e)) throw Error(`No LoRA with name ${e} exists.`);
+		if (this.loraMap.has(t)) throw Error(`LoRA with name ${t} already exists.`);
+		let n = this.loraMap.get(e);
+		this.loraMap.set(t, n), this.loraMap.delete(e), this.config.loraConfig && (this.config.loraConfig.delete(e), this.config.loraConfig.set(t, {
+			rank: n.rank,
+			alpha: n.alpha,
+			variables: Array.from(n.variables)
+		}));
+	}
+	mergeLoRA(e) {
+		let t = this.loraMap.get(e);
+		if (!t) throw Error(`No LoRA with name ${e} exists.`);
+		t.merge(), this.deleteLoRA(e);
+	}
+	attachLoRA(e) {
+		if (this.loraLayer) {
+			if (this.loraLayer.name === e) return;
+			this.detachLoRA();
+		}
+		let t = this.loraMap.get(e);
+		if (!t) throw Error(`No LoRA with name ${e} exists.`);
+		t.attach(), this.loraLayer = t, this.config.loraName = e;
+	}
+	detachLoRA() {
+		if (!this.loraLayer) throw Error("No LoRA layer is attached to this model.");
+		this.loraLayer.detach(), this.loraLayer = void 0, this.config.loraName = void 0;
+	}
+	hasLoRA(e) {
+		return e ? this.loraMap.has(e) : !!this.loraLayer;
+	}
+	listLoRAs() {
+		return Array.from(this.loraMap.keys());
+	}
+	get lora() {
+		return this.loraLayer || null;
+	}
+	getNumParams() {
+		return t(this.config);
+	}
+	validateInput(e) {
+		if (e.shape.length !== 2) throw Error(`Invalid input shape: expected [batch_size, sequence_length], got ${e.shape}`);
+		if (e.shape[1] > this.config.blockSize) throw Error(`Input sequence length ${e.shape[1]} isn't block size ${this.config.blockSize}`);
+		if (e.dtype !== "int32") throw Error(`Input tensor must be of type int32, got ${e.dtype}`);
+	}
+};
+//#endregion
+export { r as default };

package/dist/models.d.ts

@@ -0,0 +1,4 @@
+export { default as NanoGPTV1 } from './models/NanoGPTV1';
+export { default as NanoGPTV2 } from './models/NanoGPTV2';
+export { default as waitForModel } from './utilities/waitForModel';
+export { default as Model } from './models/model';

package/dist/models.js

@@ -0,0 +1,14 @@
+import { n as e } from "./chunk-CWhphoD1.js";
+import t from "./models/model.js";
+import n from "./models/NanoGPTV1.js";
+import r from "./models/NanoGPTV2.js";
+import i from "./utilities/waitForModel.js";
+//#region lib/models.ts
+var a = /* @__PURE__ */ e({
+	Model: () => t,
+	NanoGPTV1: () => n,
+	NanoGPTV2: () => r,
+	waitForModel: () => i
+});
+//#endregion
+export { t as Model, n as NanoGPTV1, r as NanoGPTV2, a as t, i as waitForModel };

package/dist/ops/dot16.js

@@ -1,4 +1,4 @@
-import { o as e } from "../tfjs_backend-CydPRQTc.js";
+import { o as e } from "../tfjs_backend-h5weiy1O.js";
 import { isPackedTensor as t } from "../utilities/packed.js";
 import { transpose16 as n } from "./transpose16.js";
 import { reshape16 as r } from "./reshape16.js";

package/dist/ops/matMulGelu.js

@@ -1,5 +1,5 @@
 import { ii as e } from "../dist-Da20xy8E.js";
-import "../matMulGelu-BAIgQaRx.js";
+import "../matMulGelu-CBoqTZM7.js";
 import "./cpu/matMulGelu.js";
 import "./grads/matMulGelu.js";
 //#region lib/ops/matMulGelu.ts

package/dist/ops/webgl/adamAdjust.js

@@ -1,5 +1,5 @@
 import { Ii as e } from "../../dist-Da20xy8E.js";
-import { t } from "../../Reshape-dnm9bO3B.js";
+import { t } from "../../Reshape-BlylqwWy.js";
 //#region lib/ops/webgl/adamAdjust.ts
 var n = class {
 	variableNames = ["moments", "value"];

package/dist/ops/webgl/fusedSoftmax.js

@@ -1,6 +1,6 @@
 import { Ii as e, Tn as t, nc as n } from "../../dist-Da20xy8E.js";
-import { t as r } from "../../Reshape-dnm9bO3B.js";
-import { a as i, r as a, t as o } from "../../RealDiv-CNsvC4AU.js";
+import { t as r } from "../../Reshape-BlylqwWy.js";
+import { a as i, r as a, t as o } from "../../RealDiv-CSnvtN2E.js";
 //#region lib/ops/webgl/fusedSoftmax.ts
 var s = class {
 	variableNames = ["logits", "maxLogits"];

package/dist/ops/webgl/gelu.js

@@ -1,7 +1,7 @@
 import { Ii as e } from "../../dist-Da20xy8E.js";
-import { _ as t, i as n } from "../../kernel_funcs_utils-D-mATnGR.js";
+import { i as t, s as n } from "../../kernel_funcs_utils-ql8Y8qPn.js";
 //#region lib/ops/webgl/gelu.ts
-var r = .7978845608028654, i = .044715, a = n({ opSnippet: t + `
+var r = .7978845608028654, i = .044715, a = t({ opSnippet: n + `
     float x3 = x * x * x;
     float inner = x + ${i} * x3;
     inner = ${r} * inner;

package/dist/ops/webgl/log.js

@@ -1,14 +1,14 @@
 import { Ii as e } from "../../dist-Da20xy8E.js";
-import { y as t } from "../../shared-CfTzpULd.js";
-import { i as n, t as r } from "../../kernel_funcs_utils-D-mATnGR.js";
+import { i as t, t as n } from "../../kernel_funcs_utils-ql8Y8qPn.js";
+import { y as r } from "../../shared-CfTzpULd.js";
 //#region lib/ops/webgl/log.ts
 e({
 	kernelName: "Log",
 	backendName: "webgl",
-	kernelFunc: n({
-		opSnippet: r + "\n  return x < 0.0 ? NAN : log(x);\n",
+	kernelFunc: t({
+		opSnippet: n + "\n  return x < 0.0 ? NAN : log(x);\n",
 		packedOpSnippet: "\n  vec4 result = log(x);\n  bvec4 isNaN = isnan(x);\n  result.r = isNaN.r ? x.r : (x.r < 0.0 ? NAN : result.r);\n  result.g = isNaN.g ? x.g : (x.g < 0.0 ? NAN : result.g);\n  result.b = isNaN.b ? x.b : (x.b < 0.0 ? NAN : result.b);\n  result.a = isNaN.a ? x.a : (x.a < 0.0 ? NAN : result.a);\n  return result;\n",
-		cpuKernelImpl: t
+		cpuKernelImpl: r
 	})
 });
 //#endregion

package/dist/ops/webgl/matMulGelu.js

@@ -1,2 +1,2 @@
-import { n as e, r as t, t as n } from "../../matMulGelu-BAIgQaRx.js";
+import { n as e, r as t, t as n } from "../../matMulGelu-CBoqTZM7.js";
 export { n as MATMUL_SHARED_DIM_THRESHOLD, e as batchMatMulGeluImpl, t as batchMatMulKernel };

package/dist/ops/webgl/matMulMul.js

@@ -1,5 +1,5 @@
 import { Ii as e } from "../../dist-Da20xy8E.js";
-import { n as t } from "../../matMulGelu-BAIgQaRx.js";
+import { n as t } from "../../matMulGelu-CBoqTZM7.js";
 //#region lib/ops/webgl/matMulMul.ts
 var n = "\n    return a * b;\n";
 function r(e) {