@genai-fi/nanogpt 1.4.1 → 1.5.0

package/dist/api/beamer.d.ts

@@ -20,6 +20,7 @@ export default class BeamAPI {
     off<E extends keyof BeamEvents>(event: E, listener: BeamEvents[E]): void;
     cancel(id?: string): void;
     private startJob;
+    more(id: string): string | undefined;
     create(conversation: Conversation[], options: BeamerOptions): string;
 }
 export {};

package/dist/api/beamer.js

@@ -25,10 +25,10 @@ var r = class {
 		} else for (let e of this._jobs.values()) e.canceled = !0, e.beamer.cancel();
 	}
 	startJob(e) {
-		this._busyCount++, this.ee.emit("status", "busy"), e.beamer.beam(e.conversation, e.options, (t) => {
+		this._busyCount++, this.ee.emit("status", "busy"), e.beamer.beam((t) => {
 			this.ee.emit("progress", e.id, t);
 		}).then((t) => {
-			this._busyCount--, this._busyCount === 0 && this.ee.emit("status", "ready"), this.ee.emit("done", e.id, t), this._jobs.delete(e.id);
+			this._busyCount--, this._busyCount === 0 && this.ee.emit("status", "ready"), this.ee.emit("done", e.id, t);
 			let n = this._queue.shift();
 			n && this.startJob(n);
 		}).catch((t) => {
@@ -37,15 +37,25 @@ var r = class {
 			n && this.startJob(n);
 		});
 	}
+	more(e) {
+		let t = this._jobs.get(e);
+		if (t) {
+			if (this._busyCount > 0) return this._queue.push(t), e;
+			this.startJob(t);
+		}
+	}
 	create(e, r) {
-		let i = t(), a = {
+		let i = t(), a = new n(this._model, this._tokeniser);
+		a.begin(e, r);
+		let o = {
 			id: i,
-			beamer: new n(this._model, this._tokeniser),
+			beamer: a,
 			canceled: !1,
 			conversation: e,
 			options: r
 		};
-		return this._jobs.set(i, a), this._busyCount > 0 ? (this._queue.push(a), i) : (this.startJob(a), i);
+		for (let [e, t] of this._jobs.entries()) t.beamer.dispose(), this._jobs.delete(e);
+		return this._jobs.set(i, o), this._busyCount > 0 ? (this._queue.push(o), i) : (this.startJob(o), i);
 	}
 };
 //#endregion

package/dist/inference/Beamer.d.ts

@@ -6,13 +6,25 @@ export default class Beamer {
     private readonly tokeniser;
     private actualTokeniser;
     private active;
+    private options;
+    private conversation;
+    private heap;
+    private hasRoot;
+    private totalCandidates;
+    private seen;
     constructor(model: Model<ModelForwardAttributes>, tokeniser: ITokeniser);
     private shouldTerminate;
     private initialise;
     private tokeniseConversation;
     private createInitialContext;
     private appendTokenToContext;
-    private batchedNextProbabilities;
     cancel(): void;
-    beam(conversation: Conversation[], options: BeamerOptions, onStep?: (beams: IBeam[]) => void): Promise<IBeam[]>;
+    begin(conversation: Conversation[], options: BeamerOptions): void;
+    private expand;
+    private pushNextChild;
+    private entryToBeam;
+    next(): Promise<IBeam | null>;
+    private _next;
+    beam(onStep?: (beams: IBeam[]) => void): Promise<IBeam[]>;
+    dispose(): void;
 }

package/dist/inference/Beamer.js

@@ -2,18 +2,26 @@ import e from "../tokeniser/CharTokeniser.js";
 import { L as t, V as n, Y as r, _n as i, _r as a, di as o, xt as s } from "../dist-Da20xy8E.js";
 import c from "../utilities/topP.js";
 import { CHARS as l, padArray as u } from "./utilities.js";
+import d from "../utilities/maxHeap.js";
 //#region lib/inference/Beamer.ts
-function d(e, t) {
-	return e.map((e, t) => ({
+function f(e, t) {
+	let n = e.map((e, t) => ({
 		token: t,
 		prob: e
-	})).filter((e) => e.prob > 0).sort((e, t) => t.prob - e.prob).slice(0, Math.max(1, t));
+	})).filter((e) => e.prob > 0).sort((e, t) => t.prob - e.prob);
+	return t === void 0 ? n : n.slice(0, Math.max(1, t));
 }
-var f = class {
+var p = class {
 	model;
 	tokeniser;
 	actualTokeniser;
 	active = !1;
+	options = null;
+	conversation = [];
+	heap = new d((e) => e.score);
+	hasRoot = !1;
+	totalCandidates = 0;
+	seen = /* @__PURE__ */ new Set();
 	constructor(e, t) {
 		this.model = e, this.tokeniser = t, this.actualTokeniser = t;
 	}
@@ -51,100 +59,150 @@ var f = class {
 			contextLength: c
 		};
 	}
-	async batchedNextProbabilities(e, t, i) {
-		if (e.length === 0) return [];
-		let a = e[0], l = e.slice();
-		for (; l.length < t;) l.push(a);
-		let u = n(l.map((e) => e.context)), d = this.model.forward({
+	cancel() {
+		this.active = !1;
+	}
+	begin(e, t) {
+		if (!(!t || t.beams < 1 || t.maxBeamLength < 1)) {
+			if (this.options) throw Error("Beamer has already been initialized.");
+			this.conversation = e, this.options = t, this.initialise(t);
+		}
+	}
+	async expand(e, t) {
+		if (e.length === 0) return;
+		if (e.length > t) throw Error(`Number of nodes (${e.length}) exceeds fixed batch size (${t}).`);
+		let i = [...e];
+		if (e.length < t) for (; i.length < t;) i.push(e[e.length - 1]);
+		let a = n(i.map((e) => e.context)), l = this.model.forward({
 			training: !1,
 			mixedPrecision: !0
-		}, u), f = l.map((e) => Math.max(0, e.contextLength - 1)), p = o(() => {
-			let e = s(f, this.model.config.blockSize).expandDims(2);
-			return r(d.mul(e).sum(1));
-		}), m = e.length, h = m < t ? p.slice([0, 0], [m, this.model.config.vocabSize]) : p, g = await h.array();
-		h !== p && h.dispose(), p.dispose(), d.dispose(), u.dispose();
-		let _ = i.topP ?? 1;
-		return g.map((e) => c(e, _));
+		}, a), u = i.map((e) => Math.max(0, e.contextLength - 1)), d = o(() => {
+			let e = s(u, this.model.config.blockSize).expandDims(2);
+			return r(l.mul(e).sum(1));
+		}), p = e.length, m = p < t ? d.slice([0, 0], [p, this.model.config.vocabSize]) : d, h = await m.array();
+		m !== d && m.dispose(), d.dispose(), l.dispose(), a.dispose();
+		let g = this.options?.topP ?? 1, _ = h.map((e) => f(c(e, g), this.options?.topK));
+		for (let t = 0; t < e.length; t++) e[t].sortedChildren = _[t], this.totalCandidates += _[t].length - 1;
 	}
-	cancel() {
-		this.active = !1;
+	pushNextChild(e) {
+		if (!e.sortedChildren || e.nextChild >= e.sortedChildren.length) return;
+		let t = e.nextChild++, n = e.sortedChildren[t].token, r = e.sortedChildren[t].prob, i = this.options.maxBeamLength, a = this.shouldTerminate(this.options.allowSpecial ?? !1, n), o = this.actualTokeniser.decode([n]), s = e.tokens.length + 1 > i, c = this.options.endOnWhiteSpace === !0 && s && /\s/.test(o), l = this.options.endOnWhiteSpace ? e.tokens.length + 1 >= i + 10 : e.tokens.length + 1 >= i;
+		this.heap.push({
+			score: e.score * r,
+			parent: e,
+			idx: t,
+			tokenId: n,
+			tokenText: o,
+			probability: r,
+			terminated: a || c || l,
+			reason: a ? "token" : c ? "whitespace" : l ? "length" : void 0
+		});
 	}
-	async beam(e, t, n) {
-		if (!t || t.beams < 1 || t.maxBeamLength < 1) return [];
-		let r = Math.max(8, t.beams);
-		this.initialise(t), this.active = !0;
-		let i = await this.tokeniseConversation(e), a = this.createInitialContext(i), o = [{
-			tokens: [],
-			score: 1,
-			text: "",
-			terminated: !1,
-			context: a.context,
-			contextLength: a.contextLength,
-			candidates: 0
-		}], s = t.maxBeamLength, c = t.endOnWhiteSpace === !0 ? Math.max(s, t.maxLength ?? s + this.model.config.blockSize) : s, l = Date.now(), u = 0;
-		for (let e = 0; e < c && this.active; e++) {
-			let e = o.filter((e) => !e.terminated);
-			if (e.length === 0) break;
-			let i = await this.batchedNextProbabilities(e, r, t), a = [];
-			for (let n = 0; n < e.length; n++) {
-				let o = e[n], c = i[n], l = d(c, t.topK ? Math.min(t.topK, r) : Math.max(1, r));
-				for (let e of l) {
-					let n = this.actualTokeniser.decode([e.token]), r = this.shouldTerminate(t.allowSpecial ?? !1, e.token), i = o.tokens.length + 1 > s, c = t.endOnWhiteSpace === !0 && i && /\s/.test(n), l = c && o.tokens.length > 0 && /^\s/.test(n);
-					a.push({
-						parent: o,
-						token: e.token,
-						tokenText: n,
-						score: o.score * e.prob,
-						terminatedByToken: r,
-						terminatedByWhitespace: c,
-						dropLeadingWhitespaceTerminator: l
-					});
-				}
+	entryToBeam(e) {
+		let t = e.reason === "whitespace" && e.parent.tokens.length > 0 && /^\s/.test(e.tokenText), n = e.parent.tokens.map((t, n) => ({
+			token: t,
+			score: e.parent.probabilities[n],
+			text: this.actualTokeniser.decode([t]),
+			terminated: e.terminated,
+			confidence: null,
+			logits: null,
+			scores: null,
+			hiddenStates: null,
+			attention: null,
+			loss: null,
+			multinomialRand: null
+		}));
+		n.push({
+			token: e.tokenId,
+			score: e.probability,
+			text: e.tokenText,
+			terminated: e.terminated,
+			confidence: null,
+			logits: null,
+			scores: null,
+			hiddenStates: null,
+			attention: null,
+			loss: null,
+			multinomialRand: null
+		});
+		let r = {
+			tokens: t ? e.parent.tokens.slice() : [...e.parent.tokens, e.tokenId],
+			score: e.score,
+			text: t ? e.parent.text : e.parent.text + e.tokenText,
+			candidates: this.totalCandidates,
+			_output: n
+		};
+		if (!this.options?.keepDuplicates) {
+			if (this.seen.has(r.text.trim())) return this.totalCandidates--, null;
+			this.seen.add(r.text.trim());
+		}
+		return r;
+	}
+	async next() {
+		this.active = !0;
+		let e = await this._next();
+		return this.active = !1, e;
+	}
+	async _next() {
+		if (!this.options || !this.conversation) throw Error("Beamer not properly initialized with options and conversation.");
+		if (!this.hasRoot) {
+			let e = await this.tokeniseConversation(this.conversation), t = this.createInitialContext(e), n = {
+				tokens: [],
+				text: "",
+				score: 1,
+				context: t.context,
+				contextLength: t.contextLength,
+				terminated: !1,
+				nextChild: 0,
+				sortedChildren: void 0,
+				probabilities: []
+			};
+			this.totalCandidates = 1, await this.expand([n], 8), this.pushNextChild(n), this.hasRoot = !0;
+		}
+		for (; this.heap.size && this.active;) {
+			if (this.heap.peek().terminated) {
+				let e = this.heap.pop();
+				this.pushNextChild(e.parent);
+				let t = this.entryToBeam(e);
+				if (!t) continue;
+				return t;
 			}
-			let c = o.filter((e) => e.terminated), f = c.slice();
-			a.sort((e, t) => t.score - e.score);
-			for (let e of a) {
-				if (f.length >= r) break;
-				let t = e.terminatedByToken || e.terminatedByWhitespace, n = !e.dropLeadingWhitespaceTerminator, i = n ? e.parent.tokens.concat(e.token) : e.parent.tokens.slice(), a = e.terminatedByToken || e.dropLeadingWhitespaceTerminator ? e.parent.text : e.parent.text + e.tokenText, o = n ? this.appendTokenToContext(e.parent.context, e.parent.contextLength, e.token) : {
-					context: e.parent.context.clone(),
-					contextLength: e.parent.contextLength
-				}, s = e.parent._output?.slice() ?? [];
-				n && s.push({
-					token: e.token,
-					score: e.score,
-					text: e.tokenText,
-					terminated: t,
-					confidence: null,
-					logits: null,
-					scores: null,
-					hiddenStates: null,
-					attention: null,
-					loss: null,
-					multinomialRand: null
-				}), f.push({
-					tokens: i,
-					score: e.score,
-					text: a,
-					terminated: t,
-					context: o.context,
-					contextLength: o.contextLength,
-					_output: s,
-					candidates: u
+			let e = [];
+			for (; e.length < 8 && this.heap.size && !this.heap.peek().terminated;) {
+				let t = this.heap.pop();
+				this.pushNextChild(t.parent), e.push({
+					text: t.parent.text + t.tokenText,
+					tokens: [...t.parent.tokens, t.tokenId],
+					score: t.score,
+					...this.appendTokenToContext(t.parent.context, t.parent.contextLength, t.tokenId),
+					terminated: t.terminated,
+					nextChild: 0,
+					sortedChildren: void 0,
+					probabilities: [...t.parent.probabilities, t.probability]
 				});
 			}
-			if (u += a.length - f.length + c.length, o.forEach((e) => e.context.dispose()), f.length === 0) break;
-			o = f, n && Date.now() - l >= 40 && (n(o.slice(0, t.beams)), l = Date.now());
+			await this.expand(e, 8);
+			for (let t of e) this.pushNextChild(t);
 		}
-		this.active = !1, u += o.length;
-		let f = o.slice().sort((e, t) => t.score - e.score), p = [], m = /* @__PURE__ */ new Set();
-		for (let e of f) {
-			let n = e.text.trim();
-			if (m.has(n) || (m.add(n), p.push(e)), p.length >= t.beams) break;
+		return null;
+	}
+	async beam(e) {
+		this.active = !0;
+		let t = [];
+		for (; t.length < (this.options?.beams ?? 8);) {
+			let n = await this._next();
+			if (!n) break;
+			t.push(n), e?.(t);
 		}
-		return o.forEach((e) => {
-			e.context.dispose(), e.candidates = u;
-		}), p;
+		return this.active = !1, t.forEach((e) => {
+			e.candidates = this.totalCandidates;
+		}), t;
+	}
+	dispose() {
+		this.active = !1, this.heap.values.forEach((e) => {
+			e.parent.context.dispose();
+		}), this.heap.clear();
 	}
 };
 //#endregion
-export { f as default };
+export { p as default };

package/dist/inference/Generator.js

@@ -4,13 +4,13 @@ import { A as n, Ct as r, I as i, Y as a, _r as o, an as s, di as c, oi as l, yt
 import d from "../utilities/multinomialCPU.js";
 import f from "../utilities/topP.js";
 import { sparseSoftmaxCrossEntropy as p } from "../training/sparseCrossEntropy.js";
-import m from "./tokenisePrompt.js";
-import { CHARS as h, getTokenConfidence as g, padArray as _ } from "./utilities.js";
+import m, { tensorFromTokens as h } from "./tokenisePrompt.js";
+import { CHARS as g, getTokenConfidence as _, padArray as v } from "./utilities.js";
 //#region lib/inference/Generator.ts
-function v(e) {
+function y(e) {
 	return Array.isArray(e);
 }
-var y = class extends e {
+var b = class extends e {
 	model;
 	tokeniser;
 	active = !1;
@@ -34,7 +34,7 @@ var y = class extends e {
 		return t === this.actualTokeniser.eosToken || t === n;
 	}
 	async _generateToken(e, t, o) {
-		let m = o?.temperature ?? 1, h = o?.topK, _ = o?.topP, v = o?.usePadding ?? !1, y = {
+		let m = o?.temperature ?? 1, h = o?.topK, g = o?.topP, v = o?.usePadding ?? !1, y = {
 			training: !1,
 			attentionScores: o?.outputAttention ? { attentionOut: [] } : void 0,
 			cache: t,
@@ -68,10 +68,10 @@ var y = class extends e {
 				])), e.dispose());
 			}), c.dispose(), [f.div(m).squeeze([1]), h];
 		}), S, C, w, T = Math.random();
-		if (_) {
+		if (g) {
 			let e = a(b), t = await e.array();
 			e.dispose();
-			let n = f(t, _);
+			let n = f(t, g);
 			(o?.outputScores || o?.outputConfidence) && (C = t), S = d(n, T);
 		} else if (h) {
 			let { values: e, indices: t } = n(b, h), i = r(e, 1);
@@ -129,7 +129,7 @@ var y = class extends e {
 			outputTensor: S,
 			token: D,
 			text: O,
-			confidence: o?.outputConfidence && C ? g(C[0]) : null,
+			confidence: o?.outputConfidence && C ? _(C[0]) : null,
 			score: o?.outputScore && C ? C[0][D] : null,
 			logits: o?.outputLogits ? (await b.array())[0] : null,
 			scores: o?.outputScores && C ? C[0] : null,
@@ -152,28 +152,42 @@ var y = class extends e {
 			content: "",
 			_timestamp: Date.now()
 		}), n = !0, this.resetCache(!e?.noCache)) : (this.lastToken < 0 || t) && this.resetCache(!e?.noCache);
-		let r = this.lastToken >= 0 && this.cache ? i([this.lastToken], [1, 1], "int32") : await m(this.actualTokeniser, this.model.config.blockSize, t ? n ? this.outputConversation.slice(0, -1) : this.outputConversation : void 0, n ? e : {
+		let r = this.outputConversation[this.outputConversation.length - 1];
+		r.content.length > 0 && (!r._output || r._output.length === 0) && (r._output = this.actualTokeniser.encode(r.content).map((e) => ({
+			token: e,
+			text: this.actualTokeniser.decode([e]),
+			confidence: 1,
+			score: 1,
+			logits: null,
+			hiddenStates: null,
+			scores: null,
+			attention: null,
+			loss: null,
+			multinomialRand: null,
+			terminated: !1
+		})));
+		let a = this.lastToken >= 0 && this.cache ? i([this.lastToken], [1, 1], "int32") : h(await m(this.actualTokeniser, this.model.config.blockSize, t ? n ? this.outputConversation.slice(0, -1) : this.outputConversation : void 0, n ? e : {
 			...e,
 			continuation: !0
-		}), a = e?.maxLength ?? 1e3;
-		for (let t = 0; t < a && this.active; t++) {
-			let n = await this._generateToken(r, this.cache ? this.cache : void 0, {
+		})), s = e?.maxLength ?? 1e3;
+		for (let t = 0; t < s && this.active; t++) {
+			let n = await this._generateToken(a, this.cache ? this.cache : void 0, {
 				...e,
 				usePadding: !this.cache
 			});
-			if (this.cache) r.dispose(), r = n.outputTensor;
+			if (this.cache) a.dispose(), a = n.outputTensor;
 			else {
-				let e = r;
-				r = o([r, n.outputTensor], 1), e.dispose();
+				let e = a;
+				a = o([a, n.outputTensor], 1), e.dispose();
 			}
-			let i = this.outputConversation[this.outputConversation.length - 1];
+			let r = this.outputConversation[this.outputConversation.length - 1];
 			if (this.cache || n.outputTensor.dispose(), n.terminated) {
-				i._completed = !0;
+				r._completed = !0;
 				break;
 			}
-			t === a - 1 && a > 1 && (i._completed = !0, n.terminated = !0), i.content += n.text, i._output ||= [], i._output.push(n);
+			t === s - 1 && s > 1 && (r._completed = !0, n.terminated = !0), r.content += n.text, r._output ||= [], r._output.push(n);
 		}
-		return r.dispose(), this.outputConversation;
+		return a.dispose(), this.outputConversation;
 	}
 	resetCache(e) {
 		this.cache && (this.cache.forEach((e) => {
@@ -199,7 +213,7 @@ var y = class extends e {
 			};
 			this.cache = e, this.lastToken = -1;
 		}
-		let r = this.tokeniser.trained ? this.tokeniser : new t(_(h, this.tokeniser.vocabSize));
+		let r = this.tokeniser.trained ? this.tokeniser : new t(v(g, this.tokeniser.vocabSize));
 		this.actualTokeniser = r, n?.loraName ? this.model.attachLoRA(n.loraName) : this.model.hasLoRA() && this.model.detachLoRA();
 	}
 	async step(e, t) {
@@ -207,7 +221,7 @@ var y = class extends e {
 			...t,
 			maxLength: 1
 		};
-		return v(e) ? this.generate(e, n) : this.generate({
+		return y(e) ? this.generate(e, n) : this.generate({
 			...e,
 			...n
 		});
@@ -256,4 +270,4 @@ var y = class extends e {
 	}
 };
 //#endregion
-export { y as default, v as isConversation };
+export { b as default, y as isConversation };

package/dist/inference/tokenisePrompt.d.ts

@@ -1,4 +1,5 @@
 import { Conversation, ITokeniser } from '../tokeniser/type';
 import { IGenerateOptions } from './types';
 import { Tensor } from '@tensorflow/tfjs-core';
-export default function tokenisePrompt(tokeniser: ITokeniser, blockSize: number, prompt?: Conversation[], options?: IGenerateOptions): Promise<Tensor>;
+export declare function tensorFromTokens(tokens: number[]): Tensor;
+export default function tokenisePrompt(tokeniser: ITokeniser, blockSize: number, prompt?: Conversation[], options?: IGenerateOptions): Promise<number[]>;

package/dist/inference/tokenisePrompt.js

@@ -1,13 +1,16 @@
 import { I as e } from "../dist-Da20xy8E.js";
 //#region lib/inference/tokenisePrompt.ts
-async function t(t, n, r, i) {
-	if (r) {
-		let a = r.length > 0 && (r[r.length - 1].role === "text" || r[r.length - 1].role === "assistant"), o;
-		return i?.nonConversational ? o = a && i?.continuation ? [t.bosToken, ...t.encode(r[r.length - 1].content)] : t.encodeAsSequence(r, !0) : a && i?.continuation ? (o = t.encodeConversation(r, !0), o[o.length - 1] === t.getSpecialTokenIndex("<|assistant_start|>") && o[o.length - 2] === t.getSpecialTokenIndex("<|assistant_end|>") && (o.pop(), o.pop())) : o = t.encodeConversation(r, !0), o.length > n && (o = o.slice(-n)), i?.debug && console.log("Tokenised prompt:", t.decode(o)), e([o], [1, o.length], "int32");
+function t(t) {
+	return e([t], [1, t.length], "int32");
+}
+async function n(e, t, n, r) {
+	if (n) {
+		let i = n.length > 0 && (n[n.length - 1].role === "text" || n[n.length - 1].role === "assistant"), a;
+		return r?.nonConversational ? a = i && r?.continuation ? [e.bosToken, ...e.encode(n[n.length - 1].content)] : e.encodeAsSequence(n, !0) : i && r?.continuation ? (a = e.encodeConversation(n, !0), a[a.length - 1] === e.getSpecialTokenIndex("<|assistant_start|>") && a[a.length - 2] === e.getSpecialTokenIndex("<|assistant_end|>") && (a.pop(), a.pop())) : a = e.encodeConversation(n, !0), a.length > t && (a = a.slice(-t)), r?.debug && console.log("Tokenised prompt:", e.decode(a)), a;
 	} else {
-		let n = i?.nonConversational ? void 0 : t.getSpecialTokenIndex("<|assistant_start|>"), r = n === void 0 ? [t.bosToken] : [t.bosToken, n];
-		return e([r], [1, r.length], "int32");
+		let t = r?.nonConversational ? void 0 : e.getSpecialTokenIndex("<|assistant_start|>");
+		return t === void 0 ? [e.bosToken] : [e.bosToken, t];
 	}
 }
 //#endregion
-export { t as default };
+export { n as default, t as tensorFromTokens };

package/dist/inference/types.d.ts

@@ -58,6 +58,7 @@ export interface BeamerOptions extends IGenerateOptions {
     maxBeamLength: number;
     beams: number;
     endOnWhiteSpace?: boolean;
+    keepDuplicates?: boolean;
 }
 export interface IBeam {
     tokens: number[];

package/dist/layers/GQA.d.ts

@@ -0,0 +1,44 @@
+import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
+import { Tensor } from '@tensorflow/tfjs-core';
+import { GPTConfig } from '../../models/config';
+import { RMSNormConfig } from './RMSNorm';
+export interface KVCache {
+    k?: Tensor;
+    v?: Tensor;
+    length: number;
+    cumulativeLength: number;
+}
+export interface AttentionScores {
+    meanOfHeads?: boolean;
+    attentionOut?: Tensor[];
+}
+interface AttentionForwardAttributes extends ForwardAttributes {
+    attentionScores?: AttentionScores;
+    pastKV?: KVCache;
+    seed?: number;
+    ropePositionOffset?: number;
+}
+export interface GQAConfig extends RMSNormConfig {
+    useQKNorm?: boolean;
+    ropeInterleaved?: boolean;
+    kvHeads?: number;
+}
+export default class GQA extends BaseLayer<AttentionForwardAttributes> {
+    private readonly attentionConfig;
+    private divisor;
+    private index;
+    private units;
+    private projUnits;
+    private ATTN;
+    private PROJ;
+    private qNorm;
+    private kNorm;
+    constructor(index: number, config: GPTConfig, attentionConfig: GQAConfig, parent?: BaseLayer);
+    protected build(): void;
+    private getAttentionScores;
+    private getQKV;
+    private getOutputProjection;
+    private updateCache;
+    forward(attr: AttentionForwardAttributes, x: Tensor): Tensor;
+}
+export {};

package/dist/layers/{CausalSelfAttention.js → GQA.js}

@@ -2,18 +2,18 @@ import { T as e, di as t, oi as n, pt as r } from "../dist-Da20xy8E.js";
 import { isPackedTensor as i } from "../utilities/packed.js";
 import { transpose16 as a } from "../ops/transpose16.js";
 import { reshape16 as o } from "../ops/reshape16.js";
-import { t as s } from "../matMul16-DGtlziMn.js";
+import { t as s } from "../matMul16-CO1QhsRt.js";
 import { t as c } from "../pack16-BhuXNUS7.js";
 import { attentionMask as l } from "../ops/attentionMask.js";
 import u from "./BaseLayer.js";
-import { t as d } from "../rope-C8EAiSGY.js";
+import { t as d } from "../rope-BPu_uooz.js";
 import { appendCache as f } from "../ops/appendCache.js";
 import { softmax16 as p } from "../ops/softmax16.js";
 import { dot16 as m } from "../ops/dot16.js";
 import { qkv as h } from "../ops/qkv.js";
-import { normRMS as g } from "../ops/normRMS.js";
-import { dropout16 as _ } from "../ops/dropout16.js";
-//#region lib/layers/CausalSelfAttention.ts
+import { dropout16 as g } from "../ops/dropout16.js";
+import _ from "./RMSNorm.js";
+//#region lib/layers/GQA.ts
 var v = class extends u {
 	attentionConfig;
 	divisor;
@@ -22,8 +22,10 @@ var v = class extends u {
 	projUnits;
 	ATTN;
 	PROJ;
+	qNorm = null;
+	kNorm = null;
 	constructor(e, t, n, r) {
-		super(t, r), this.attentionConfig = n, this.index = e, this.units = t.nEmbed * 3, this.projUnits = t.nEmbed, this.ATTN = `block_${this.index}_cAttn`, this.PROJ = `block_${this.index}_cProj`, this.addVariable(this.ATTN), this.addVariable(this.PROJ), this.divisor = 1 / Math.sqrt(t.nEmbed / t.nHead);
+		super(t, r), this.attentionConfig = n, this.index = e, this.units = t.nEmbed * 3, this.projUnits = t.nEmbed, this.ATTN = `block_${this.index}_cAttn`, this.PROJ = `block_${this.index}_cProj`, this.addVariable(this.ATTN), this.addVariable(this.PROJ), this.divisor = 1 / Math.sqrt(t.nEmbed / t.nHead), (this.attentionConfig.useQKNorm ?? !1) && (this.qNorm = new _(t, n, `block_${this.index}_qNorm`, this), this.kNorm = new _(t, n, `block_${this.index}_kNorm`, this));
 	}
 	build() {
 		this.hasVariable(this.ATTN) === !1 && this.setVariable(this.ATTN, e(r([this.config.nEmbed, this.units], 0, .02), !0, this.ATTN)), this.hasVariable(this.PROJ) === !1 && this.setVariable(this.PROJ, e(r([this.projUnits, this.config.nEmbed], 0, .02), !0, this.PROJ));
@@ -62,20 +64,20 @@ var v = class extends u {
 	forward(e, r) {
 		return t(() => {
 			this.startMemory();
-			let [t, i, a] = this.getQKV(r), o = e.pastKV ? e.pastKV.cumulativeLength : e.ropePositionOffset || 0, c = e.ropeCache, l = c ? d(t, c, o) : t, u = c ? d(i, c, o) : i, f = this.attentionConfig.useQKNorm ?? !1, p = f ? g(l) : l;
-			f && l.dispose();
-			let m = f ? g(u) : u;
-			f && u.dispose(), c && (t.dispose(), i.dispose());
+			let [t, i, a] = this.getQKV(r), o = e.pastKV ? e.pastKV.cumulativeLength : e.ropePositionOffset || 0, c = e.ropeCache, l = this.attentionConfig.ropeInterleaved ?? !0, u = c ? d(t, c, o, !1, l) : t, f = c ? d(i, c, o, !1, l) : i, p = this.qNorm ? this.qNorm.call(e, u) : u;
+			this.qNorm && u.dispose();
+			let m = this.kNorm ? this.kNorm.call(e, f) : f;
+			this.kNorm && f.dispose(), c && (t.dispose(), i.dispose());
 			let h = e.pastKV ? e.pastKV.length : 0;
 			e.pastKV && !e.training && this.updateCache(m, a, e.pastKV);
-			let v = e.pastKV?.k ? e.pastKV.k : m, y = e.pastKV?.v ? e.pastKV.v : a, b;
-			b = h > 0 ? this.getAttentionScores(p, v, h) : this.getAttentionScores(p, v), p.dispose(), e.pastKV || v.dispose();
-			let x = s(b, y), S = e.attentionScores !== void 0 && e.attentionScores.attentionOut !== void 0;
-			S || b.dispose(), e.pastKV || y.dispose();
-			let C = this.getOutputProjection(x);
-			if (x.dispose(), S && e.attentionScores && e.attentionScores.attentionOut !== void 0) {
-				let t = b.shape[1], r = b.shape[2];
-				e.attentionScores.attentionOut?.push(n(b.slice([
+			let _ = e.pastKV?.k ? e.pastKV.k : m, v = e.pastKV?.v ? e.pastKV.v : a, y;
+			y = h > 0 ? this.getAttentionScores(p, _, h) : this.getAttentionScores(p, _), p.dispose(), e.pastKV || _.dispose();
+			let b = s(y, v), x = e.attentionScores !== void 0 && e.attentionScores.attentionOut !== void 0;
+			x || y.dispose(), e.pastKV || v.dispose();
+			let S = this.getOutputProjection(b);
+			if (b.dispose(), x && e.attentionScores && e.attentionScores.attentionOut !== void 0) {
+				let t = y.shape[1], r = y.shape[2];
+				e.attentionScores.attentionOut?.push(n(y.slice([
 					0,
 					0,
 					0,
@@ -91,7 +93,7 @@ var v = class extends u {
 					-1
 				])));
 			}
-			return this.endMemory("CausalSelfAttention"), e.dropout && e.dropout > 0 ? _(C, e.dropout) : C;
+			return this.endMemory("CausalSelfAttention"), e.dropout && e.dropout > 0 ? g(S, e.dropout) : S;
 		});
 	}
 };

package/dist/layers/HeadLayer.d.ts

@@ -0,0 +1,12 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
+import { GPTConfig } from '../../models/config';
+export default class HeadLayer extends BaseLayer {
+    private vocabSize;
+    private embedDim;
+    private initializer;
+    private WEIGHTS;
+    constructor(config: GPTConfig, name: string, parent?: BaseLayer);
+    project(inputs: Tensor): Tensor;
+    forward(_: ForwardAttributes, x: Tensor): Tensor;
+}