@genai-fi/nanogpt 0.1.2 → 0.1.4

package/dist/NanoGPTModel.d.ts

@@ -33,6 +33,7 @@ export default class NanoGPT {
     set trainable(value: boolean);
     private validateInput;
     private calculateLoss;
+    private computeAttentionRollout;
     forward(idx: TF.Tensor, targets?: TF.Tensor, training?: boolean, includeAttention?: boolean): {
         logits: TF.Tensor;
         loss?: TF.Tensor;

package/dist/NanoGPTModel.js

@@ -54,8 +54,8 @@ class $ {
   }
   inputPhase(t, e = !1) {
     return this.tf.tidy(() => {
-      const [, s] = t.shape, n = this.wte.embed(t), i = this.tf.range(0, s, 1, "int32"), o = this.wpe.apply(i), h = n.add(o);
-      return this.drop.apply(h, { training: e });
+      const [, s] = t.shape, i = this.wte.embed(t), n = this.tf.range(0, s, 1, "int32"), h = this.wpe.apply(n), o = i.add(h);
+      return this.drop.apply(o, { training: e });
     });
   }
   setSkipMask(t) {
@@ -90,44 +90,60 @@ class $ {
       throw console.error("Error computing loss:", s), new Error(`Loss computation failed: ${s}`);
     }
   }
-  forward(t, e, s = !1, n = !1) {
+  // Attention rollout per Abnar & Zuidema (2020)
+  // Expects list of (B, T, T) attention matrices already averaged over heads.
+  computeAttentionRollout(t) {
+    return this.tf.tidy(() => {
+      if (t.length === 0)
+        throw new Error("No attentions for rollout");
+      const e = t[0].shape[0], s = t[0].shape[1], i = this.tf.eye(s, s).expandDims(0);
+      let n = i.tile([e, 1, 1]);
+      for (const h of t) {
+        let o = h.add(i);
+        o = o.div(o.sum(-1, !0)), n = o.matMul(n);
+      }
+      return n;
+    });
+  }
+  forward(t, e, s = !1, i = !1) {
     return this.validateInput(t), this.tf.tidy(() => {
-      let i = this.inputPhase(t, s), o;
-      n && (o = this.tf.zeros([i.shape[0], i.shape[1], i.shape[1]]));
-      for (const l of this.blocks) {
-        const { output: r, attention: f } = l.call(i, s, n);
-        i = r, f && o && (o = o.add(f));
+      let n = this.inputPhase(t, s);
+      const h = [];
+      for (const c of this.blocks) {
+        const { output: p, attention: a } = c.call(n, s, i);
+        n = p, i && a && h.push(a);
       }
-      o && (o = o.div(this.blocks.length)), i = this.lnF.apply(i);
-      const h = this.wte.project(i);
-      let a;
-      return e && (a = this.calculateLoss(h, e)), { logits: h, loss: a, attention: n ? o : void 0 };
+      let o;
+      i && h.length > 0 && (o = this.computeAttentionRollout(h)), n = this.lnF.apply(n);
+      const l = this.wte.project(n);
+      let r;
+      return e && (r = this.calculateLoss(l, e)), { logits: l, loss: r, attention: i ? o : void 0 };
     });
   }
   generate(t, e) {
-    const s = e?.temperature ?? 1, n = e?.topK, i = e?.usePadding ?? !1, o = e?.includeAttention ?? !1;
+    const s = e?.temperature ?? 1, i = e?.topK, n = e?.usePadding ?? !1, h = e?.includeAttention ?? !1;
     return this.tf.tidy(() => {
-      const h = t, a = h.shape[1], l = a <= this.config.blockSize ? h : h.slice(
-        [0, a - this.config.blockSize],
-        [h.shape[0], this.config.blockSize]
-      ), r = i ? this.config.blockSize - l.shape[1] : 0, f = r > 0 ? this.tf.pad(l, [
+      const o = t, l = o.shape[1], r = l <= this.config.blockSize ? o : o.slice(
+        [0, l - this.config.blockSize],
+        [o.shape[0], this.config.blockSize]
+      ), c = n ? this.config.blockSize - r.shape[1] : 0, p = c > 0 ? this.tf.pad(r, [
         [0, 0],
-        [0, r]
-      ]) : l, { logits: g, attention: p } = this.forward(f, void 0, !1, o), d = g.shape[1] - 1 - r, m = g.slice([0, d, 0], [g.shape[0], 1, g.shape[2]]), u = p ? p.slice([0, d, 0], [p.shape[0], 1, p.shape[2]]) : void 0, b = m.div(s);
-      let c;
-      if (n) {
-        const { values: k, indices: w } = this.tf.topk(b, n), E = this.tf.multinomial(k.squeeze([1]), 1);
-        c = this.tf.gather(w.squeeze([1]), E, 1);
+        [0, c]
+      ]) : r, { logits: a, attention: g } = this.forward(p, void 0, !1, h), d = a.shape[1] - 1 - c, m = a.slice([0, d, 0], [a.shape[0], 1, a.shape[2]]), u = g ? g.slice([0, d, 0], [g.shape[0], 1, g.shape[2]]) : void 0, b = m.div(s);
+      let f;
+      if (i) {
+        const { values: k, indices: w } = this.tf.topk(b, i), E = this.tf.multinomial(k.squeeze([1]), 1);
+        f = this.tf.gather(w.squeeze([1]), E, 1);
       } else
-        c = this.tf.multinomial(b.squeeze([1]), 1);
-      return c = c.reshape([1, 1]), { output: c, attention: u?.squeeze([1]) };
+        f = this.tf.multinomial(b.squeeze([1]), 1);
+      return f = f.reshape([1, 1]), { output: f, attention: u?.squeeze([1]) };
     });
   }
   getNumParams() {
     const t = this.config.vocabSize * this.config.nEmbed + this.config.blockSize * this.config.nEmbed, e = this.config.nLayer * (4 * this.config.nEmbed * this.config.nEmbed + // qkv + proj
     2 * this.config.nEmbed), s = this.config.nLayer * (4 * this.config.nEmbed * this.config.nEmbed + // fc
-    this.config.nEmbed * 4 * this.config.nEmbed), n = this.config.nEmbed + this.config.vocabSize * this.config.nEmbed;
-    return t + e + s + n;
+    this.config.nEmbed * 4 * this.config.nEmbed), i = this.config.nEmbed + this.config.vocabSize * this.config.nEmbed;
+    return t + e + s + i;
   }
 }
 export {

package/dist/TeachableLLM.js

@@ -1,13 +1,13 @@
-import m from "./NanoGPTModel.js";
-import { defaultConfig as d } from "./config.js";
-import { saveModel as l } from "./utilities/save.js";
-import { loadModel as u } from "./utilities/load.js";
-import _ from "./Generator.js";
-import c from "./Trainer.js";
-import { E as f } from "./index-SOhdqzHq.js";
+import d from "./NanoGPTModel.js";
+import { defaultConfig as m } from "./config.js";
+import { saveModel as u } from "./utilities/save.js";
+import { loadModel as l } from "./utilities/load.js";
+import f from "./Generator.js";
+import _ from "./Trainer.js";
+import { E as c } from "./index-SOhdqzHq.js";
 import { dummyPassAsync as a } from "./utilities/dummy.js";
 import g from "./tokeniser/CharTokeniser.js";
-class n extends f {
+class n extends c {
   _config;
   _model;
   tf;
@@ -43,27 +43,27 @@ class n extends f {
   saveModel() {
     if (!this._model || !this._tokeniser)
       throw new Error("Model or tokeniser is not initialized.");
-    return l(this._model, this._tokeniser);
+    return u(this._model, this._tokeniser);
   }
   static loadModel(t, r) {
     const e = new n(t);
-    return u(t, r).then(({ model: i, tokeniser: s }) => {
-      e._model = i, e._tokeniser = s, e._config = i.config, e.setStatus("warmup"), a(i).then(() => {
+    return l(t, r).then(({ model: i, tokeniser: o }) => {
+      e._model = i, e._tokeniser = o, e._config = i.config, e.setStatus("warmup"), a(i).then(() => {
         e.setStatus("ready");
-      }).catch((o) => {
-        e.setStatus("error"), e.emit("error", o);
+      }).catch((s) => {
+        e.setStatus("error"), e.emit("error", s);
       });
     }).catch((i) => {
       e.setStatus("error"), e.emit("error", i);
     }), e;
   }
   static create(t, r = {}) {
-    const e = { ...d, ...r }, i = new g(e.vocabSize), s = new m(t, e), o = new n(t, i, s);
-    return o.setStatus("warmup"), a(s).then(() => {
-      o.setStatus("ready");
+    const e = { ...m, ...r }, i = new g(e.vocabSize), o = new d(t, e), s = new n(t, i, o);
+    return s.setStatus("warmup"), a(o).then(() => {
+      s.setStatus("ready");
     }).catch((h) => {
-      o.setStatus("error"), o.emit("error", h);
-    }), o;
+      s.setStatus("error"), s.emit("error", h);
+    }), s;
   }
   getNumParams() {
     if (!this._model)
@@ -73,7 +73,7 @@ class n extends f {
   trainer() {
     if (!this._model || !this._tokeniser)
       throw new Error("Model or tokeniser is not initialized.");
-    const t = new c(this._model, this._tokeniser);
+    const t = new _(this._model, this._tokeniser);
     return t.on("start", () => this.setStatus("training")), t.on("stop", () => this.setStatus("ready")), t.on("log", async (r) => {
       const e = this.listeners("trainStep");
       for (const i of e)
@@ -86,8 +86,12 @@ class n extends f {
   generator() {
     if (!this._model || !this._tokeniser)
       throw new Error("Model or tokeniser is not initialized.");
-    const t = new _(this._model, this._tokeniser);
-    return t.on("start", () => this.setStatus("busy")), t.on("stop", () => this.setStatus("ready")), t;
+    const t = new f(this._model, this._tokeniser);
+    return t.on("start", () => {
+      this.status === "ready" && this.setStatus("busy");
+    }), t.on("stop", () => {
+      this.status === "busy" && this.setStatus("ready");
+    }), t;
   }
   generateText(t, r) {
     return this.generator().generate(t, r);

package/dist/Trainer.js

@@ -1,6 +1,6 @@
 import { E as l } from "./index-SOhdqzHq.js";
 import o from "./training/FullTrainer.js";
-class m extends l {
+class d extends l {
   trainer;
   constructor(a, t) {
     super(), this.trainer = new o(a.tf, a, t, 1e-3);
@@ -8,27 +8,28 @@ class m extends l {
   stop() {
   }
   async train(a, t) {
-    const { trainDataset: r, validationDataset: e } = await this.trainer.createTrainValidationSplit(
+    const { trainDataset: e, validationDataset: r } = await this.trainer.createTrainValidationSplit(
       a,
       t?.batchSize || 32,
       t?.validationSplit || 0.1
     );
-    this.emit("start"), await this.trainer.trainOnDataset(
-      r,
+    this.trainer.setLearningRate(t?.learningRate || 1e-3), this.emit("start"), await this.trainer.trainOnDataset(
+      e,
       {
         prompt: t?.prompt,
         logInterval: t?.logInterval || 10,
         desiredLoss: t?.desiredLoss || 0.01,
+        maxSteps: t?.maxSteps || 1e3,
         onStep: async (i) => {
           const s = this.listeners("log");
           for (const n of s)
             await n(i);
         }
       },
-      e
+      r
     ), this.emit("stop");
   }
 }
 export {
-  m as default
+  d as default
 };

package/dist/training/FullTrainer.js

@@ -1,17 +1,18 @@
 import { generateText as L } from "../utilities/generate.js";
-import f from "./Trainer.js";
-const w = {
+import w from "./Trainer.js";
+const g = {
   desiredLoss: 0.01,
-  logInterval: 1
+  logInterval: 1,
+  maxSteps: 1e3
 };
-class g extends f {
+class S extends w {
   constructor(r, i, o, n = 3e-4) {
     super(r, i, o, n);
   }
   // Train for multiple epochs using Dataset API - FIXED memory leaks
   async trainOnDataset(r, i, o) {
-    const { desiredLoss: n, logInterval: h, onStep: l, prompt: c } = {
-      ...w,
+    const { desiredLoss: n, logInterval: c, onStep: l, prompt: p, maxSteps: d } = {
+      ...g,
       ...i
     }, s = {
       pass: 0,
@@ -24,19 +25,21 @@ class g extends f {
       validationLosses: []
     };
     this.dummyPass(), this.model.trainable = !0;
-    const d = Date.now(), m = await r.iterator();
+    const m = Date.now();
+    this.running = !0;
+    const u = await r.iterator();
     try {
-      for (; !(s.lastLoss < n); ) {
-        const e = await m.next();
+      for (; this.running && !(s.lastLoss < n); ) {
+        const e = await u.next();
         if (e.done) break;
-        const p = e.value, u = this.trainBatch(s, p), a = {
+        const h = e.value, f = this.trainBatch(s, h), a = {
           loss: s.lastLoss,
           step: s.step,
-          time: Date.now() - d,
-          batchSize: p.xs.shape[0]
+          time: Date.now() - m,
+          batchSize: h.xs.shape[0]
         };
-        if (this.model.log.push(a), s.step % h === 0) {
-          if (await u, o)
+        if (this.model.log.push(a), s.step % c === 0) {
+          if (await f, o)
             try {
               const t = await this.evaluateOnDataset(o, 5);
               s.validationLosses.push(t), a.valLoss = t;
@@ -44,8 +47,8 @@ class g extends f {
               console.error("Validation error:", t);
             }
           if (l) {
-            if (c) {
-              const t = await L(this.tokenizer, this.model, c, 100, {
+            if (p) {
+              const t = await L(this.tokenizer, this.model, p, 100, {
                 temperature: 0.8
               });
               a.example = t;
@@ -53,13 +56,14 @@ class g extends f {
             await l(a);
           }
         }
+        s.step >= d && this.stop();
       }
     } catch (e) {
       throw console.error("Training error:", e), this.tf.dispose(), e;
     }
-    return this.tf.dispose(), { losses: s.losses, validationLosses: s.validationLosses };
+    return this.tf.dispose(), this.running = !1, { losses: s.losses, validationLosses: s.validationLosses };
   }
 }
 export {
-  g as default
+  S as default
 };

package/dist/training/LayerTrainer.js

@@ -5,7 +5,8 @@ const w = {
   desiredLoss: 0.01,
   logInterval: 1,
   stepsPerLayer: 400,
-  maxPasses: 3
+  maxPasses: 3,
+  maxSteps: 1e3
 };
 class b extends S {
   trainingPattern = [];
@@ -37,20 +38,20 @@ class b extends S {
       validationLosses: []
     };
     this.dummyPass();
-    const f = Date.now();
+    const m = Date.now();
     this.startPass = 0, this.startLayer = 0;
-    const m = await r.iterator();
+    const f = await r.iterator();
     this.applyTrainingPattern(t.layerStep % this.trainingPattern.length);
     try {
       for (; !(t.lastLoss < p); ) {
-        const n = await m.next();
+        const n = await f.next();
         if (n.done) break;
         const y = n.value, P = this.trainBatch(t, y);
         t.stepSinceLayerChange++;
         const o = {
           loss: t.lastLoss,
           step: t.step,
-          time: Date.now() - f,
+          time: Date.now() - m,
           batchSize: y.xs.shape[0],
           pass: t.pass,
           layer: t.layerStep % this.model.config.nLayer

package/dist/training/Trainer.d.ts

@@ -20,6 +20,7 @@ export interface TrainingOptions {
     desiredLoss: number;
     logInterval: number;
     prompt?: string;
+    maxSteps: number;
     onStep?: (log: TrainingLogEntry) => Promise<void> | void;
 }
 export default abstract class GPTTrainer {
@@ -29,7 +30,10 @@ export default abstract class GPTTrainer {
     protected datasetBuilder: DatasetBuilder;
     protected tf: typeof TF;
     protected learningRate: number;
+    protected running: boolean;
     constructor(tf: typeof TF, model: NanoGPT, tokenizer: ITokeniser, learningRate?: number);
+    setLearningRate(learningRate: number): void;
+    stop(): void;
     getOptimizer(): AdamExt;
     resetOptimizer(config?: AdamConfig): void;
     private printGradients;

package/dist/training/Trainer.js

@@ -9,6 +9,13 @@ class y {
   datasetBuilder;
   tf;
   learningRate;
+  running = !1;
+  setLearningRate(t) {
+    this.learningRate = t, this.resetOptimizer({ learningRateFactor: 1, beta1: 0.9, beta2: 0.99, epsilon: 1e-8 });
+  }
+  stop() {
+    this.running = !1;
+  }
   getOptimizer() {
     return this.optimizer;
   }

package/package.json

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