@genai-fi/nanogpt 0.23.0 → 1.0.0

package/dist/{Generator.d.ts → inference/Generator.d.ts}

@@ -1,31 +1,16 @@
-import { Conversation, ITokeniser } from './tokeniser/type';
+import { Conversation, ITokeniser } from '../tokeniser/type';
 import { default as EE } from 'eventemitter3';
-import { default as Model, ModelForwardAttributes } from './models/model';
-import { GenerateOptions } from './inference/types';
+import { default as Model, ModelForwardAttributes } from '../models/model';
+import { IGenerateOptions, GeneratorConversation, IGeneratorOutput } from './types';
 export declare function isConversation(data: unknown): data is Conversation[];
-export interface IGenerateOptions extends GenerateOptions {
-    maxLength?: number;
-    noCache?: boolean;
-    allowSpecial?: boolean;
-    nonConversational?: boolean;
-    continuation?: boolean;
-}
 export interface IGenerator extends EE<'start' | 'stop' | 'tokens' | 'reset'> {
-    generate(prompt: Conversation[], options?: IGenerateOptions): Promise<Conversation[]>;
-    generate(options?: IGenerateOptions): Promise<Conversation[]>;
-    step(prompt: Conversation[], options?: IGenerateOptions): Promise<Conversation[]>;
-    step(options?: IGenerateOptions): Promise<Conversation[]>;
+    generate(prompt: Conversation[], options?: IGenerateOptions): Promise<GeneratorConversation[]>;
+    generate(options?: IGenerateOptions): Promise<GeneratorConversation[]>;
+    step(prompt: Conversation[], options?: IGenerateOptions): Promise<GeneratorConversation[]>;
+    step(options?: IGenerateOptions): Promise<GeneratorConversation[]>;
     stop(): void;
-    getConversation(): Conversation[];
-    getAttentionData(): number[][][][][];
-    getProbabilitiesData(): number[][][];
-    getEmbeddingsData(): {
-        name: string;
-        tensor: number[][];
-    }[][];
-    getTokens(): number[];
-    getLastLoss(): number | null;
-    getLastMultinomialRand(): number | null;
+    getConversation(): GeneratorConversation[];
+    getRawOutput(): IGeneratorOutput[];
     dispose(): void;
     reset(): void;
 }
@@ -42,18 +27,13 @@ export default class Generator extends EE<'start' | 'stop' | 'tokens' | 'reset'>
     private outputConversation;
     private actualTokeniser;
     private lastToken;
-    private attentionData;
-    private probabilitiesData;
-    private embeddingsData;
-    private tokens;
     private lastLoss;
-    private lastMultinomialRand;
+    private rawOutput;
     private jobQueue;
     private processingJob;
     private startTime;
+    private tokenCount;
     constructor(model: Model<ModelForwardAttributes>, tokeniser: ITokeniser);
-    private tokenisePrompt;
-    private processResponse;
     /** Generate logits and select a token. */
     private _generateToken;
     /** Generate multiple tokens in a loop and produce text */
@@ -62,21 +42,13 @@ export default class Generator extends EE<'start' | 'stop' | 'tokens' | 'reset'>
     reset(): void;
     dispose(): void;
     private initialise;
-    step(prompt: Conversation[], options?: IGenerateOptions): Promise<Conversation[]>;
-    step(options?: IGenerateOptions): Promise<Conversation[]>;
-    generate(prompt: Conversation[], options?: IGenerateOptions): Promise<Conversation[]>;
-    generate(options?: IGenerateOptions): Promise<Conversation[]>;
+    step(prompt: Conversation[], options?: IGenerateOptions): Promise<GeneratorConversation[]>;
+    step(options?: IGenerateOptions): Promise<GeneratorConversation[]>;
+    generate(prompt: Conversation[], options?: IGenerateOptions): Promise<GeneratorConversation[]>;
+    generate(options?: IGenerateOptions): Promise<GeneratorConversation[]>;
     private startJob;
     getQueueLength(): number;
     stop(): void;
     getConversation(): Conversation[];
-    getAttentionData(): number[][][][][];
-    getProbabilitiesData(): number[][][];
-    getEmbeddingsData(): {
-        name: string;
-        tensor: number[][];
-    }[][];
-    getTokens(): number[];
-    getLastLoss(): number | null;
-    getLastMultinomialRand(): number | null;
+    getRawOutput(): IGeneratorOutput[];
 }

