@genai-fi/nanogpt 0.4.5 → 0.5.1

package/dist/NanoGPTModel.js

@@ -1,17 +1,16 @@
-import { defaultConfig as x } from "./config.js";
-import W from "./layers/TransformerBlock.js";
-import { E as F, D as P, T as q, r as T, p as D } from "./TiedEmbedding-DznFwzcB.js";
-import K from "./layers/RoPECache.js";
-import N from "./layers/RMSNorm.js";
-import { estimateParameterCount as R } from "./utilities/parameters.js";
-import { createSoftmaxCrossEntropyWithGrad as A } from "./training/sparseCrossEntropy.js";
-import B from "./layers/BaseLayer.js";
-import { o as $, h as E, p as G, E as v, a9 as O, aa as j, ab as Q, t as w, a8 as V, f as C } from "./index--6vO-cOz.js";
-import { r as _ } from "./reshape-z51Eu-re.js";
-import { r as X } from "./range-C_vpUjBu.js";
-import { e as H } from "./tfjs_backend-DuKis_xG.js";
-import { g as J } from "./gather-C5D8PxwA.js";
-import { s as U } from "./softmax-Dsxflvdl.js";
+import { defaultConfig as L } from "./config.js";
+import v from "./layers/TransformerBlock.js";
+import { E as T, D as q, T as K, r as P, p as _ } from "./TiedEmbedding-DsDRvLB0.js";
+import F from "./layers/RoPECache.js";
+import D from "./layers/RMSNorm.js";
+import { estimateParameterCount as O } from "./utilities/parameters.js";
+import { createSoftmaxCrossEntropyWithGrad as N } from "./training/sparseCrossEntropy.js";
+import { B as R } from "./BaseLayer-BhrMN8JO.js";
+import { o as E, i as d, q as B, E as y, aa as G, ab as V, ac as j, t as w, a9 as A, f as z, F as W } from "./index-iNhkcAEQ.js";
+import { r as C } from "./reshape-DxTPgnwL.js";
+import { r as H } from "./range-BsFU-SNG.js";
+import { g as J } from "./gather-Bxe1Qip8.js";
+import { s as Q } from "./softmax-BjsptB07.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -28,13 +27,13 @@ import { s as U } from "./softmax-Dsxflvdl.js";
  * limitations under the License.
  * =============================================================================
  */
