@genai-fi/nanogpt 0.5.6 → 0.6.1

package/dist/Generator.js

@@ -1,12 +1,15 @@
 import { E as u } from "./index-Dwqa6Zy2.js";
-import "./index-Du-bmOP8.js";
+import "./index-CamYe_M8.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 "@tensorflow/tfjs";
+import "./random_width-CLMQG5Jn.js";
+import "./register_all_kernels-COt6wLD0.js";
+import "./index-Tf7vU29b.js";
+import "./dataset-BFFipD1c.js";
 import "./ops/cpu/rope.js";
 import "./ops/webgl/rope.js";
 import "./ops/grads/rope.js";
@@ -21,21 +24,19 @@ import "./ops/grads/matMulGelu.js";
 import "./ops/cpu/normRMS.js";
 import "./ops/webgl/normRMS.js";
 import "./ops/grads/normRMS.js";
-import "./random_width-C-v-35bY.js";
 import "./ops/cpu/gatherSub.js";
 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 "./dataset-iqT4Otvb.js";
-import "./index-Tf7vU29b.js";
 import "./papaparse.min-C8l2Kvo1.js";
 import "./ops/cpu/gelu.js";
 import "./ops/webgl/gelu.js";
 import "./ops/grads/gelu.js";
-import { t as d } from "./tensor2d-CRWjDyUe.js";
-import { c as g } from "./concat-DdKPyAtw.js";
+import "./ops/webgl/log.js";
+import { t as d } from "./tensor2d-jO8JY5Jd.js";
+import { c as g } from "./concat-XOK9ANZu.js";
 const k = [
   ...Array.from({ length: 95 }, (a, t) => String.fromCharCode(t + 32)),
   // ASCII
@@ -51,7 +52,7 @@ const k = [
 function T(a, t) {
   return a.length === t ? a : a.length > t ? a.slice(0, t) : a.concat(Array(t - a.length).fill(""));
 }
-class ot extends u {
+class rt extends u {
   constructor(t, o) {
     super(), this.model = t, this.tokeniser = o;
   }
@@ -123,5 +124,5 @@ class ot extends u {
   }
 }
 export {
-  ot as default
+  rt as default
 };

package/dist/NanoGPTModel.js

@@ -1,18 +1,19 @@
 import { defaultConfig as F } from "./config.js";
 import O from "./layers/TransformerBlock.js";
-import { T as N, r as R } from "./TiedEmbedding-8S8xn8e6.js";
-import A from "./layers/RoPECache.js";
-import G from "./layers/RMSNorm.js";
-import { estimateParameterCount as j } from "./utilities/parameters.js";
-import { createSoftmaxCrossEntropyWithGrad as B } from "./training/sparseCrossEntropy.js";
-import V from "./layers/BaseLayer.js";
-import { E as H, D as W, p as J } from "./random_width-C-v-35bY.js";
-import { o as x, j as y, u as Q, E as I, a9 as U, aa as X, ab as Y, t as z, a8 as Z, f as L, H as tt } from "./index-Du-bmOP8.js";
-import { r as T } from "./reshape-BH7eBpwq.js";
-import { r as et } from "./range-Bvs1hidm.js";
-import { s as q } from "./softmax-DhWoBa7r.js";
-import { t as ot } from "./ops-CJNniCAV.js";
-import { g as st } from "./gather-D6MsdXqc.js";
+import { T as _, r as D } from "./TiedEmbedding-C1HBot-5.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 G from "./layers/BaseLayer.js";
+import { E as B, D as V, p as j } from "./random_width-CLMQG5Jn.js";
+import { q as W, w as H, E as J, a6 as Q, t as z, a7 as U, f as v, n as X } from "./index-CamYe_M8.js";
+import { m as Y, t as Z } from "./register_all_kernels-COt6wLD0.js";
+import { r as L } from "./reshape-C45vIIRU.js";
+import { r as tt } from "./range-DqYjKnuG.js";
+import { s as M } from "./softmax-DvMvui-_.js";
+import { t as et } from "./ops-DqtYemmV.js";
+import { g as ot } from "./gather-DKtUaTtA.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -29,69 +30,17 @@ import { g as st } from "./gather-D6MsdXqc.js";
  * limitations under the License.
  * =============================================================================
  */
-function nt(l, t) {
-  let e = y(l, "a", "mod"), o = y(t, "b", "mod");
-  [e, o] = Q(e, o);
-  const n = { a: e, b: o };
-  return I.runKernel(U, n);
-}
-const it = /* @__PURE__ */ x({ mod_: nt });
-/**
- * @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 rt(l, t, e, o = !1) {
-  const n = y(l, "logits", "multinomial"), s = n.size, i = n.rank;
+function st(u, t, e, o = !1) {
+  const r = H(u, "logits", "multinomial"), s = r.size, n = r.rank;
   if (s < 2)
     throw new Error(`Error in multinomial: you need at least 2 outcomes, but got ${s}.`);
-  if (i > 2)
-    throw new Error(`Rank of probabilities must be 1 or 2, but is ${i}`);
+  if (n > 2)
+    throw new Error(`Rank of probabilities must be 1 or 2, but is ${n}`);
   e = e || Math.random();
-  const r = { logits: i === 1 ? T(n, [1, -1]) : n }, u = { numSamples: t, seed: e, normalized: o }, c = I.runKernel(X, r, u);
-  return i === 1 ? T(c, [c.size]) : c;
-}
-const C = /* @__PURE__ */ x({ multinomial_: rt });
-/**
- * @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 ct(l, t = 1, e = !0) {
-  const o = y(l, "x", "topk");
-  if (o.rank === 0)
-    throw new Error("topk() expects the input to be of rank 1 or higher");
-  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 > 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 }, [p, r] = I.runKernel(Y, s, i);
-  return { values: p, indices: r };
+  const i = { logits: n === 1 ? L(r, [1, -1]) : r }, h = { numSamples: t, seed: e, normalized: o }, c = J.runKernel(Q, i, h);
+  return n === 1 ? L(c, [c.size]) : c;
 }
-const at = /* @__PURE__ */ x({ topk_: ct });
+const C = /* @__PURE__ */ W({ multinomial_: st });
 /**
  * @license
  * Copyright 2018 Google LLC
@@ -101,13 +50,13 @@ const at = /* @__PURE__ */ x({ topk_: ct });
  * https://opensource.org/licenses/MIT.
  * =============================================================================
  */
-function lt(l) {
-  return new W(l);
+function nt(u) {
+  return new V(u);
 }
-function pt(l) {
-  return new H(l);
+function it(u) {
+  return new B(u);
 }
-class xt extends V {
+class St extends G {
   wte;
   // Token embeddings
   wpe;
@@ -121,15 +70,15 @@ class xt extends V {
   log = [];
   // Training log
   constructor(t = {}) {
-    super({ gpt: { ...F, ...t }, layerConfig: {} }), this.wte = new N(this.config, "token_embedding", this), this.config.gpt.useRope === !1 ? this.wpe = pt({
+    super({ gpt: { ...F, ...t }, layerConfig: {} }), this.wte = new _(this.config, "token_embedding", this), this.config.gpt.useRope === !1 ? this.wpe = it({
       inputDim: this.config.gpt.blockSize,
       outputDim: this.config.gpt.nEmbed,
       name: "positional_embedding",
-      embeddingsInitializer: R({ mean: 0, stddev: 0.02 })
-    }) : (this.ropeCache = new A(this.config.gpt), this.config.layerConfig.ropeCache = this.ropeCache), this.drop = lt({ rate: this.config.gpt.dropout }), this.blocks = [];
+      embeddingsInitializer: D({ mean: 0, stddev: 0.02 })
+    }) : (this.ropeCache = new K(this.config.gpt), this.config.layerConfig.ropeCache = this.ropeCache), this.drop = nt({ rate: this.config.gpt.dropout }), this.blocks = [];
     for (let e = 0; e < this.config.gpt.nLayer; e++)
       this.blocks.push(new O(e, this.config, this));
-    this.lnF = new G(this.config, "final_rms_norm", this);
+    this.lnF = new N(this.config, "final_rms_norm", this);
   }
   get checkpointing() {
     return this.config.layerConfig.checkpointing === !0;
@@ -139,12 +88,12 @@ class xt extends V {
   }
   inputPhase(t, e, o = !1) {
     return z(() => {
-      const n = this.wte.embed(t);
+      const r = this.wte.embed(t);
       if (this.config.gpt.useRope === !1) {
-        const [, s] = t.shape, i = this.config.gpt.blockSize, p = et(0, s, 1, "int32"), r = it(Z(p, L(e, "int32")), L(i, "int32")), u = this.wpe.apply(r), c = n.add(u);
+        const [, s] = t.shape, n = this.config.gpt.blockSize, p = tt(0, s, 1, "int32"), i = Y(U(p, v(e, "int32")), v(n, "int32")), h = this.wpe.apply(i), c = r.add(h);
         return this.drop.apply(c, { training: o });
       } else
-        return this.drop.apply(n, { training: o });
+        return this.drop.apply(r, { training: o });
     });
   }
   setSkipMask(t) {
@@ -169,7 +118,7 @@ class xt extends V {
   }
   calculateLoss(t, e) {
     try {
-      return B()(t, e).mean();
+      return A()(t, e).mean();
     } catch (o) {
       throw console.error("Error computing loss:", o), new Error(`Loss computation failed: ${o}`);
     }
@@ -209,35 +158,35 @@ class xt extends V {
   forward(t, e, o) {
     return this.validateInput(e), z(() => {
       this.startMemory();
-      const n = t.cache?.[0]?.length ?? 0;
-      let s = this.inputPhase(e, n, t.training);
+      const r = t.cache?.[0]?.length ?? 0;
+      let s = this.inputPhase(e, r, 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 r = 0; r < this.blocks.length; r++) {
-        const u = this.blocks[r], c = Math.random() * 1e9, g = {
+      for (let i = 0; i < this.blocks.length; i++) {
+        const h = this.blocks[i], c = Math.random() * 1e9, g = {
           training: t.training,
           seed: c,
           attentionScores: t.attentionScores,
-          pastKV: t.cache ? t.cache[r] : void 0
-        }, S = this.config.layerConfig.checkpointing && t.training ? u.callCheckpoint(g, s) : u.call(g, s);
-        s.dispose(), s = S;
+          pastKV: t.cache ? t.cache[i] : void 0
+        }, E = this.config.layerConfig.checkpointing && t.training ? h.callCheckpoint(g, s) : h.call(g, s);
+        s.dispose(), s = E;
       }
       s = this.lnF.call(t, s);
-      const i = this.wte.project(s);
+      const n = this.wte.project(s);
       s.dispose();
       let p;
-      return o && (p = this.calculateLoss(i, o)), this.endMemory("Forward"), p ? [i, p] : [i];
+      return o && (p = this.calculateLoss(n, o)), this.endMemory("Forward"), p ? [n, p] : [n];
     });
   }
   generate(t, e, o) {
-    const n = o?.temperature ?? 1, s = o?.topK, i = o?.topP, p = o?.usePadding ?? !1;
+    const r = o?.temperature ?? 1, s = o?.topK, n = o?.topP, p = o?.usePadding ?? !1;
     return z(() => {
-      const r = t, u = r.shape[1], c = u <= this.config.gpt.blockSize ? r : r.slice(
-        [0, u - this.config.gpt.blockSize],
-        [r.shape[0], this.config.gpt.blockSize]
-      ), g = p ? this.config.gpt.blockSize - c.shape[1] : 0, S = g > 0 ? J(c, [
+      const i = t, h = i.shape[1], c = h <= this.config.gpt.blockSize ? i : i.slice(
+        [0, h - this.config.gpt.blockSize],
+        [i.shape[0], this.config.gpt.blockSize]
+      ), g = p ? this.config.gpt.blockSize - c.shape[1] : 0, E = g > 0 ? j(c, [
         [0, 0],
         [0, g]
       ]) : c, f = {
@@ -246,41 +195,41 @@ class xt extends V {
           attentionOut: []
         } : void 0,
         cache: e
-      }, [d] = this.forward(f, S), M = d.shape[1] - 1 - g, K = d.slice([0, M, 0], [d.shape[0], 1, d.shape[2]]);
-      f.attentionScores?.attentionOut && f.attentionScores.attentionOut.forEach((h, b) => {
-        h.shape[1] !== 1 && (f.attentionScores.attentionOut[b] = tt(
-          h.slice([0, M, 0], [h.shape[0], 1, h.shape[2]])
-        ), h.dispose());
+      }, [d] = this.forward(f, E), $ = d.shape[1] - 1 - g, q = d.slice([0, $, 0], [d.shape[0], 1, d.shape[2]]);
+      f.attentionScores?.attentionOut && f.attentionScores.attentionOut.forEach((l, b) => {
+        l.shape[1] !== 1 && (f.attentionScores.attentionOut[b] = X(
+          l.slice([0, $, 0], [l.shape[0], 1, l.shape[2]])
+        ), l.dispose());
       }), d.dispose();
-      const w = K.div(n);
+      const w = q.div(r);
       let m;
-      if (i) {
-        const h = q(w.squeeze([1])), b = h.arraySync()[0];
-        h.dispose();
-        const E = b.map((a, k) => ({ prob: a, index: k })).sort((a, k) => k.prob - a.prob);
-        let v = 0;
-        const $ = new Array(E.length).fill(0);
-        for (const a of E)
-          if (v += a.prob, $[a.index] = a.prob, v >= i)
+      if (n) {
+        const l = M(w.squeeze([1])), b = l.arraySync()[0];
+        l.dispose();
+        const y = b.map((a, k) => ({ prob: a, index: k })).sort((a, k) => k.prob - a.prob);
+        let P = 0;
+        const S = new Array(y.length).fill(0);
+        for (const a of y)
+          if (P += a.prob, S[a.index] = a.prob, P >= n)
             break;
-        const _ = $.reduce((a, k) => a + k, 0), D = $.map((a) => a / _);
-        m = C(ot(D), 1, void 0, !0);
+        const x = S.reduce((a, k) => a + k, 0), T = S.map((a) => a / x);
+        m = C(et(T), 1, void 0, !0);
       } else if (s) {
-        const { values: h, indices: b } = at(w, s), E = C(h.squeeze([1]), 1);
-        m = st(b.squeeze([1]), E, 1);
+        const { values: l, indices: b } = Z(w, s), y = C(l.squeeze([1]), 1);
+        m = ot(b.squeeze([1]), y, 1);
       } else
         m = C(w.squeeze([1]), 1);
-      let P;
-      return o?.includeProbabilities && (P = q(w.squeeze([1]))), m = m.reshape([1, 1]), { output: m, probabilities: P, attention: f.attentionScores?.attentionOut };
+      let I;
+      return o?.includeProbabilities && (I = M(w.squeeze([1]))), m = m.reshape([1, 1]), { output: m, probabilities: I, attention: f.attentionScores?.attentionOut };
     });
   }
   getNumParams() {
-    return j(this.config.gpt);
+    return R(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 {
-  xt as default
+  St as default
 };