package/dist/inference/Generator.js

@@ -0,0 +1,271 @@
+import { t as e } from "../eventemitter3-D_qV3Lof.js";
+import { SPECIALS as t } from "../tokeniser/BaseTokeniser.js";
+import n from "../tokeniser/CharTokeniser.js";
+import { A as r, Ct as i, I as a, Y as o, _r as s, an as c, di as l, oi as u, yt as d } from "../dist-Da20xy8E.js";
+import f from "../utilities/multinomialCPU.js";
+import p from "../utilities/topP.js";
+import { sparseSoftmaxCrossEntropy as m } from "../training/sparseCrossEntropy.js";
+import h from "./tokenisePrompt.js";
+import { getTokenConfidence as g } from "./utilities.js";
+//#region lib/inference/Generator.ts
+function _(e) {
+	return Array.isArray(e);
+}
+var v = [
+	...t,
+	...Array.from({ length: 95 }, (e, t) => String.fromCharCode(t + 32)),
+	..."áéíóúüñ¿¡",
+	..."äöÄÖÅå",
+	..."αβγδεζηθικλμνξοπρστυφχψωΑΒΓΔΕΖΗΘΙΚΛΜΝΞΟΠΡΣΤΥΦΧΨΩ",
+	..."абвгдеёжзийклмнопрстуфхцчшщъыьэюяАБВГДЕЁЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЫЬЭЮЯ"
+];
+function y(e, t) {
+	return e.length === t ? e : e.length > t ? e.slice(0, t) : e.concat(Array(t - e.length).fill(""));
+}
+var b = class extends e {
+	model;
+	tokeniser;
+	active = !1;
+	cache = null;
+	initialPrompt = null;
+	outputConversation = [];
+	actualTokeniser;
+	lastToken = -1;
+	lastLoss = null;
+	rawOutput = [];
+	jobQueue = [];
+	processingJob = !1;
+	startTime = null;
+	tokenCount = 0;
+	constructor(e, t) {
+		super(), this.model = e, this.tokeniser = t, this.actualTokeniser = t;
+	}
+	async _generateToken(e, t, n) {
+		let s = n?.temperature ?? 1, h = n?.topK, _ = n?.topP, v = n?.usePadding ?? !1, y = {
+			training: !1,
+			attentionScores: n?.outputAttention ? { attentionOut: [] } : void 0,
+			cache: t,
+			outputEmbeddings: !!n?.outputHiddenStates
+		}, [b, x] = l(() => {
+			let t = e, r = t.shape[1], i = r <= this.model.config.blockSize ? t : t.slice([0, r - this.model.config.blockSize], [t.shape[0], this.model.config.blockSize]), o = v ? this.model.config.blockSize - i.shape[1] : 0, c = o > 0 ? d(i, [[0, 0], [0, o]]) : i, l = this.model.forward(y, c), f = l.shape[1] - 1 - o, p = l.slice([
+				0,
+				f,
+				0
+			], [
+				l.shape[0],
+				1,
+				l.shape[2]
+			]), h;
+			if (n?.targets) {
+				let e = n.targets.shift();
+				if (e !== void 0) {
+					let t = a([[e]], [1, 1], "int32"), n = m(p, t);
+					h = n.mean(), t.dispose(), n.dispose();
+				}
+			}
+			return y.attentionScores?.attentionOut && y.attentionScores.attentionOut.forEach((e, t) => {
+				e.shape[1] !== 1 && (y.attentionScores.attentionOut[t] = u(e.slice([
+					0,
+					f,
+					0
+				], [
+					e.shape[0],
+					1,
+					e.shape[2]
+				])), e.dispose());
+			}), l.dispose(), [p.div(s).squeeze([1]), h];
+		}), S, C, w, T = Math.random();
+		if (_) {
+			let e = o(b), t = await e.array();
+			e.dispose();
+			let r = p(t, _);
+			(n?.outputScores || n?.outputConfidence) && (C = t), S = f(r, T);
+		} else if (h) {
+			let { values: e, indices: t } = r(b, h), n = i(e, 1);
+			S = c(t, n, 1), e.dispose(), t.dispose(), n.dispose();
+		} else if (S = i(b, 1), n?.outputScores || n?.outputConfidence) {
+			let e = o(b);
+			C = await e.array(), e.dispose();
+		}
+		if (y.embeddings) {
+			let e = (n?.outputHiddenStates === "all" ? y.embeddings : y.embeddings.filter((e) => e.name.startsWith("block_output_"))).map(async (e) => {
+				let t = e.tensor.shape[1], r = e.tensor.slice([
+					0,
+					t - 1,
+					0
+				], [
+					e.tensor.shape[0],
+					1,
+					e.tensor.shape[2]
+				]);
+				e.tensor.dispose();
+				let i = r.squeeze([1]);
+				if (r.dispose(), n?.outputHiddenStates === "softmax") {
+					let t = this.model.project(i);
+					i.dispose();
+					let n = o(t, -1);
+					t.dispose();
+					let r = {
+						name: e.name,
+						tensor: await n.array()
+					};
+					return n.dispose(), r;
+				} else if (n?.outputHiddenStates === "logits") {
+					let t = this.model.project(i);
+					i.dispose();
+					let n = {
+						name: e.name,
+						tensor: await t.array()
+					};
+					return t.dispose(), n;
+				} else {
+					let t = await i.array();
+					return i.dispose(), {
+						name: e.name,
+						tensor: t
+					};
+				}
+			});
+			w = await Promise.all(e);
+		}
+		let E = S.reshape([1, 1]);
+		S.dispose(), S = E;
+		let D = (await S.array())[0][0], O = this.actualTokeniser.decode([D]);
+		this.lastToken = D;
+		let k = !n?.allowSpecial && this.tokeniser.isSpecialToken(D), A = {
+			outputTensor: S,
+			token: D,
+			text: O,
+			confidence: n?.outputConfidence && C ? g(C[0]) : null,
+			score: n?.outputScore && C ? C[0][D] : null,
+			logits: n?.outputLogits ? (await b.array())[0] : null,
+			scores: n?.outputScores && C ? C[0] : null,
+			hiddenStates: w ? w.map((e) => e.tensor[0]) : null,
+			attention: n?.outputAttention ? await Promise.all(y.attentionScores?.attentionOut?.map((e) => e.array()) ?? []) : null,
+			loss: this.lastLoss,
+			multinomialRand: T,
+			terminated: k
+		};
+		if (b.dispose(), x) {
+			let e = await x.array();
+			x.dispose(), A.loss = e;
+		}
+		return this.rawOutput.push(A), (!n?.chunkSize || this.tokenCount++ % n.chunkSize === 0) && (this.emit("tokens", A), n?._onChunk && await n._onChunk(A)), A;
+	}
+	async _generate(e, t) {
+		let n = !1;
+		this.outputConversation.length === 0 || this.outputConversation[this.outputConversation.length - 1]._completed || this.outputConversation[this.outputConversation.length - 1].role !== "assistant" && e?.nonConversational !== !0 || this.outputConversation[this.outputConversation.length - 1].role !== "text" && e?.nonConversational === !0 ? (this.outputConversation.push({
+			role: e?.nonConversational === !0 ? "text" : "assistant",
+			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 ? a([this.lastToken], [1, 1], "int32") : await h(this.actualTokeniser, this.model.config.blockSize, t ? n ? this.outputConversation.slice(0, -1) : this.outputConversation : void 0, e), i = e?.maxLength ?? 1e3;
+		for (let t = 0; t < i && this.active; t++) {
+			let n = await this._generateToken(r, this.cache ? this.cache : void 0, {
+				...e,
+				usePadding: !this.cache
+			});
+			if (this.cache) r.dispose(), r = n.outputTensor;
+			else {
+				let e = r;
+				r = s([r, n.outputTensor], 1), e.dispose();
+			}
+			let a = this.outputConversation[this.outputConversation.length - 1];
+			if (this.cache || n.outputTensor.dispose(), n.terminated) {
+				a._completed = !0;
+				break;
+			}
+			t === i - 1 && i > 1 && (a._completed = !0, n.terminated = !0), a.content += n.text, a._output ||= [], a._output.push(n);
+		}
+		return r.dispose(), this.outputConversation;
+	}
+	resetCache(e) {
+		this.cache && (this.cache.forEach((e) => {
+			e && (e.k && e.k.dispose(), e.v && e.v.dispose(), e.k = void 0, e.v = void 0, e.cumulativeLength = 0, e.length = 0);
+		}), e || (this.cache = null)), this.lastToken = -1;
+	}
+	reset() {
+		this.resetCache(), this.outputConversation = [], this.initialPrompt = null, this.rawOutput = [], this.lastLoss = null, this.emit("reset");
+	}
+	dispose() {
+		this.reset();
+	}
+	initialise(e, t) {
+		if (this.cache && t?.noCache && this.reset(), this.initialPrompt = e || null, this.lastToken === -1 ? this.outputConversation = (this.initialPrompt || []).slice() : e && e.length > this.outputConversation.length && (this.outputConversation = (this.initialPrompt || []).slice(), this.resetCache()), !this.cache && !t?.noCache && (this.model.config.modelType !== "GenAI_NanoGPT_v1" || this.model.config.useRope)) {
+			let e = Array(this.model.config.nLayer);
+			for (let t = 0; t < this.model.config.nLayer; t++) e[t] = {
+				k: void 0,
+				v: void 0,
+				length: 0,
+				cumulativeLength: 0
+			};
+			this.cache = e, this.lastToken = -1;
+		}
+		let r = this.tokeniser.trained ? this.tokeniser : new n(y(v, this.tokeniser.vocabSize));
+		this.actualTokeniser = r, t?.loraName ? this.model.attachLoRA(t.loraName) : this.model.hasLoRA() && this.model.detachLoRA();
+	}
+	async step(e, t) {
+		let n = {
+			...t,
+			maxLength: 1
+		};
+		return _(e) ? this.generate(e, n) : this.generate({
+			...e,
+			...n
+		});
+	}
+	async generate(e, t) {
+		let n;
+		if (Array.isArray(e) ? n = e : typeof e == "object" && (t = e), this.processingJob) {
+			if (this.jobQueue.length > 10) throw Error("Job queue is too long, rejecting new job");
+			return new Promise((e, r) => {
+				this.jobQueue.push({
+					prompt: n,
+					options: t,
+					resolve: e,
+					reject: r
+				});
+			});
+		}
+		this.processingJob = !0, this.startTime = Date.now();
+		try {
+			let e = await this.startJob(n, t);
+			if (this.processingJob = !1, this.jobQueue.length > 0) {
+				let e = this.jobQueue.shift();
+				this.generate(e.prompt || [], e.options).then(e.resolve).catch(e.reject);
+			}
+			return e;
+		} catch (e) {
+			throw this.processingJob = !1, e;
+		}
+	}
+	async startJob(e, t) {
+		this.initialise(e, t), this.active = !0, this.model.metaData.generationSettings = t, t?.maxLength !== 1 && this.emit("start");
+		let n = await this._generate(t, !!e);
+		if (this.active = !1, this.startTime !== null) {
+			let e = Date.now(), n = e - this.startTime;
+			this.startTime = null, this.model.metaData.actionLog = this.model.metaData.actionLog || [], this.model.metaData.actionLog.push({
+				action: "generate",
+				timestamp: e,
+				duration: n,
+				tokensProcessed: this.rawOutput.length,
+				options: t || {}
+			});
+		}
+		return this.emit("stop"), n;
+	}
+	getQueueLength() {
+		return this.jobQueue.length;
+	}
+	stop() {
+		this.active = !1;
+	}
+	getConversation() {
+		return this.outputConversation;
+	}
+	getRawOutput() {
+		return this.rawOutput;
+	}
+};
+//#endregion
+export { b as default, _ as isConversation };

package/dist/inference/tokenisePrompt.d.ts

@@ -0,0 +1,4 @@
+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>;

package/dist/inference/tokenisePrompt.js

@@ -0,0 +1,13 @@
+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", o;
+		return o = i?.nonConversational ? a && i?.continuation ? [t.bosToken, ...t.encode(r[r.length - 1].content)] : t.encodeAsSequence(r, !0) : t.encodeConversation(r, !0), o.length > n && (o = o.slice(-n)), e([o], [1, o.length], "int32");
+	} 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");
+	}
+}
+//#endregion
+export { t as default };

