@genai-fi/nanogpt 0.2.12 → 0.3.1

package/dist/Generator.js

@@ -1,49 +1,52 @@
 import { E as u } from "./index-Dwqa6Zy2.js";
-class f extends u {
+import "./index-pWA4_lUh.js";
+import { t as d } from "./tensor2d-ZuQSh2D-.js";
+import { c as k } from "./concat-BIZS_td9.js";
+class w extends u {
   constructor(s, e) {
     super(), this.model = s, this.tokeniser = e;
   }
   active = !1;
   async tokenisePrompt(s) {
     const e = s ? await this.tokeniser.tokenise([s], !0) : [[this.tokeniser.eosToken]];
-    return this.model.tf.tensor2d(e, [1, e[0].length], "int32");
+    return d(e, [1, e[0].length], "int32");
   }
   async generateNoCache(s, e) {
-    let t = await this.tokenisePrompt(s), i = s || "";
-    const o = e?.maxLength ?? 1e3;
-    for (let a = 0; a < o && this.active; a++) {
+    let t = await this.tokenisePrompt(s), o = s || "";
+    const n = e?.maxLength ?? 1e3;
+    for (let a = 0; a < n && this.active; a++) {
       const {
-        output: n,
+        output: i,
         attention: c,
         probabilities: l
       } = this.model.generate(t, void 0, e), h = t;
-      t = this.model.tf.concat([t, n], 1), h.dispose();
-      const r = await this.processResponse(n, c, l);
-      if (n.dispose(), r === null)
+      t = k([t, i], 1), h.dispose();
+      const r = await this.processResponse(i, c, l);
+      if (i.dispose(), r === null)
         break;
-      i += r;
+      o += r;
     }
-    return t.dispose(), i;
+    return t.dispose(), o;
   }
   async processResponse(s, e, t) {
-    const i = (await s.array())[0][0];
-    if (i === this.tokeniser.eosToken)
+    const o = (await s.array())[0][0];
+    if (o === this.tokeniser.eosToken)
       return null;
-    const o = await this.tokeniser.decode([i]);
+    const n = await this.tokeniser.decode([o]);
     let a;
     e && (a = await e.array(), e.dispose());
-    let n;
-    return t && (n = await t.array(), t.dispose()), this.emit("tokens", [i], o, a, n), o;
+    let i;
+    return t && (i = await t.array(), t.dispose()), this.emit("tokens", [o], n, a, i), n;
   }
   async generateCache(s, e) {
-    let t = await this.tokenisePrompt(s), i = s || "";
-    const o = new Array(this.model.config.nLayer).fill(void 0), a = e?.maxLength ?? 1e3;
-    for (let n = 0; n < a && this.active; n++) {
+    let t = await this.tokenisePrompt(s), o = s || "";
+    const n = new Array(this.model.config.nLayer).fill(void 0), a = e?.maxLength ?? 1e3;
+    for (let i = 0; i < a && this.active; i++) {
       const {
         output: c,
         attention: l,
         probabilities: h
-      } = this.model.generate(t, o, {
+      } = this.model.generate(t, n, {
         ...e,
         usePadding: !1
       });
@@ -51,20 +54,22 @@ class f extends u {
       const r = await this.processResponse(c, l, h);
       if (r === null)
         break;
-      i += r;
+      o += r;
     }
-    return t.dispose(), i;
+    return n.forEach((i) => {
+      i && (i.k.dispose(), i.v.dispose());
+    }), t.dispose(), o;
   }
   async generate(s, e) {
     const t = s && s.length > this.model.config.blockSize ? s.slice(-this.model.config.blockSize) : s;
     this.active = !0, this.emit("start");
-    const o = await (this.model.config.useRope && !e?.noCache ? this.generateCache(t, e) : this.generateNoCache(t, e));
-    return this.active = !1, this.emit("stop"), o;
+    const n = await (this.model.config.useRope && !e?.noCache ? this.generateCache(t, e) : this.generateNoCache(t, e));
+    return this.active = !1, this.emit("stop"), n;
   }
   stop() {
     this.active = !1;
   }
 }
 export {
-  f as default
+  w as default
 };

package/dist/NanoGPTModel.d.ts

@@ -1,8 +1,8 @@
-import { default as TF } from '@tensorflow/tfjs';
 import { GPTConfig } from './config';
 import { KVCache } from './layers/CausalSelfAttention';
 import { default as MemoryProfiler } from './utilities/profile';
 import { default as BaseLayer } from './layers/BaseLayer';
+import { Tensor, Variable } from '@tensorflow/tfjs-core';
 export interface TrainingLogEntry {
     loss: number;
     valLoss?: number;
@@ -26,12 +26,11 @@ export default class NanoGPT extends BaseLayer {
     private blocks;
     private lnF;
     private ropeCache?;
-    readonly tf: typeof TF;
     log: TrainingLogEntry[];
-    constructor(tf: typeof TF, config?: Partial<GPTConfig>);
-    get variables(): TF.Variable[];
-    saveWeights(): Map<string, TF.Tensor[]>;
-    loadWeights(weights: Map<string, TF.Tensor[]>): void;
+    constructor(config?: Partial<GPTConfig>);
+    get variables(): Variable[];
+    saveWeights(): Map<string, Tensor[]>;
+    loadWeights(weights: Map<string, Tensor[]>): void;
     private inputPhase;
     setSkipMask(mask: boolean[]): void;
     setTrainableMask(mask: boolean[]): void;
@@ -40,15 +39,15 @@ export default class NanoGPT extends BaseLayer {
     private validateInput;
     private calculateLoss;
     private computeAttentionRollout;
-    forward(idx: TF.Tensor, targets?: TF.Tensor, training?: boolean, includeAttention?: boolean, cache?: (KVCache | undefined)[]): {
-        logits: TF.Tensor;
-        loss?: TF.Tensor;
-        attention?: TF.Tensor;
+    forward(idx: Tensor, targets?: Tensor, training?: boolean, includeAttention?: boolean, cache?: (KVCache | undefined)[]): {
+        logits: Tensor;
+        loss?: Tensor;
+        attention?: Tensor;
     };
-    generate(idx: TF.Tensor, cache?: (KVCache | undefined)[], options?: GenerateOptions): {
-        output: TF.Tensor;
-        attention?: TF.Tensor;
-        probabilities?: TF.Tensor;
+    generate(idx: Tensor, cache?: (KVCache | undefined)[], options?: GenerateOptions): {
+        output: Tensor;
+        attention?: Tensor;
+        probabilities?: Tensor;
     };
     getNumParams(): number;
     dispose(): void;

package/dist/NanoGPTModel.js

@@ -1,12 +1,98 @@
-import { defaultConfig as v } from "./config.js";
-import z from "./layers/TransformerBlock.js";
-import S from "./layers/TiedEmbedding.js";
-import _ from "./layers/RoPECache.js";
-import I from "./layers/RMSNorm.js";
-import { estimateParameterCount as F } from "./utilities/parameters.js";
-import { createSoftmaxCrossEntropyWithGrad as L } from "./training/sparseCrossEntropy.js";
-import P from "./layers/BaseLayer.js";
-class A extends P {
+import { defaultConfig as F } from "./config.js";
+import L from "./layers/TransformerBlock.js";
+import P from "./layers/TiedEmbedding.js";
+import C from "./layers/RoPECache.js";
+import q from "./layers/RMSNorm.js";
+import { estimateParameterCount as K } from "./utilities/parameters.js";
+import { createSoftmaxCrossEntropyWithGrad as N } from "./training/sparseCrossEntropy.js";
+import T from "./layers/BaseLayer.js";
+import { r as R, e as D, p as A } from "./random_width-oeUIlUZj.js";
+import { o as y, h as E, p as B, E as z, W as G, X as O, Y as Q, t as w, Z as X, f as _ } from "./index-pWA4_lUh.js";
+import { e as j, a as U } from "./exports_layers-tbTBcwMM.js";
+import { r as S } from "./reshape-C8CR_Bad.js";
+import { r as V } from "./range-CcDl05lo.js";
+import { g as Y } from "./gather-BPGW8RsB.js";
+import { s as Z } from "./softmax-Be_lsqUc.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 H(m, t) {
+  let e = E(m, "a", "mod"), o = E(t, "b", "mod");
+  [e, o] = B(e, o);
+  const i = { a: e, b: o };
+  return z.runKernel(G, i);
+}
+const J = /* @__PURE__ */ y({ mod_: H });
+/**
+ * @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 tt(m, t, e, o = !1) {
+  const i = E(m, "logits", "multinomial"), s = i.size, r = i.rank;
+  if (s < 2)
+    throw new Error(`Error in multinomial: you need at least 2 outcomes, but got ${s}.`);
+  if (r > 2)
+    throw new Error(`Rank of probabilities must be 1 or 2, but is ${r}`);
+  e = e || Math.random();
+  const n = { logits: r === 1 ? S(i, [1, -1]) : i }, h = { numSamples: t, seed: e, normalized: o }, a = z.runKernel(O, n, h);
+  return r === 1 ? S(a, [a.size]) : a;
+}
+const I = /* @__PURE__ */ y({ multinomial_: tt });
+/**
+ * @license
+ * Copyright 2018 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 et(m, t = 1, e = !0) {
+  const o = E(m, "x", "topk");
+  if (o.rank === 0)
+    throw new Error("topk() expects the input to be of rank 1 or higher");
+  const i = o.shape[o.shape.length - 1];
+  if (t < 0)
+    throw new Error(`'k' passed to topk() must be >= 0 but got ${t}`);
+  if (t > i)
+    throw new Error(`'k' passed to topk() must be <= the last dimension (${i}) but got ${t}`);
+  const s = { x: o }, r = { k: t, sorted: e }, [l, n] = z.runKernel(Q, s, r);
+  return { values: l, indices: n };
+}
+const ot = /* @__PURE__ */ y({ topk_: et });
+class kt extends T {
   config;
   wte;
   // Token embeddings
@@ -18,23 +104,22 @@ class A extends P {
   lnF;
   // Final layer norm
   ropeCache;
-  tf;
   log = [];
   // Training log
-  constructor(t, e = {}) {
-    super(), this.tf = t, this.config = { ...v, ...e }, this.wte = new S(t, {
+  constructor(t = {}) {
+    super(), this.config = { ...F, ...t }, this.wte = new P({
       vocabSize: this.config.vocabSize,
       embedDim: this.config.nEmbed,
       name: "token_embedding"
-    }), this.config.useRope === !1 ? this.wpe = this.tf.layers.embedding({
+    }), this.config.useRope === !1 ? this.wpe = j({
       inputDim: this.config.blockSize,
       outputDim: this.config.nEmbed,
       name: "positional_embedding",
-      embeddingsInitializer: this.tf.initializers.randomNormal({ mean: 0, stddev: 0.02 })
-    }) : this.ropeCache = new _(t, this.config), this.drop = this.tf.layers.dropout({ rate: this.config.dropout }), this.blocks = [];
-    for (let o = 0; o < this.config.nLayer; o++)
-      this.blocks.push(new z(this.tf, o, this.config, this.ropeCache));
-    this.lnF = new I(t, [this.config.nEmbed], 1e-8, "final_rms_norm");
+      embeddingsInitializer: R({ mean: 0, stddev: 0.02 })
+    }) : this.ropeCache = new C(this.config), this.drop = U({ rate: this.config.dropout }), this.blocks = [];
+    for (let e = 0; e < this.config.nLayer; e++)
+      this.blocks.push(new L(e, this.config, this.ropeCache));
+    this.lnF = new q([this.config.nEmbed], 1e-8, "final_rms_norm");
   }
   get variables() {
     return [
@@ -58,14 +143,11 @@ class A extends P {
     this.lnF.setWeights(t.get("final_rms_norm") || []);
   }
   inputPhase(t, e, o = !1) {
-    return this.tf.tidy(() => {
+    return w(() => {
       const i = this.wte.embed(t);
       if (this.config.useRope === !1) {
-        const [, s] = t.shape, l = this.config.blockSize, r = this.tf.range(0, s, 1, "int32"), n = this.tf.mod(
-          this.tf.add(r, this.tf.scalar(e, "int32")),
-          this.tf.scalar(l, "int32")
-        ), h = this.wpe.apply(n), c = i.add(h);
-        return this.drop.apply(c, { training: o });
+        const [, s] = t.shape, r = this.config.blockSize, l = V(0, s, 1, "int32"), n = J(X(l, _(e, "int32")), _(r, "int32")), h = this.wpe.apply(n), a = i.add(h);
+        return this.drop.apply(a, { training: o });
       } else
         return this.drop.apply(i, { training: o });
     });
@@ -103,7 +185,7 @@ class A extends P {
   }
   calculateLoss(t, e) {
     try {
-      return L()(t, e).mean();
+      return N()(t, e).mean();
     } catch (o) {
       throw console.error("Error computing loss:", o), new Error(`Loss computation failed: ${o}`);
     }
@@ -111,89 +193,79 @@ class A extends P {
   // Attention rollout per Abnar & Zuidema (2020)
   // Expects list of (B, T, T) attention matrices already averaged over heads.
   computeAttentionRollout(t) {
-    return this.tf.tidy(() => {
+    return w(() => {
       if (t.length === 0)
         throw new Error("No attentions for rollout");
       const [e, o, i] = t[0].shape;
       for (const s of t) {
-        const [l, r, n] = s.shape;
-        if (l !== e || r !== o || n !== i)
+        const [r, l, n] = s.shape;
+        if (r !== e || l !== o || n !== i)
           throw new Error(
-            `Inconsistent attention shapes in rollout: expected [${e},${o},${i}] got [${l},${r},${n}]`
+            `Inconsistent attention shapes in rollout: expected [${e},${o},${i}] got [${r},${l},${n}]`
           );
       }
       if (o === i) {
-        const s = this.tf.eye(i, i).expandDims(0);
-        let l = s.tile([e, 1, 1]);
-        for (const r of t) {
-          const n = r.add(s);
-          l = n.div(n.sum(-1, !0)).matMul(l);
-        }
-        return l;
-      }
-      if (o === 1) {
-        let s = null;
-        const l = this.tf.tensor1d([i - 1], "int32"), r = this.tf.oneHot(l, i).reshape([1, 1, i]).tile([e, 1, 1]);
-        l.dispose();
-        for (const n of t) {
-          let h = n.add(r);
-          h = h.div(h.sum(-1, !0)), s == null ? s = h : (s = s.mul(h), s = s.div(s.sum(-1, !0)));
+        const s = D(i, i).expandDims(0);
+        let r = s.tile([e, 1, 1]);
+        for (const l of t) {
+          const n = l.add(s);
+          r = n.div(n.sum(-1, !0)).matMul(r);
         }
-        return s;
+        return r;
       }
       throw new Error(`Unsupported attention shapes for rollout: [B=${e}, Q=${o}, K=${i}]`);
     });
   }
   forward(t, e, o = !1, i = !1, s) {
-    return this.validateInput(t), this.tf.tidy(() => {
+    return this.validateInput(t), w(() => {
       this.startMemory();
-      const l = s?.[0]?.length ?? 0;
-      let r = this.inputPhase(t, l, o);
+      const r = s?.[0]?.length ?? 0;
+      let l = this.inputPhase(t, r, o);
       const n = [];
       if (s && s.length !== this.blocks.length)
         throw console.error("Cache", s), new Error(`Cache length ${s.length} does not match number of blocks ${this.blocks.length}`);
-      for (let a = 0; a < this.blocks.length; a++) {
-        const d = r, g = this.blocks[a], {
-          output: m,
-          attention: b,
+      for (let c = 0; c < this.blocks.length; c++) {
+        const u = l, d = this.blocks[c], {
+          output: b,
+          attention: k,
           cache: f
-        } = g.call(r, o, i, s ? s[a] : void 0);
-        r = m, d.dispose(), i && b && n.push(b), s && f ? (s[a]?.k.dispose(), s[a]?.v.dispose(), s[a] = f) : f && (f.k.dispose(), f.v.dispose());
+        } = d.call(l, o, i, s ? s[c] : void 0);
+        l = b, u.dispose(), i && k && n.push(k), s && f ? (s[c]?.k.dispose(), s[c]?.v.dispose(), s[c] = f) : f && (f.k.dispose(), f.v.dispose());
       }
       let h;
-      i && n.length > 0 && (h = this.computeAttentionRollout(n)), r = this.lnF.apply(r);
-      const c = this.wte.project(r);
+      i && n.length > 0 && (h = this.computeAttentionRollout(n)), l = this.lnF.apply(l);
+      const a = this.wte.project(l);
       let p;
-      return e && (p = this.calculateLoss(c, e)), this.endMemory("Forward"), { logits: c, loss: p, attention: i ? h : void 0 };
+      return e && (p = this.calculateLoss(a, e)), this.endMemory("Forward"), { logits: a, loss: p, attention: i ? h : void 0 };
     });
   }
   generate(t, e, o) {
-    const i = o?.temperature ?? 1, s = o?.topK, l = o?.usePadding ?? !1, r = o?.includeAttention ?? !1;
-    return this.tf.tidy(() => {
-      const n = t, h = n.shape[1], c = h <= this.config.blockSize ? n : n.slice(
+    const i = o?.temperature ?? 1, s = o?.topK, r = o?.usePadding ?? !1, l = o?.includeAttention ?? !1;
+    return w(() => {
+      const n = t, h = n.shape[1], a = h <= this.config.blockSize ? n : n.slice(
         [0, h - this.config.blockSize],
         [n.shape[0], this.config.blockSize]
-      ), p = l ? this.config.blockSize - c.shape[1] : 0, a = p > 0 ? this.tf.pad(c, [
+      ), p = r ? this.config.blockSize - a.shape[1] : 0, c = p > 0 ? A(a, [
         [0, 0],
         [0, p]
-      ]) : c, { logits: d, attention: g } = this.forward(a, void 0, !1, r, e), m = d.shape[1] - 1 - p, b = d.slice([0, m, 0], [d.shape[0], 1, d.shape[2]]), f = g ? g.slice([0, m, 0], [g.shape[0], 1, g.shape[2]]) : void 0, k = b.div(i);
-      let u;
+      ]) : a, { logits: u, attention: d } = this.forward(c, void 0, !1, l, e), b = u.shape[1] - 1 - p, k = u.slice([0, b, 0], [u.shape[0], 1, u.shape[2]]), f = d ? d.slice([0, b, 0], [d.shape[0], 1, d.shape[2]]) : void 0, $ = k.div(i);
+      let g;
       if (s) {
-        const { values: y, indices: E } = this.tf.topk(k, s), $ = this.tf.multinomial(y.squeeze([1]), 1);
-        u = this.tf.gather(E.squeeze([1]), $, 1);
+        const { values: M, indices: x } = ot($, s), W = I(M.squeeze([1]), 1);
+        g = Y(x.squeeze([1]), W, 1);
       } else
-        u = this.tf.multinomial(k.squeeze([1]), 1);
-      let w;
-      return o?.includeProbabilities && (w = this.tf.softmax(k.squeeze([1]))), u = u.reshape([1, 1]), { output: u, attention: f?.squeeze([1]), probabilities: w };
+        g = I($.squeeze([1]), 1);
+      let v;
+      return o?.includeProbabilities && (v = Z($.squeeze([1]))), g = g.reshape([1, 1]), { output: g, attention: f?.squeeze([1]), probabilities: v };
     });
   }
   getNumParams() {
-    return F(this.config);
+    return K(this.config);
   }
   dispose() {
     this.wte.dispose(), this.wpe && this.wpe.dispose(), this.drop.dispose(), this.blocks.forEach((t) => t.dispose()), this.lnF.dispose();
   }
 }
 export {
-  A as default
+  kt as default
 };

package/dist/TeachableLLM.d.ts

@@ -1,20 +1,20 @@
-import { default as TF } from '@tensorflow/tfjs';
 import { GPTConfig } from './config';
 import { ITokeniser } from './tokeniser/type';
 import { default as NanoGPT } from './NanoGPTModel';
 import { SaveOptions } from './utilities/save';
 import { default as Generator, IGenerateOptions } from './Generator';
 import { default as Trainer, ITrainerOptions } from './Trainer';
-import { default as EE } from 'eventemitter3';
 import { default as MemoryProfiler } from './utilities/profile';
 type TeachableLLMStatus = 'warmup' | 'awaitingTokens' | 'ready' | 'training' | 'loading' | 'busy' | 'error';
-export default class TeachableLLM extends EE<'status' | 'error' | 'trainStep'> {
+export default class TeachableLLM {
+    private ee;
     private _config?;
     private _model?;
-    readonly tf: typeof TF;
     private _tokeniser?;
     private _status;
-    constructor(tf: typeof TF, tokeniser?: ITokeniser, model?: NanoGPT);
+    constructor(tokeniser?: ITokeniser, model?: NanoGPT);
+    get vocab(): string[];
+    get loaded(): boolean;
     get config(): GPTConfig;
     get model(): NanoGPT;
     get tokeniser(): ITokeniser;
@@ -22,16 +22,25 @@ export default class TeachableLLM extends EE<'status' | 'error' | 'trainStep'> {
     get ready(): boolean;
     private setStatus;
     saveModel(options?: SaveOptions): Promise<Blob>;
-    static loadModel(tf: typeof TF, data: Blob | Buffer | string): TeachableLLM;
-    static create(tf: typeof TF, config?: Partial<GPTConfig>): TeachableLLM;
+    static loadModel(data: Blob | Buffer | string): TeachableLLM;
+    static create(tokeniserType: 'char' | 'bpe', config?: Partial<GPTConfig>): TeachableLLM;
     getProfiler(): MemoryProfiler | undefined;
     get enableProfiler(): boolean;
     set enableProfiler(value: boolean);
     getNumParams(): number;
     trainer(): Trainer;
     train(text: string[], options?: ITrainerOptions): Promise<void>;
+    trainTokeniser(text: string[]): Promise<number>;
     generator(): Generator;
     generateText(prompt?: string, options?: IGenerateOptions): Promise<string>;
     dispose(): void;
+    on(event: 'status', listener: (status: TeachableLLMStatus) => void): void;
+    on(event: 'error', listener: (error: Error) => void): void;
+    on(event: 'trainStep', listener: (step: number) => void): void;
+    on(event: 'loaded', listener: () => void): void;
+    off(event: 'status', listener: (status: TeachableLLMStatus) => void): void;
+    off(event: 'error', listener: (error: Error) => void): void;
+    off(event: 'trainStep', listener: (step: number) => void): void;
+    off(event: 'loaded', listener: () => void): void;
 }
 export {};