@genai-fi/nanogpt 0.7.2 → 0.8.0

package/dist/Generator.d.ts

@@ -1,19 +1,51 @@
-import { default as NanoGPT, GenerateOptions } from './NanoGPTModel';
 import { ITokeniser } from './tokeniser/type';
 import { default as EE } from 'eventemitter3';
+import { default as Model, ModelForwardAttributes } from './models/model';
+export interface GenerateOptions {
+    temperature?: number;
+    topK?: number;
+    topP?: number;
+    usePadding?: boolean;
+    attentionScores?: boolean;
+    includeProbabilities?: boolean;
+    embeddings?: boolean;
+}
 export interface IGenerateOptions extends GenerateOptions {
     maxLength?: number;
     noCache?: boolean;
 }
+/**
+ * Text generator using a NanoGPT model and a tokeniser.
+ * This uses the forward method of the model to generate text token by token, including options for temperature, top-k, and top-p sampling.
+ */
 export default class Generator extends EE<'start' | 'stop' | 'tokens'> {
     private readonly model;
     private readonly tokeniser;
     private active;
-    constructor(model: NanoGPT, tokeniser: ITokeniser);
+    private cache;
+    private initialPrompt;
+    private outputText;
+    private actualTokeniser;
+    private lastToken;
+    private attentionData;
+    private probabilitiesData;
+    private embeddingsData;
+    private tokens;
+    constructor(model: Model<ModelForwardAttributes>, tokeniser: ITokeniser);
     private tokenisePrompt;
-    private generateNoCache;
     private processResponse;
-    private generateCache;
+    /** Generate logits and select a token. */
+    private _generateToken;
+    /** Generate multiple tokens in a loop and produce text */
+    private _generate;
+    reset(): void;
+    dispose(): void;
+    private initialise;
+    step(prompt?: string, options?: IGenerateOptions): Promise<string>;
     generate(prompt?: string, options?: IGenerateOptions): Promise<string>;
     stop(): void;
+    getText(): string;
+    getAttentionData(): number[][][][];
+    getProbabilitiesData(): number[][][];
+    getTokens(): number[];
 }

package/dist/Generator.js

@@ -1,15 +1,15 @@
-import { E as u } from "./index-Dwqa6Zy2.js";
-import "./index-BoWRt-10.js";
+import { E as z } from "./index-Dwqa6Zy2.js";
+import { B as A, C as L, E as C, a5 as I, t as O, k as R } from "./index-CUQrfsw_.js";
 import "./ops/cpu/attentionMask.js";
 import "./ops/webgl/attentionMask.js";
 import "./ops/grads/attentionMask.js";
 import "./ops/cpu/qkv.js";
 import "./ops/webgl/qkv.js";
 import "./ops/grads/qkv.js";
-import "./random_width-sZORGo5k.js";
-import "./register_all_kernels-BwDSRN-f.js";
+import { p as _ } from "./random_width-D8Pwy_na.js";
+import { t as K } from "./register_all_kernels-DUshvVWP.js";
 import "./index-Tf7vU29b.js";
-import "./dataset-CtdBYwjo.js";
+import "./dataset-CJmEGu6D.js";
 import "./ops/cpu/rope.js";
 import "./ops/webgl/rope.js";
 import "./ops/grads/rope.js";
@@ -29,7 +29,7 @@ import "./ops/webgl/gatherSub.js";
 import "./ops/cpu/scatterSub.js";
 import "./ops/webgl/scatterSub.js";
 import "./jszip.min-CjP2V1VV.js";
-import f from "./tokeniser/CharTokeniser.js";
+import M from "./tokeniser/CharTokeniser.js";
 import "./ops/cpu/adamAdjust.js";
 import "./ops/webgl/adamAdjust.js";
 import "./ops/cpu/adamMoments.js";
@@ -37,12 +37,42 @@ import "./ops/webgl/adamMoments.js";
 import "./papaparse.min-C8l2Kvo1.js";
 import "./ops/cpu/gelu.js";
 import "./ops/webgl/gelu.js";