package/dist/inference/types.d.ts

@@ -1,16 +1,52 @@
 import { Conversation } from '../tokeniser/type';
+import { Tensor } from '@tensorflow/tfjs-core';
+export interface IGeneratorOutput {
+    outputTensor: Tensor;
+    token: number;
+    text: string;
+    confidence: number | null;
+    score: number | null;
+    logits: number[] | null;
+    scores: number[] | null;
+    hiddenStates: number[][] | null;
+    attention: number[][][][] | null;
+    loss: number | null;
+    multinomialRand: number | null;
+    terminated: boolean;
+}
 export interface GeneratorConversation extends Conversation {
     _completed?: boolean;
     _timestamp?: number;
+    _output?: IGeneratorOutput[];
 }
-export interface GenerateOptions {
+export interface IGenerateOptions {
+    input?: Conversation[] | string; /** Optional prompt. */
+    previous_response_id?: string;
     temperature?: number;
-    topK?: number;
-    topP?: number;
+    topK?: number; /** Select the top K best tokens. */
+    topP?: number; /** Select the top P best tokens (nucleus sampling). */
     usePadding?: boolean;
-    attentionScores?: boolean;
-    includeProbabilities?: boolean;
-    embeddings?: 'embedding' | 'logits' | 'softmax' | 'all';
-    targets?: number[];
-    loraName?: string;
+    outputAttention?: boolean; /** All attention weights. */
+    outputScores?: boolean; /** Softmax scores of the final logits (all tokens). */
+    outputConfidence?: boolean; /** Entropy-based confidence in [0, 1] from the softmax distribution. */
+    outputScore?: boolean; /** Selected token score (probability) */
+    outputLogits?: boolean; /** Final output logits. */
+    outputHiddenStates?: 'embedding' | 'logits' | 'softmax' | 'all';
+    outputLoss?: boolean;
+    outputMultinomialRand?: boolean;
+    targets?: number[]; /** Optional target tokens for loss calculation. */
+    loraName?: string; /** Optional LoRA name to use for inference. */
+    maxLength?: number; /** Maximum length of the generated text. */
+    noCache?: boolean; /** Do not use key/value cache for attention. */
+    allowSpecial?: boolean; /** Keep special tokens in the output. */
+    nonConversational?: boolean; /** Do not use turn taking tokens. */
+    continuation?: boolean; /** Continue the previous response without adding a new turn. */
+    chunkSize?: number; /** Number of tokens to generate before returning a partial response. */
+    background?: boolean; /** Run the generation in the background and return immediately. */
+    _onChunk?: (output: IGeneratorOutput) => void | Promise<void>;
+}
+export interface IGeneratorResponse {
+    output: GeneratorConversation[] | null;
+    id: string;
+    done: boolean;
 }

package/dist/inference/utilities.d.ts

@@ -0,0 +1,9 @@
+import { IGeneratorOutput } from './types';
+/**
+ * Confidence derived from normalized entropy of the probability distribution.
+ * Returns a value in [0, 1], where 1 is fully confident (delta-like distribution)
+ * and 0 is maximally uncertain (uniform distribution).
+ */
+export declare function getTokenConfidence(probabilities: number[]): number;
+export declare function getAttention(output: IGeneratorOutput, layer: number, head: number): number[] | null;
+export declare function getHiddenState(output: IGeneratorOutput, layer: number): number[] | null;

package/dist/inference/utilities.js

@@ -0,0 +1,20 @@
+//#region lib/inference/utilities.ts
+function e(e) {
+	if (!e.length) return 0;
+	let t = 0;
+	for (let n of e) n > 0 && (t -= n * Math.log(n));
+	let n = Math.log(e.length);
+	if (n <= 0) return 1;
+	let r = t / n;
+	return Math.min(1, Math.max(0, 1 - r));
+}
+function t(e, t, n) {
+	if (!e.attention || !e.attention[t] || !e.attention[t][n]) return null;
+	let r = e.attention[t][n].length;
+	return e.attention[t][n][r - 1];
+}
+function n(e, t) {
+	return !e.hiddenStates || !e.hiddenStates[t] ? null : e.hiddenStates[t];
+}
+//#endregion
+export { t as getAttention, n as getHiddenState, e as getTokenConfidence };