@genai-fi/nanogpt 0.1.3 → 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/package.json

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