@genai-fi/nanogpt 1.2.0 → 1.2.1

package/dist/api/beamer.d.ts

@@ -13,9 +13,13 @@ export default class BeamAPI {
     private _model;
     private _tokeniser;
     private _busyCount;
+    private _jobs;
+    private _queue;
     constructor(model: Model<ModelForwardAttributes, GPTConfig>, tokeniser: ITokeniser);
     on<E extends keyof BeamEvents>(event: E, listener: BeamEvents[E]): void;
     off<E extends keyof BeamEvents>(event: E, listener: BeamEvents[E]): void;
+    cancel(id?: string): void;
+    private startJob;
     create(conversation: Conversation[], options: BeamerOptions): string;
 }
 export {};

package/dist/api/beamer.js

@@ -7,6 +7,8 @@ var r = class {
 	_model;
 	_tokeniser;
 	_busyCount = 0;
+	_jobs = /* @__PURE__ */ new Map();
+	_queue = [];
 	constructor(t, n) {
 		this._model = t, this._tokeniser = n, this.ee = new e();
 	}
@@ -16,15 +18,34 @@ var r = class {
 	off(e, t) {
 		this.ee.off(e, t);
 	}
+	cancel(e) {
+		if (e) {
+			let t = this._jobs.get(e);
+			t && (t.canceled = !0, t.beamer.cancel());
+		} 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.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);
+			let n = this._queue.shift();
+			n && this.startJob(n);
+		}).catch((t) => {
+			this._busyCount--, this._busyCount === 0 && this.ee.emit("status", "ready"), this.ee.emit("error", t), this._jobs.delete(e.id);
+			let n = this._queue.shift();
+			n && this.startJob(n);
+		});
+	}
 	create(e, r) {
-		let i = t();
-		return this._busyCount++, this.ee.emit("status", "busy"), new n(this._model, this._tokeniser).beam(e, r, (e) => {
-			this.ee.emit("progress", i, e);
-		}).then((e) => {
-			this._busyCount--, this._busyCount === 0 && this.ee.emit("status", "ready"), this.ee.emit("done", i, e);
-		}).catch((e) => {
-			this._busyCount--, this._busyCount === 0 && this.ee.emit("status", "ready"), this.ee.emit("error", e);
-		}), i;
+		let i = t(), a = {
+			id: i,
+			beamer: new n(this._model, this._tokeniser),
+			canceled: !1,
+			conversation: e,
+			options: r
+		};
+		return this._jobs.set(i, a), this._busyCount > 0 ? (this._queue.push(a), i) : (this.startJob(a), i);
 	}
 };
 //#endregion

package/dist/inference/Beamer.d.ts

@@ -1,11 +1,11 @@
 import { Conversation, ITokeniser } from '../tokeniser/type';
 import { default as Model, ModelForwardAttributes } from '../models/model';
-import { default as EE } from 'eventemitter3';
 import { BeamerOptions, IBeam } from './types';