-function Y(c, t) {
-  let e = E(c, "a", "mod"), o = E(t, "b", "mod");
-  [e, o] = G(e, o);
-  const i = { a: e, b: o };
-  return v.runKernel(O, i);
+function U(h, t) {
+  let e = d(h, "a", "mod"), o = d(t, "b", "mod");
+  [e, o] = B(e, o);
+  const n = { a: e, b: o };
+  return y.runKernel(G, n);
 }
-const Z = /* @__PURE__ */ $({ mod_: Y });
+const X = /* @__PURE__ */ E({ mod_: U });
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -51,17 +50,17 @@ const Z = /* @__PURE__ */ $({ mod_: Y });
  * limitations under the License.
  * =============================================================================
  */
-function tt(c, t, e, o = !1) {
-  const i = E(c, "logits", "multinomial"), s = i.size, l = i.rank;
+function Y(h, t, e, o = !1) {
+  const n = d(h, "logits", "multinomial"), s = n.size, i = n.rank;
   if (s < 2)
     throw new Error(`Error in multinomial: you need at least 2 outcomes, but got ${s}.`);
-  if (l > 2)
-    throw new Error(`Rank of probabilities must be 1 or 2, but is ${l}`);
+  if (i > 2)
+    throw new Error(`Rank of probabilities must be 1 or 2, but is ${i}`);
   e = e || Math.random();
-  const n = { logits: l === 1 ? _(i, [1, -1]) : i }, h = { numSamples: t, seed: e, normalized: o }, a = v.runKernel(j, n, h);
-  return l === 1 ? _(a, [a.size]) : a;
+  const c = { logits: i === 1 ? C(n, [1, -1]) : n }, l = { numSamples: t, seed: e, normalized: o }, a = y.runKernel(V, c, l);
+  return i === 1 ? C(a, [a.size]) : a;
 }
-const M = /* @__PURE__ */ $({ multinomial_: tt });
+const I = /* @__PURE__ */ E({ multinomial_: Y });
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -78,19 +77,19 @@ const M = /* @__PURE__ */ $({ multinomial_: tt });
  * limitations under the License.
  * =============================================================================
  */
-function et(c, t = 1, e = !0) {
-  const o = E(c, "x", "topk");
+function Z(h, t = 1, e = !0) {
+  const o = d(h, "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];
+  const n = 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 }, l = { k: t, sorted: e }, [r, n] = v.runKernel(Q, s, l);
-  return { values: r, indices: n };
+  if (t > n)
+    throw new Error(`'k' passed to topk() must be <= the last dimension (${n}) but got ${t}`);
+  const s = { x: o }, i = { k: t, sorted: e }, [r, c] = y.runKernel(j, s, i);
+  return { values: r, indices: c };
 }
-const ot = /* @__PURE__ */ $({ topk_: et });
+const tt = /* @__PURE__ */ E({ topk_: Z });
 /**
  * @license
  * Copyright 2018 Google LLC
@@ -100,13 +99,13 @@ const ot = /* @__PURE__ */ $({ topk_: et });
  * https://opensource.org/licenses/MIT.
  * =============================================================================
  */
-function st(c) {
-  return new P(c);
+function et(h) {
+  return new q(h);
 }
-function it(c) {
-  return new F(c);
+function ot(h) {
+  return new T(h);
 }
-class wt extends B {
+class dt extends R {
   wte;
   // Token embeddings
   wpe;
@@ -120,55 +119,30 @@ class wt extends B {
   log = [];
   // Training log
   constructor(t = {}) {
-    super({ gpt: { ...x, ...t }, layerConfig: {} }), this.wte = new q({
-      vocabSize: this.config.gpt.vocabSize,
-      embedDim: this.config.gpt.nEmbed,
-      name: "token_embedding"
-    }), this.config.gpt.useRope === !1 ? this.wpe = it({
+    super({ gpt: { ...L, ...t }, layerConfig: {} }), this.wte = new K(this.config, "token_embedding", this), this.config.gpt.useRope === !1 ? this.wpe = ot({
       inputDim: this.config.gpt.blockSize,
       outputDim: this.config.gpt.nEmbed,
       name: "positional_embedding",
-      embeddingsInitializer: T({ mean: 0, stddev: 0.02 })
-    }) : (this.ropeCache = new K(this.config.gpt), this.config.layerConfig.ropeCache = this.ropeCache), this.drop = st({ rate: this.config.gpt.dropout }), this.blocks = [];
+      embeddingsInitializer: P({ mean: 0, stddev: 0.02 })
+    }) : (this.ropeCache = new F(this.config.gpt), this.config.layerConfig.ropeCache = this.ropeCache), this.drop = et({ rate: this.config.gpt.dropout }), this.blocks = [];
     for (let e = 0; e < this.config.gpt.nLayer; e++)
-      this.blocks.push(new W(e, this.config));
-    this.lnF = new N(this.config, "final_rms_norm");
+      this.blocks.push(new v(e, this.config, this));
+    this.lnF = new D(this.config, "final_rms_norm", this);
   }
   get checkpointing() {
-    return this.config.layerConfig.checkpointAttention === !0 || this.config.layerConfig.checkpointMLP === !0;
+    return this.config.layerConfig.checkpointing === !0;
   }
   set checkpointing(t) {
-    this.config.layerConfig.checkpointAttention = t, this.config.layerConfig.checkpointMLP = t;
-  }
-  get variables() {
-    return [
-      //...this.wpe.trainableWeights.map((v) => v.read() as TF.Variable),
-      ...this.blocks.flatMap((t) => t.variables),
-      ...this.lnF.trainableWeights.map((t) => t),
-      ...this.wte.variables
-    ];
-  }
-  saveWeights() {
-    const t = /* @__PURE__ */ new Map();
-    t.set("token_embedding", this.wte.getWeights()), this.wpe && t.set("positional_embedding", this.wpe.getWeights());
-    for (let e = 0; e < this.blocks.length; e++)
-      this.blocks[e].saveWeights(t);
-    return t.set("final_rms_norm", this.lnF.getWeights()), t;
-  }
-  loadWeights(t) {
-    this.wte.setWeights(t.get("token_embedding") || []), this.wpe && this.wpe.setWeights(t.get("positional_embedding") || []);
-    for (let e = 0; e < this.blocks.length; e++)
-      this.blocks[e].loadWeights(t);
-    this.lnF.setWeights(t.get("final_rms_norm") || []);
+    this.config.layerConfig.checkpointing = t;
   }
   inputPhase(t, e, o = !1) {
     return w(() => {
-      const i = this.wte.embed(t);
+      const n = this.wte.embed(t);
       if (this.config.gpt.useRope === !1) {
-        const [, s] = t.shape, l = this.config.gpt.blockSize, r = X(0, s, 1, "int32"), n = Z(V(r, C(e, "int32")), C(l, "int32")), h = this.wpe.apply(n), a = i.add(h);
+        const [, s] = t.shape, i = this.config.gpt.blockSize, r = H(0, s, 1, "int32"), c = X(A(r, z(e, "int32")), z(i, "int32")), l = this.wpe.apply(c), a = n.add(l);
         return this.drop.apply(a, { training: o });
       } else
-        return this.drop.apply(i, { training: o });
+        return this.drop.apply(n, { training: o });
     });
   }
   setSkipMask(t) {
@@ -183,11 +157,6 @@ class wt extends B {
     for (let e = 0; e < this.blocks.length; e++)
       this.blocks[e].trainable = t[e];
   }
-  set trainable(t) {
-    for (const e of this.blocks)
-      e.trainable = t;
-    this.lnF.trainable = t;
-  }
   validateInput(t) {
     if (t.shape.length !== 2)
       throw new Error(`Invalid input shape: expected [batch_size, sequence_length], got ${t.shape}`);
@@ -198,84 +167,107 @@ class wt extends B {
   }
   calculateLoss(t, e) {
     try {
-      return A()(t, e).mean();
+      return N()(t, e).mean();
     } catch (o) {
       throw console.error("Error computing loss:", o), new Error(`Loss computation failed: ${o}`);
     }
   }
   // Attention rollout per Abnar & Zuidema (2020)
   // Expects list of (B, T, T) attention matrices already averaged over heads.
-  computeAttentionRollout(t) {
-    return w(() => {
-      if (t.length === 0)
-        throw new Error("No attentions for rollout");
-      const [e, o, i] = t[0].shape;
-      for (const n of t) {
-        const [h, a, p] = n.shape;
-        if (h !== e || a !== o || p !== i)
-          throw new Error(
-            `Inconsistent attention shapes in rollout: expected [${e},${o},${i}] got [${h},${a},${p}]`
-          );
-      }
-      const s = t.map((n) => n.slice([0, 0, 0], [e, o, o])), l = H(o, o).expandDims(0);
-      let r = l.tile([e, 1, 1]);
-      for (const n of s) {
-        const h = n.add(l);
-        r = h.div(h.sum(-1, !0)).matMul(r);
-      }
-      return r;
-    });
-  }
-  forward(t, e, o = !1, i = !1, s) {
-    return this.validateInput(t), w(() => {
+  /*private computeAttentionRollout(attentions: Tensor[]): Tensor {
+          return tidy(() => {
+              if (attentions.length === 0) {
+                  throw new Error('No attentions for rollout');
+              }
+              const [B, Q, K] = attentions[0].shape as number[];
+  
+              // Validate shapes are consistent
+              for (const a of attentions) {
+                  const [b2, q2, k2] = a.shape as number[];
+                  if (b2 !== B || q2 !== Q || k2 !== K) {
+                      throw new Error(
+                          `Inconsistent attention shapes in rollout: expected [${B},${Q},${K}] got [${b2},${q2},${k2}]`
+                      );
+                  }
+              }
+  
+              // Always slice to [B, Q, Q] for rollout
+              const attentionsSliced = attentions.map((att) => att.slice([0, 0, 0], [B, Q, Q]));
+  
+              const ey = eye(Q, Q).expandDims(0); // (1,Q,Q)
+              let rollout = ey.tile([B, 1, 1]); // (B,Q,Q)
+              for (const att of attentionsSliced) {
+                  const a = att.add(ey);
+                  const aNorm = a.div(a.sum(-1, true)); // (B,Q,Q)
+                  rollout = aNorm.matMul(rollout); // (B,Q,Q)
+              }
+              return rollout;
+          });
+      }*/
+  forward(t, e, o) {
+    return this.validateInput(e), w(() => {
       this.startMemory();
-      const l = s?.[0]?.length ?? 0;
-      let r = this.inputPhase(t, l, 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 g = 0; g < this.blocks.length; g++) {
-        const u = r, m = this.blocks[g], {
-          output: b,
-          attention: k,
-          cache: f
-        } = m.call(r, o, i, s ? s[g] : void 0);
-        r = b, u.dispose(), i && k && n.push(k), s && f ? (s[g]?.k.dispose(), s[g]?.v.dispose(), s[g] = f) : f && (f.k.dispose(), f.v.dispose());
+      const n = t.cache?.[0]?.length ?? 0;
+      let s = this.inputPhase(e, n, t.training);
+      if (t.cache && t.cache.length !== this.blocks.length)
+        throw console.error("Cache", t.cache), new Error(
+          `Cache length ${t.cache.length} does not match number of blocks ${this.blocks.length}`
+        );
+      for (let c = 0; c < this.blocks.length; c++) {
+        const l = this.blocks[c], a = Math.random() * 1e9, u = {
+          training: t.training,
+          seed: a,
+          attentionScores: t.attentionScores,
+          pastKV: t.cache ? t.cache[c] : void 0
+        }, p = this.config.layerConfig.checkpointing && t.training ? l.callCheckpoint(u, s) : l.call(u, s);
+        s.dispose(), s = p;
       }
-      let h;
-      i && n.length > 0 && (h = this.computeAttentionRollout(n)), r = this.lnF.apply(r);
-      const a = this.wte.project(r);
-      let p;
-      return e && (p = this.calculateLoss(a, e)), this.endMemory("Forward"), { logits: a, loss: p, attention: i ? h : void 0 };
+      s = this.lnF.call(t, s);
+      const i = this.wte.project(s);
+      s.dispose();
+      let r;
+      return o && (r = this.calculateLoss(i, o)), this.endMemory("Forward"), r ? [i, r] : [i];
     });
   }
   generate(t, e, o) {
-    const i = o?.temperature ?? 1, s = o?.topK, l = o?.usePadding ?? !1, r = o?.includeAttention ?? !1;
+    const n = o?.temperature ?? 1, s = o?.topK, i = o?.usePadding ?? !1;
     return w(() => {
-      const n = t, h = n.shape[1], a = h <= this.config.gpt.blockSize ? n : n.slice(
-        [0, h - this.config.gpt.blockSize],
-        [n.shape[0], this.config.gpt.blockSize]
-      ), p = l ? this.config.gpt.blockSize - a.shape[1] : 0, g = p > 0 ? D(a, [
+      const r = t, c = r.shape[1], l = c <= this.config.gpt.blockSize ? r : r.slice(
+        [0, c - this.config.gpt.blockSize],
+        [r.shape[0], this.config.gpt.blockSize]
+      ), a = i ? this.config.gpt.blockSize - l.shape[1] : 0, u = a > 0 ? _(l, [
         [0, 0],
-        [0, p]
-      ]) : a, { logits: u, attention: m } = this.forward(g, void 0, !1, r, e), b = u.shape[1] - 1 - p, k = u.slice([0, b, 0], [u.shape[0], 1, u.shape[2]]), f = m ? m.slice([0, b, 0], [m.shape[0], 1, m.shape[2]]) : void 0, y = k.div(i);
-      let d;
+        [0, a]
+      ]) : l, p = {
+        training: !1,
+        attentionScores: o?.attentionScores ? {
+          attentionOut: []
+        } : void 0,
+        cache: e
+      }, [f] = this.forward(p, u), S = f.shape[1] - 1 - a, M = f.slice([0, S, 0], [f.shape[0], 1, f.shape[2]]);
+      p.attentionScores?.attentionOut && p.attentionScores.attentionOut.forEach((g, k) => {
+        g.shape[1] !== 1 && (p.attentionScores.attentionOut[k] = W(
+          g.slice([0, S, 0], [g.shape[0], 1, g.shape[2]])
+        ), g.dispose());
+      }), f.dispose();
+      const b = M.div(n);
+      let m;
       if (s) {
-        const { values: S, indices: I } = ot(y, s), L = M(S.squeeze([1]), 1);
-        d = J(I.squeeze([1]), L, 1);
+        const { values: g, indices: k } = tt(b, s), x = I(g.squeeze([1]), 1);
+        m = J(k.squeeze([1]), x, 1);
       } else
-        d = M(y.squeeze([1]), 1);
-      let z;
-      return o?.includeProbabilities && (z = U(y.squeeze([1]))), d = d.reshape([1, 1]), { output: d, attention: f?.squeeze([1]), probabilities: z };
+        m = I(b.squeeze([1]), 1);
+      let $;
+      return o?.includeProbabilities && ($ = Q(b.squeeze([1]))), m = m.reshape([1, 1]), { output: m, probabilities: $, attention: p.attentionScores?.attentionOut };
     });
   }
   getNumParams() {
-    return R(this.config.gpt);
+    return O(this.config.gpt);
   }
   dispose() {
     this.wte.dispose(), this.wpe && this.wpe.dispose(), this.drop.dispose(), this.blocks.forEach((t) => t.dispose()), this.lnF.dispose();
   }
 }
 export {
-  wt as default
+  dt as default
 };

package/dist/{Reshape-CiAY8ltP.js → Reshape-BE5rA4rT.js}

@@ -1,5 +1,5 @@
-import { ac as f, ad as g, n as p, ae as C, j as x } from "./index--6vO-cOz.js";
-import { u as I } from "./gpgpu_math-CUzjlO9A.js";
+import { ad as $, ae as g, p, af as C, k as x } from "./index-iNhkcAEQ.js";
+import { u as I } from "./gpgpu_math-C0zyxKFi.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -17,7 +17,7 @@ import { u as I } from "./gpgpu_math-CUzjlO9A.js";
  * =============================================================================
  */
 function R(t, e, o = "index") {
-  const s = f(e);
+  const s = $(e);
   return s.map((n, r) => {
     const i = `int ${t[r]} = ${o} / ${n}`, u = r === s.length - 1 ? `int ${t[r + 1]} = ${o} - ${t[r]} * ${n}` : `index -= ${t[r]} * ${n}`;
     return `${i}; ${u};`;
@@ -38,7 +38,7 @@ function S(t, e, o = "index") {
   }).join("");
 }
 function F(t) {
-  const e = f(t).map((o) => o.toString());
+  const e = $(t).map((o) => o.toString());
   return `
   int getFlatIndex(ivec3 coords) {
     return coords.x * ${e[0]} + coords.y * ${e[1]} + coords.z;
@@ -82,7 +82,7 @@ function m(t) {
 function d(t) {
   return t % 2 === 0;
 }
-function $(t, e) {
+function f(t, e) {
   if (t = t.slice(-2), e = e.slice(-2), g(t, e) || !t.length || !e.length || t[0] === 0 || t[1] === 0 || e[0] === 0 || e[1] === 0)
     return !0;
   if (t.length !== e.length) {
@@ -201,12 +201,12 @@ function D(t, e, o) {
  * limitations under the License.
  * =============================================================================
  */
-function k(t) {
+function U(t) {
   const { inputs: e, backend: o, attrs: s } = t, { x: n } = e, { shape: r } = s, i = o, u = p(n.shape), a = C(r, u), c = p(a);
   x(u === c, () => `The new shape (${a}) has ${c} elements and the old shape (${n.shape}) has ${u} elements. The new shape and old shape must have the same number of elements.`);
   const l = i.texData.get(n.dataId);
-  return l.isPacked && !$(n.shape, a) && !(l.texture !== null && $(l.shape, a)) ? D(n, a, i) : (i.incRef(n.dataId), { dataId: n.dataId, shape: a, dtype: n.dtype });
+  return l.isPacked && !f(n.shape, a) && !(l.texture !== null && f(l.shape, a)) ? D(n, a, i) : (i.incRef(n.dataId), { dataId: n.dataId, shape: a, dtype: n.dtype });
 }
 export {
-  k as r
+  U as r
 };

package/dist/TeachableLLM.js

@@ -11,7 +11,7 @@ import g from "./tokeniser/bpe.js";
 import "./papaparse.min-C8l2Kvo1.js";
 import "./index-Tf7vU29b.js";
 import "./jszip.min-CjP2V1VV.js";
-import "./index--6vO-cOz.js";
+import "./index-iNhkcAEQ.js";
 import "./ops/cpu/scatterSub.js";
 import "./ops/webgl/scatterSub.js";
 import "./ops/cpu/gatherSub.js";