-import "./gelu-C-dPj6Ku.js";
+import "./gelu-Bd3UBBxg.js";
 import "./ops/webgl/log.js";
-import { t as d } from "./tensor2d-wxPAnDQy.js";
-import { c as g } from "./concat-CsxrgovM.js";
-const k = [
-  ...Array.from({ length: 95 }, (a, t) => String.fromCharCode(t + 32)),
+import $ from "./utilities/multinomialCPU.js";
+import { r as x } from "./reshape-DEfQGSin.js";
+import { t as P } from "./tensor2d-BHy261cI.js";
+import { s as v } from "./softmax-DLoZWYBx.js";
+import { g as q } from "./gather-C1siEkdp.js";
+import { c as G } from "./concat-BmDqqFsa.js";
+/**
+ * @license
+ * Copyright 2020 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+function N(h, t, e, i = !1) {
+  const o = L(h, "logits", "multinomial"), s = o.size, n = o.rank;
+  if (s < 2)
+    throw new Error(`Error in multinomial: you need at least 2 outcomes, but got ${s}.`);
+  if (n > 2)
+    throw new Error(`Rank of probabilities must be 1 or 2, but is ${n}`);
+  e = e || Math.random();
+  const a = { logits: n === 1 ? x(o, [1, -1]) : o }, p = { numSamples: t, seed: e, normalized: i }, l = C.runKernel(I, a, p);
+  return n === 1 ? x(l, [l.size]) : l;
+}
+const S = /* @__PURE__ */ A({ multinomial_: N }), B = [
+  ...Array.from({ length: 95 }, (h, t) => String.fromCharCode(t + 32)),
   // ASCII
   // Spanish accented letters and punctuation
   ..."áéíóúüñ¿¡",
@@ -53,80 +83,164 @@ const k = [
   // Cyrillic letters
   ..."абвгдеёжзийклмнопрстуфхцчшщъыьэюяАБВГДЕЁЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЫЬЭЮЯ"
 ];
-function w(a, t) {
-  return a.length === t ? a : a.length > t ? a.slice(0, t) : a.concat(Array(t - a.length).fill(""));
+function H(h, t) {
+  return h.length === t ? h : h.length > t ? h.slice(0, t) : h.concat(Array(t - h.length).fill(""));
 }
-class pt extends u {
-  constructor(t, o) {
-    super(), this.model = t, this.tokeniser = o;
+class Mt extends z {
+  constructor(t, e) {
+    super(), this.model = t, this.tokeniser = e, this.actualTokeniser = e;
   }
   active = !1;
-  async tokenisePrompt(t, o) {
-    const r = o ? await t.tokenise([o], !0) : [[t.eosToken]];
-    return d(r, [1, r[0].length], "int32");
-  }
-  async generateNoCache(t, o, r) {
-    let i = await this.tokenisePrompt(t, o), s = o || "";
-    const n = r?.maxLength ?? 1e3;
-    for (let m = 0; m < n && this.active; m++) {
-      const {
-        output: e,
-        attention: p,
-        probabilities: c
-      } = await this.model.generate(i, void 0, r), h = i;
-      i = g([i, e], 1), h.dispose();
-      const l = await this.processResponse(t, e, p, c);
-      if (e.dispose(), l === null)
-        break;
-      s += l;
-    }
-    return i.dispose(), s;
+  cache = null;
+  initialPrompt = null;
+  outputText = "";
+  actualTokeniser;
+  lastToken = -1;
+  attentionData = [];
+  probabilitiesData = [];
+  embeddingsData = [];
+  tokens = [];
+  async tokenisePrompt(t, e) {
+    const i = e ? await t.tokenise([e], !0) : [[t.eosToken]];
+    return P(i, [1, i[0].length], "int32");
   }
-  async processResponse(t, o, r, i) {
-    const s = (await o.array())[0][0];
-    if (s === this.tokeniser.eosToken)
+  async processResponse(t, e, i, o) {
+    const s = (await e.array())[0][0];
+    if (this.lastToken = s, s === this.tokeniser.eosToken)
       return null;
     const n = await t.decode([s]);
-    let m;
-    r && (m = await Promise.all(r.map((p) => p.array().then((c) => c))), r.forEach((p) => p.dispose()));
-    let e;
-    return i && (e = await i.array(), i.dispose()), this.emit("tokens", [s], n, m, e), n;
+    if (i) {
+      const d = await Promise.all(i.map((a) => a.array().then((p) => p)));
+      i.forEach((a) => a.dispose()), this.attentionData.push(d);
+    }
+    if (o) {
+      const d = await o.array();
+      o.dispose(), this.probabilitiesData.push(d);
+    }
+    return this.tokens.push(s), this.emit("tokens", [s], n), n;
   }
-  async generateCache(t, o, r) {
-    let i = await this.tokenisePrompt(t, o), s = o || "";
-    const n = new Array(this.model.config.gpt.nLayer);
-    for (let e = 0; e < this.model.config.gpt.nLayer; e++)
-      n[e] = { k: void 0, v: void 0, length: 0, cumulativeLength: 0 };
-    const m = r?.maxLength ?? 1e3;
-    for (let e = 0; e < m && this.active; e++) {
+  /** Generate logits and select a token. */
+  async _generateToken(t, e, i) {
+    const o = i?.temperature ?? 1, s = i?.topK, n = i?.topP, d = i?.usePadding ?? !1, a = {
+      training: !1,
+      attentionScores: i?.attentionScores ? {
+        attentionOut: []
+      } : void 0,
+      cache: e,
+      outputEmbeddings: i?.embeddings ?? !1
+    }, p = O(() => {
+      const r = t, m = r.shape[1], u = m <= this.model.config.blockSize ? r : r.slice(
+        [0, m - this.model.config.blockSize],
+        [r.shape[0], this.model.config.blockSize]
+      ), k = d ? this.model.config.blockSize - u.shape[1] : 0, b = k > 0 ? _(u, [
+        [0, 0],
+        [0, k]
+      ]) : u, [f] = this.model.forward(a, b), y = f.shape[1] - 1 - k, c = f.slice([0, y, 0], [f.shape[0], 1, f.shape[2]]);
+      return a.attentionScores?.attentionOut && a.attentionScores.attentionOut.forEach((T, E) => {
+        T.shape[1] !== 1 && (a.attentionScores.attentionOut[E] = R(
+          T.slice([0, y, 0], [T.shape[0], 1, T.shape[2]])
+        ), T.dispose());
+      }), f.dispose(), c.div(o).squeeze([1]);
+    });
+    let l;
+    if (n) {
+      const r = v(p), m = await r.array();
+      r.dispose();
+      const u = m[0].map((c, g) => ({ prob: c, index: g })).sort((c, g) => g.prob - c.prob);
+      let k = 0;
+      const b = new Array(u.length).fill(0);
+      for (const c of u)
+        if (k += c.prob, b[c.index] = c.prob, k >= n)
+          break;
+      const f = b.reduce((c, g) => c + g, 0), y = b.map((c) => c / f);
+      l = $(y);
+    } else if (s) {
+      const { values: r, indices: m } = K(p, s), u = S(r, 1);
+      l = q(m, u, 1), r.dispose(), m.dispose(), u.dispose();
+    } else
+      l = S(p, 1);
+    let w;
+    i?.includeProbabilities && (w = v(p)), a.embeddings && this.embeddingsData.push(
+      await Promise.all(
+        a.embeddings.map(async (r) => {
+          const m = await r.array();
+          return r.dispose(), m;
+        })
+      )
+    );
+    const D = l.reshape([1, 1]);
+    return l.dispose(), l = D, p.dispose(), { output: l, probabilities: w, attention: a.attentionScores?.attentionOut };
+  }
+  /** Generate multiple tokens in a loop and produce text */
+  async _generate(t) {
+    let e = this.lastToken >= 0 && this.cache ? P([this.lastToken], [1, 1], "int32") : await this.tokenisePrompt(this.actualTokeniser, this.outputText);
+    const i = t?.maxLength ?? 1e3;
+    for (let o = 0; o < i && this.active; o++) {
       const {
-        output: p,
-        probabilities: c,
-        attention: h
-      } = await this.model.generate(i, n, {
-        ...r,
-        usePadding: !1
+        output: s,
+        probabilities: n,
+        attention: d
+      } = await this._generateToken(e, this.cache ? this.cache : void 0, {
+        ...t,
+        usePadding: !this.cache
       });
-      i.dispose(), i = p;
-      const l = await this.processResponse(t, p, h, c);
-      if (l === null)
+      if (this.cache)
+        e.dispose(), e = s;
+      else {
+        const p = e;
+        e = G([e, s], 1), p.dispose();
+      }
+      const a = await this.processResponse(this.actualTokeniser, s, d, n);
+      if (this.cache || s.dispose(), a === null)
         break;
-      s += l;
+      this.outputText += a;
     }
-    return n.forEach((e) => {
-      e && (e.k && e.k.dispose(), e.v && e.v.dispose());
-    }), i.dispose(), s;
+    return e.dispose(), this.outputText;
+  }
+  reset() {
+    this.cache && (this.cache.forEach((t) => {
+      t && (t.k && t.k.dispose(), t.v && t.v.dispose());
+    }), this.cache = null), this.outputText = "", this.initialPrompt = null, this.lastToken = -1, this.attentionData = [], this.probabilitiesData = [], this.tokens = [];
   }
-  async generate(t, o) {
-    const r = t && t.length > this.model.config.gpt.blockSize ? t.slice(-this.model.config.gpt.blockSize) : t;
-    this.active = !0, this.emit("start");
-    const i = this.tokeniser.trained ? this.tokeniser : new f(w(k, this.tokeniser.vocabSize)), n = await (this.model.config.gpt.useRope && !o?.noCache ? this.generateCache(i, r, o) : this.generateNoCache(i, r, o));
-    return this.active = !1, this.emit("stop"), n;
+  dispose() {
+    this.reset();
+  }
+  initialise(t, e) {
+    const i = t && t.length > this.model.config.blockSize ? t.slice(-this.model.config.blockSize) : t ?? null;
+    if (this.cache && e?.noCache && this.reset(), this.initialPrompt = i || null, this.lastToken === -1 && (this.outputText = this.initialPrompt || ""), !this.cache && !e?.noCache && this.model.config.useRope) {
+      const s = new Array(this.model.config.nLayer);
+      for (let n = 0; n < this.model.config.nLayer; n++)
+        s[n] = { k: void 0, v: void 0, length: 0, cumulativeLength: 0 };
+      this.cache = s, this.lastToken = -1;
+    }
+    const o = this.tokeniser.trained ? this.tokeniser : new M(H(B, this.tokeniser.vocabSize));
+    this.actualTokeniser = o;
+  }
+  async step(t, e) {
+    const i = { ...e, maxLength: 1 };
+    return this.generate(t, i);
+  }
+  async generate(t, e) {
+    this.initialise(t, e), this.active = !0, this.emit("start");
+    const o = await this._generate(e);
+    return this.active = !1, this.emit("stop"), o;
   }
   stop() {
     this.active = !1;
   }
+  getText() {
+    return this.outputText;
+  }
+  getAttentionData() {
+    return this.attentionData;
+  }
+  getProbabilitiesData() {
+    return this.probabilitiesData;
+  }
+  getTokens() {
+    return this.tokens;
+  }
 }
 export {
-  pt as default
+  Mt as default
 };

package/dist/{RealDiv-Dy0p8Bvo.js → RealDiv-N8TpOMYv.js}

@@ -1,10 +1,10 @@
-import { aq as T, ac as E, p as O, j as V, ay as B, Y as F, U, az as j } from "./index-BoWRt-10.js";
-import { r as $ } from "./Reshape-DH5srBP0.js";
-import { g as A, a as k, b as C, c as N, e as R } from "./axis_util-BzbKo31C.js";
-import { t as K, m as W } from "./shared-wS99K7_n.js";
-import { c as _ } from "./backend_util-TE7aTPhZ.js";
-import { f as y } from "./gpgpu_math-DGNLNL4I.js";
-import { g as G, b as L } from "./kernel_funcs_utils-BYKWV8Aa.js";
+import { as as T, af as E, p as O, j as V, aA as B, a0 as F, X as j, aB as K } from "./index-CUQrfsw_.js";
+import { r as $ } from "./Reshape-Bo8HzP8V.js";
+import { g as A, a as k, b as C, c as N, e as R } from "./axis_util-DubwyOhW.js";
+import { t as U, m as W } from "./shared-D1elLckx.js";
+import { c as _ } from "./backend_util-BJ-_jSeK.js";
+import { f as y } from "./gpgpu_math-TFLxaLkw.js";
+import { g as G, b as L } from "./kernel_funcs_utils-P9aFa232.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -105,7 +105,7 @@ class w {
  * limitations under the License.
  * =============================================================================
  */
-class q {
+class X {
   constructor(s, e) {
     this.variableNames = ["x"];
     const { windowSize: t, batchSize: n, inSize: l, outSize: r } = s;
@@ -229,7 +229,7 @@ class q {
  * limitations under the License.
  * =============================================================================
  */
-function X(a) {
+function q(a) {
   const s = [];
   for (; s.length === 0 || s[s.length - 1].outSize !== 1; ) {
     const e = s.length ? s[s.length - 1].outSize : a[1], t = _(e);
@@ -242,12 +242,12 @@ function X(a) {
   return s;
 }
 function P(a, s, e, t) {
-  const n = X(a.shape);
+  const n = q(a.shape);
   let l = a;
   for (let r = 0; r < n.length; r++) {
     const { inSize: i, windowSize: c, outSize: o } = n[r];
     let u, p;
-    e === "mean" ? u = r === 0 ? new w({ windowSize: c, inSize: i, batchSize: a.shape[0], outSize: o }, i) : new w({ windowSize: c, inSize: i, batchSize: a.shape[0], outSize: o }) : u = new q({ windowSize: c, inSize: i, batchSize: a.shape[0], outSize: o }, e), p = l, l = t.runWebGLProgram(u, [l], s), p.dataId !== a.dataId && t.disposeIntermediateTensorInfo(p);
+    e === "mean" ? u = r === 0 ? new w({ windowSize: c, inSize: i, batchSize: a.shape[0], outSize: o }, i) : new w({ windowSize: c, inSize: i, batchSize: a.shape[0], outSize: o }) : u = new X({ windowSize: c, inSize: i, batchSize: a.shape[0], outSize: o }, e), p = l, l = t.runWebGLProgram(u, [l], s), p.dataId !== a.dataId && t.disposeIntermediateTensorInfo(p);
   }
   return l;
 }
@@ -459,7 +459,7 @@ function te(a) {
       const I = e.texData.get(d.dataId).values, m = new Array(i);
       for (let v = 0; v < m.length; v++)
         m[v] = n.shape[u[v]];
-      const z = K(I, n.shape, n.dtype, u, m);
+      const z = U(I, n.shape, n.dtype, u, m);
       d = e.makeTensorInfo(m, n.dtype);
       const M = e.texData.get(d.dataId);
       M.values = z;
@@ -482,7 +482,7 @@ function te(a) {
   return p && e.disposeIntermediateTensorInfo(d), x;
 }
 const he = {
-  kernelName: U,
+  kernelName: j,
   backendName: "webgl",
   kernelFunc: te
 };
@@ -525,7 +525,7 @@ return a / b;`, se = `
 
   return result;
 `, ne = L({ opSnippet: ae, packedOpSnippet: se, checkOutOfBounds: !0 }), fe = {
-  kernelName: j,
+  kernelName: K,
   backendName: "webgl",
   kernelFunc: ne
 };

package/dist/{Reshape-DvudQDvJ.js → Reshape-B-lWQRnF.js}

@@ -1,4 +1,4 @@
-import { j as h, a3 as d, l as c, K as m } from "./index-BoWRt-10.js";
+import { j as h, a4 as d, n as c, U as m } from "./index-CUQrfsw_.js";
 /**
  * @license
  * Copyright 2021 Google LLC. All Rights Reserved.

package/dist/{Reshape-DH5srBP0.js → Reshape-Bo8HzP8V.js}

@@ -1,5 +1,5 @@
-import { j as c, a3 as C, l as f, K as R } from "./index-BoWRt-10.js";
-import { u as g, g as I, a as x, b as F, c as $, d as u, e as l, i as m } from "./gpgpu_math-DGNLNL4I.js";
+import { j as c, a4 as C, n as f, U as R } from "./index-CUQrfsw_.js";
+import { u as g, g as I, a as x, b as F, c as $, d as u, e as m, i as l } from "./gpgpu_math-TFLxaLkw.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -82,14 +82,14 @@ function v(s, t) {
 function b(s, t, i) {
   const a = [
     u(s.shape),
-    ...l(s.shape)
+    ...m(s.shape)
   ], e = {
     dtype: s.dtype,
     shape: a,
     dataId: s.dataId
   }, o = [
     u(t),
-    ...l(t)
+    ...m(t)
   ], r = new S(o, a), p = !0, n = [a], h = i.runWebGLProgram(r, [e], s.dtype, n, p);
   return { dataId: h.dataId, shape: t, dtype: h.dtype };
 }
@@ -113,7 +113,7 @@ function y(s) {
   const { inputs: t, backend: i, attrs: a } = s, { x: e } = t, { shape: o } = a, r = i, p = c(e.shape), n = C(o, p), h = c(n);
   f(p === h, () => `The new shape (${n}) has ${h} elements and the old shape (${e.shape}) has ${p} elements. The new shape and old shape must have the same number of elements.`);
   const d = r.texData.get(e.dataId);
-  return d.isPacked && !m(e.shape, n) && !(d.texture !== null && m(d.shape, n)) ? b(e, n, r) : (r.incRef(e.dataId), { dataId: e.dataId, shape: n, dtype: e.dtype });
+  return d.isPacked && !l(e.shape, n) && !(d.texture !== null && l(d.shape, n)) ? b(e, n, r) : (r.incRef(e.dataId), { dataId: e.dataId, shape: n, dtype: e.dtype });
 }
 const U = {
   kernelName: R,

package/dist/TeachableLLM.d.ts

@@ -1,11 +1,11 @@
-import { GPTConfig } from './config';
+import { GPTConfig } from './models/config';
 import { ITokeniser } from './tokeniser/type';
-import { default as NanoGPT, TrainingLogEntry } from './NanoGPTModel';
-import { SaveOptions } from './utilities/save';
+import { SaveOptions } from './loader/save';
 import { default as Generator, IGenerateOptions } from './Generator';
 import { default as Trainer, ITrainerOptions } from './Trainer';
 import { default as MemoryProfiler } from './utilities/profile';
-import { TrainingProgress } from './training/Trainer';
+import { TrainingLogEntry, TrainingProgress } from './training/Trainer';
+import { default as Model, ModelForwardAttributes } from './models/model';
 type TeachableLLMStatus = 'warmup' | 'awaitingTokens' | 'ready' | 'training' | 'loading' | 'busy' | 'error';
 interface TeachableLLMMeta {
     name?: string;
@@ -20,12 +20,12 @@ export default class TeachableLLM {
     private _status;
     private _memoryRequirements?;
     meta: TeachableLLMMeta;
-    constructor(tokeniser?: ITokeniser, model?: NanoGPT);
+    constructor(tokeniser?: ITokeniser, model?: Model<ModelForwardAttributes>);
     get vocab(): string[];
     /** Model is fully loaded */
     get loaded(): boolean;
     get config(): GPTConfig;
-    get model(): NanoGPT;
+    get model(): Model<ModelForwardAttributes>;
     get tokeniser(): ITokeniser;
     get status(): TeachableLLMStatus;
     /** Model is both ready and not busy */