-export default class Beamer extends EE {
+export default class Beamer {
     private readonly model;
     private readonly tokeniser;
     private actualTokeniser;
+    private active;
     constructor(model: Model<ModelForwardAttributes>, tokeniser: ITokeniser);
     private shouldTerminate;
     private initialise;
@@ -13,5 +13,6 @@ export default class Beamer extends EE {
     private createInitialContext;
     private appendTokenToContext;
     private batchedNextProbabilities;
+    cancel(): void;
     beam(conversation: Conversation[], options: BeamerOptions, onStep?: (beams: IBeam[]) => void): Promise<IBeam[]>;
 }

package/dist/inference/Beamer.js

@@ -1,30 +1,30 @@
-import { t as e } from "../eventemitter3-D_qV3Lof.js";
-import t from "../tokeniser/CharTokeniser.js";
-import { L as n, V as r, Y as i, _n as a, _r as o, di as s, xt as c } from "../dist-Da20xy8E.js";
-import l from "../utilities/topP.js";
-import { CHARS as u, padArray as d } from "./utilities.js";
+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";
 //#region lib/inference/Beamer.ts
-function f(e, t) {
+function d(e, t) {
 	return 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));
 }
-var p = class extends e {
+var f = class {
 	model;
 	tokeniser;
 	actualTokeniser;
+	active = !1;
 	constructor(e, t) {
-		super(), this.model = e, this.tokeniser = t, this.actualTokeniser = t;
+		this.model = e, this.tokeniser = t, this.actualTokeniser = t;
 	}
 	shouldTerminate(e, t) {
 		if (e) return !1;
 		let n = this.tokeniser.getSpecialTokenIndex("<|assistant_end|>");
 		return t === this.actualTokeniser.eosToken || t === n;
 	}
-	initialise(e) {
-		let n = this.tokeniser.trained ? this.tokeniser : new t(d(u, this.tokeniser.vocabSize));
-		this.actualTokeniser = n, e?.loraName ? this.model.attachLoRA(e.loraName) : this.model.hasLoRA() && this.model.detachLoRA();
+	initialise(t) {
+		let n = this.tokeniser.trained ? this.tokeniser : new e(u(l, this.tokeniser.vocabSize));
+		this.actualTokeniser = n, t?.loraName ? this.model.attachLoRA(t.loraName) : this.model.hasLoRA() && this.model.detachLoRA();
 	}
 	async tokeniseConversation(e) {
 		let t = this.actualTokeniser.encodeConversation(e, !1);
@@ -32,65 +32,68 @@ var p = class extends e {
 		return t;
 	}
 	createInitialContext(e) {
-		let t = this.model.config.blockSize, r = e.length > t ? e.slice(-t) : e.slice();
+		let n = this.model.config.blockSize, r = e.length > n ? e.slice(-n) : e.slice();
 		return {
-			context: n(r.concat(Array(t - r.length).fill(0)), "int32"),
+			context: t(r.concat(Array(n - r.length).fill(0)), "int32"),
 			contextLength: r.length
 		};
 	}
-	appendTokenToContext(e, t, r) {
-		let i = this.model.config.blockSize;
-		return t < i ? {
-			context: s(() => {
-				let n = c([t], i).squeeze([0]).asType("int32"), o = a(1, "int32").sub(n), s = n.mul(a(r, "int32"));
-				return e.mul(o).add(s);
+	appendTokenToContext(e, n, r) {
+		let c = this.model.config.blockSize;
+		return n < c ? {
+			context: o(() => {
+				let t = s([n], c).squeeze([0]).asType("int32"), a = i(1, "int32").sub(t), o = t.mul(i(r, "int32"));
+				return e.mul(a).add(o);
 			}),
-			contextLength: t + 1
+			contextLength: n + 1
 		} : {
-			context: s(() => o([e.slice([1], [i - 1]), n([r], "int32")], 0)),
-			contextLength: i
+			context: o(() => a([e.slice([1], [c - 1]), t([r], "int32")], 0)),
+			contextLength: c
 		};
 	}
-	async batchedNextProbabilities(e, t, n) {
+	async batchedNextProbabilities(e, t, i) {
 		if (e.length === 0) return [];
-		let a = e[0], o = e.slice();
-		for (; o.length < t;) o.push(a);
-		let u = r(o.map((e) => e.context)), d = this.model.forward({
+		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({
 			training: !1,
 			mixedPrecision: !0
-		}, u), f = o.map((e) => Math.max(0, e.contextLength - 1)), p = s(() => {
-			let e = c(f, this.model.config.blockSize).expandDims(2);
-			return i(d.mul(e).sum(1));
+		}, 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 _ = n.topP ?? 1;
-		return g.map((e) => l(e, _));
+		let _ = i.topP ?? 1;
+		return g.map((e) => c(e, _));
+	}
+	cancel() {
+		this.active = !1;
 	}
 	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.initialise(t), this.active = !0;
 		let i = await this.tokeniseConversation(e), a = this.createInitialContext(i), o = [{
 			tokens: [],
-			score: 0,
+			score: 1,
 			text: "",
 			terminated: !1,
 			context: a.context,
 			contextLength: a.contextLength
 		}], s = t.maxBeamLength, c = t.endOnWhiteSpace === !0 ? Math.max(s, t.maxLength ?? s + this.model.config.blockSize) : s, l = Date.now();
-		for (let e = 0; e < c; e++) {
+		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 = f(c, t.topK ? Math.min(t.topK, r) : Math.max(1, r));
+				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);
+					let n = this.actualTokeniser.decode([e.token]), r = this.shouldTerminate(t.allowSpecial ?? !1, e.token), i = o.tokens.length + 1 >= Math.max(2, 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 + Math.log(Math.max(e.prob, 1e-12)),
+						score: o.score * e.prob,
 						terminatedByToken: r,
 						terminatedByWhitespace: c,
 						dropLeadingWhitespaceTerminator: l
@@ -117,16 +120,14 @@ var p = class extends e {
 			if (o.forEach((e) => e.context.dispose()), c.length === 0) break;
 			o = c, n && Date.now() - l >= 40 && (n(o.slice(0, t.beams)), l = Date.now());
 		}
-		o.forEach((e) => {
-			e.score /= Math.max(1, e.text.length);
-		});
-		let u = o.slice().sort((e, t) => t.score - e.score), d = [], p = /* @__PURE__ */ new Set();
+		this.active = !1;
+		let u = o.slice().sort((e, t) => t.score - e.score), f = [], p = /* @__PURE__ */ new Set();
 		for (let e of u) {
 			let n = e.text.trim();
-			if (p.has(n) || (p.add(n), d.push(e)), d.length >= t.beams) break;
+			if (p.has(n) || (p.add(n), f.push(e)), f.length >= t.beams) break;
 		}
-		return o.forEach((e) => e.context.dispose()), d;
+		return o.forEach((e) => e.context.dispose()), f;
 	}
 };
 //#endregion
-export { p as default };
+export { f as default };

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@genai-fi/nanogpt",
-    "version": "1.2.0",
+    "version": "1.2.1",
     "type": "module",
     "main": "dist/main.js",
     "types": "dist/main.d.ts",