@genai-fi/nanogpt 0.4.2 → 0.4.4

package/dist/Generator.js

@@ -1,7 +1,7 @@
 import { E as u } from "./index-Dwqa6Zy2.js";
-import "./index-C4JCoBvj.js";
-import { t as d } from "./tensor2d-DqFGNs_K.js";
-import { c as p } from "./concat-CuRsVY-K.js";
+import "./index--6vO-cOz.js";
+import { t as d } from "./tensor2d-DUr_htjt.js";
+import { c as p } from "./concat-DvWM7HGZ.js";
 class w extends u {
   constructor(s, e) {
     super(), this.model = s, this.tokeniser = e;

package/dist/NanoGPTModel.js

@@ -1,17 +1,17 @@
 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 K, p as T } from "./TiedEmbedding-CnJ1bx4q.js";
-import D from "./layers/RoPECache.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 y, h as E, p as G, E as v, a6 as O, a7 as Q, a8 as j, t as w, a5 as U, f as C } from "./index-C4JCoBvj.js";
-import { r as _ } from "./reshape-Boe4DuIO.js";
-import { r as V } from "./range-9AzeApCc.js";
-import { e as X } from "./tfjs_backend-Cug-PH75.js";
-import { g as H } from "./gather-ZYRWhmXR.js";
-import { s as J } from "./softmax-Cujsg4ay.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";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -28,13 +28,13 @@ import { s as J } from "./softmax-Cujsg4ay.js";
  * limitations under the License.
  * =============================================================================
  */
-function Y(h, t) {
-  let e = E(h, "a", "mod"), o = E(t, "b", "mod");
+function Y(c, t) {
+  let e = E(c, "a", "mod"), o = E(t, "b", "mod");
   [e, o] = G(e, o);
-  const n = { a: e, b: o };
-  return v.runKernel(O, n);
+  const i = { a: e, b: o };
+  return v.runKernel(O, i);
 }
-const Z = /* @__PURE__ */ y({ mod_: Y });
+const Z = /* @__PURE__ */ $({ mod_: Y });
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -51,17 +51,17 @@ const Z = /* @__PURE__ */ y({ mod_: Y });
  * limitations under the License.
  * =============================================================================
  */
-function tt(h, t, e, o = !1) {
-  const n = E(h, "logits", "multinomial"), s = n.size, r = n.rank;
+function tt(c, t, e, o = !1) {
+  const i = E(c, "logits", "multinomial"), s = i.size, l = 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}`);
+  if (l > 2)
+    throw new Error(`Rank of probabilities must be 1 or 2, but is ${l}`);
   e = e || Math.random();
-  const i = { logits: r === 1 ? _(n, [1, -1]) : n }, p = { numSamples: t, seed: e, normalized: o }, l = v.runKernel(Q, i, p);
-  return r === 1 ? _(l, [l.size]) : l;
+  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 M = /* @__PURE__ */ y({ multinomial_: tt });
+const M = /* @__PURE__ */ $({ multinomial_: tt });
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -78,19 +78,19 @@ const M = /* @__PURE__ */ y({ multinomial_: tt });
  * limitations under the License.
  * =============================================================================
  */
-function et(h, t = 1, e = !0) {
-  const o = E(h, "x", "topk");
+function et(c, t = 1, e = !0) {
+  const o = E(c, "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];
+  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 > n)
-    throw new Error(`'k' passed to topk() must be <= the last dimension (${n}) but got ${t}`);
-  const s = { x: o }, r = { k: t, sorted: e }, [a, i] = v.runKernel(j, s, r);
-  return { values: a, indices: i };
+  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 };
 }
-const ot = /* @__PURE__ */ y({ topk_: et });
+const ot = /* @__PURE__ */ $({ topk_: et });
 /**
  * @license
  * Copyright 2018 Google LLC
@@ -100,11 +100,11 @@ const ot = /* @__PURE__ */ y({ topk_: et });
  * https://opensource.org/licenses/MIT.
  * =============================================================================
  */
-function st(h) {
-  return new P(h);
+function st(c) {
+  return new P(c);
 }
-function nt(h) {
-  return new F(h);
+function it(c) {
+  return new F(c);
 }
 class wt extends B {
   wte;
@@ -124,12 +124,12 @@ class wt extends B {
       vocabSize: this.config.gpt.vocabSize,
       embedDim: this.config.gpt.nEmbed,
       name: "token_embedding"
-    }), this.config.gpt.useRope === !1 ? this.wpe = nt({
+    }), this.config.gpt.useRope === !1 ? this.wpe = it({
       inputDim: this.config.gpt.blockSize,
       outputDim: this.config.gpt.nEmbed,
       name: "positional_embedding",
-      embeddingsInitializer: K({ mean: 0, stddev: 0.02 })
-    }) : (this.ropeCache = new D(this.config.gpt), this.config.layerConfig.ropeCache = this.ropeCache), this.drop = st({ rate: this.config.gpt.dropout }), this.blocks = [];
+      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 = [];
     for (let e = 0; e < this.config.gpt.nLayer; e++)
       this.blocks.push(new W(e, this.config));
     this.lnF = new N(this.config, 1e-8, "final_rms_norm");
@@ -163,12 +163,12 @@ class wt extends B {
   }
   inputPhase(t, e, o = !1) {
     return w(() => {
-      const n = this.wte.embed(t);
+      const i = this.wte.embed(t);
       if (this.config.gpt.useRope === !1) {
-        const [, s] = t.shape, r = this.config.gpt.blockSize, a = V(0, s, 1, "int32"), i = Z(U(a, C(e, "int32")), C(r, "int32")), p = this.wpe.apply(i), l = n.add(p);
-        return this.drop.apply(l, { training: o });
+        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);
+        return this.drop.apply(a, { training: o });
       } else
-        return this.drop.apply(n, { training: o });
+        return this.drop.apply(i, { training: o });
     });
   }
   setSkipMask(t) {
@@ -209,67 +209,64 @@ class wt extends B {
     return w(() => {
       if (t.length === 0)
         throw new Error("No attentions for rollout");
-      const [e, o, n] = t[0].shape;
-      for (const s of t) {
-        const [r, a, i] = s.shape;
-        if (r !== e || a !== o || i !== n)
+      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},${n}] got [${r},${a},${i}]`
+            `Inconsistent attention shapes in rollout: expected [${e},${o},${i}] got [${h},${a},${p}]`
           );
       }
-      if (o === n) {
-        const s = X(n, n).expandDims(0);
-        let r = s.tile([e, 1, 1]);
-        for (const a of t) {
-          const i = a.add(s);
-          r = i.div(i.sum(-1, !0)).matMul(r);
-        }
-        return r;
+      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);
       }
-      throw new Error(`Unsupported attention shapes for rollout: [B=${e}, Q=${o}, K=${n}]`);
+      return r;
     });
   }
-  forward(t, e, o = !1, n = !1, s) {
+  forward(t, e, o = !1, i = !1, s) {
     return this.validateInput(t), w(() => {
       this.startMemory();
-      const r = s?.[0]?.length ?? 0;
-      let a = this.inputPhase(t, r, o);
-      const i = [];
+      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 c = 0; c < this.blocks.length; c++) {
-        const u = a, m = this.blocks[c], {
+      for (let g = 0; g < this.blocks.length; g++) {
+        const u = r, m = this.blocks[g], {
           output: b,
           attention: k,
           cache: f
-        } = m.call(a, o, n, s ? s[c] : void 0);
-        a = b, u.dispose(), n && k && i.push(k), s && f ? (s[c]?.k.dispose(), s[c]?.v.dispose(), s[c] = f) : f && (f.k.dispose(), f.v.dispose());
+        } = 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());
       }
+      let h;
+      i && n.length > 0 && (h = this.computeAttentionRollout(n)), r = this.lnF.apply(r);
+      const a = this.wte.project(r);
       let p;
-      n && i.length > 0 && (p = this.computeAttentionRollout(i)), a = this.lnF.apply(a);
-      const l = this.wte.project(a);
-      let g;
-      return e && (g = this.calculateLoss(l, e)), this.endMemory("Forward"), { logits: l, loss: g, attention: n ? p : 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 n = o?.temperature ?? 1, s = o?.topK, r = o?.usePadding ?? !1, a = o?.includeAttention ?? !1;
+    const i = o?.temperature ?? 1, s = o?.topK, l = o?.usePadding ?? !1, r = o?.includeAttention ?? !1;
     return w(() => {
-      const i = t, p = i.shape[1], l = p <= this.config.gpt.blockSize ? i : i.slice(
-        [0, p - this.config.gpt.blockSize],
-        [i.shape[0], this.config.gpt.blockSize]
-      ), g = r ? this.config.gpt.blockSize - l.shape[1] : 0, c = g > 0 ? T(l, [
+      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, [
         [0, 0],
-        [0, g]
-      ]) : l, { logits: u, attention: m } = this.forward(c, void 0, !1, a, e), b = u.shape[1] - 1 - g, 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, $ = k.div(n);
+        [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;
       if (s) {
-        const { values: S, indices: I } = ot($, s), L = M(S.squeeze([1]), 1);
-        d = H(I.squeeze([1]), L, 1);
+        const { values: S, indices: I } = ot(y, s), L = M(S.squeeze([1]), 1);
+        d = J(I.squeeze([1]), L, 1);
       } else
-        d = M($.squeeze([1]), 1);
+        d = M(y.squeeze([1]), 1);
       let z;
-      return o?.includeProbabilities && (z = J($.squeeze([1]))), d = d.reshape([1, 1]), { output: d, attention: f?.squeeze([1]), probabilities: z };
+      return o?.includeProbabilities && (z = U(y.squeeze([1]))), d = d.reshape([1, 1]), { output: d, attention: f?.squeeze([1]), probabilities: z };
     });
   }
   getNumParams() {

package/dist/Reshape-CiAY8ltP.js

@@ -0,0 +1,212 @@
+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";
+/**
+ * @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 R(t, e, o = "index") {
+  const s = f(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};`;
+  }).join("");
+}
+function y(t, e) {
+  const o = t.length, s = t.map((r) => `${e}[${r}]`), n = new Array(o - 1);
+  n[o - 2] = s[o - 1];
+  for (let r = o - 3; r >= 0; --r)
+    n[r] = `(${n[r + 1]} * ${s[r + 1]})`;
+  return n;
+}
+function S(t, e, o = "index") {
+  const s = t.map((r, i) => i), n = y(s, e);
+  return n.map((r, i) => {
+    const u = `int ${t[i]} = ${o} / ${n[i]}`, a = i === n.length - 1 ? `int ${t[i + 1]} = ${o} - ${t[i]} * ${n[i]}` : `index -= ${t[i]} * ${n[i]}`;
+    return `${u}; ${a};`;
+  }).join("");
+}
+function F(t) {
+  const e = f(t).map((o) => o.toString());
+  return `
+  int getFlatIndex(ivec3 coords) {
+    return coords.x * ${e[0]} + coords.y * ${e[1]} + coords.z;
+  }
+`;
+}
+function v() {
+  return `
+  int getFlatIndex(ivec3 coords) {
+    return coords.x * outShapeStrides[0] + coords.y * outShapeStrides[1] + coords.z;
+  }
+`;
+}
+/**
+ * @license
+ * Copyright 2017 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(t, e = 2) {
+  return p(t.slice(0, t.length - e));
+}
+function m(t) {
+  if (t.length === 0)
+    throw Error("Cannot get rows and columns of an empty shape array.");
+  return [
+    t.length > 1 ? t[t.length - 2] : 1,
+    t[t.length - 1]
+  ];
+}
+function d(t) {
+  return t % 2 === 0;
+}
+function $(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) {
+    const o = t[t.length - 1], s = e[e.length - 1];
+    if (o === s || d(o) && d(s) && (t[0] === 1 || e[0] === 1))
+      return !0;
+  }
+  return t[1] === e[1] && d(t[0]) && d(e[0]);
+}
+/**
+ * @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.
+ * =============================================================================
+ */
+class b {
+  constructor(e, o) {
+    this.variableNames = ["A"], this.packedInputs = !0, this.packedOutput = !0, this.customUniforms = [{ name: "inputShape", type: "ivec3" }], this.outputShape = e, this.enableShapeUniforms = I(this.outputShape.length);
+    let s = "";
+    for (let n = 0; n < 4; n++) {
+      let r = "thisRC = rc;";
+      n % 2 === 1 && (r += "thisRC.z += 1;"), n > 1 && (r += "thisRC.y += 1;"), s += `
+        ${r}
+        ${n > 0 ? "if(thisRC.y < rows && thisRC.z < cols){" : ""}
+          int flatIndex = getFlatIndex(thisRC);
+
+          ivec3 inputRC = inputCoordsFromReshapedOutCoords(flatIndex);
+          vec2 inputRCInnerDims = vec2(float(inputRC.y),float(inputRC.z));
+
+          result[${n}] =
+            getChannel(getA(inputRC.x, inputRC.y, inputRC.z), inputRCInnerDims);
+        ${n > 0 ? "}" : ""}
+      `;
+    }
+    this.userCode = `
+      ${w(o, this.enableShapeUniforms)}
+      ${this.enableShapeUniforms ? v() : F(e)}
+
+      void main() {
+        ivec3 rc = getOutputCoords();
+
+        vec4 result = vec4(0.);
+
+        ivec3 thisRC;
+        int rows = ${this.enableShapeUniforms ? "outShape[1]" : e[1]};
+        int cols = ${this.enableShapeUniforms ? "outShape[2]" : e[2]};
+
+        ${s}
+
+        setOutput(result);
+      }
+    `;
+  }
+}
+function w(t, e) {
+  return `
+    ivec3 inputCoordsFromReshapedOutCoords(int index) {
+      ${e ? S(["r", "c", "d"], "inputShape") : R(["r", "c", "d"], t)}
+      return ivec3(r, c, d);
+    }
+  `;
+}
+/**
+ * @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 D(t, e, o) {
+  const s = [
+    h(t.shape),
+    ...m(t.shape)
+  ], n = {
+    dtype: t.dtype,
+    shape: s,
+    dataId: t.dataId
+  }, r = [
+    h(e),
+    ...m(e)
+  ], i = new b(r, s), u = !0, a = [s], c = o.runWebGLProgram(i, [n], t.dtype, a, u);
+  return { dataId: c.dataId, shape: e, dtype: c.dtype };
+}
+/**
+ * @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 k(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 });
+}
+export {
+  k 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-C4JCoBvj.js";
+import "./index--6vO-cOz.js";
 import "./ops/cpu/scatterSub.js";
 import "./ops/webgl/scatterSub.js";
 import "./ops/cpu/gatherSub.js";
@@ -31,6 +31,12 @@ import "./ops/webgl/appendCache.js";
 import "./ops/cpu/fusedSoftmax.js";
 import "./ops/webgl/fusedSoftmax.js";
 import "./ops/grads/fusedSoftmax.js";
+import "./ops/cpu/matMulGelu.js";
+import "./ops/webgl/matMulGelu.js";
+import "./ops/grads/matMulGelu.js";
+import "./ops/cpu/gelu.js";
+import "./ops/webgl/gelu.js";
+import "./ops/grads/gelu.js";
 import w from "./utilities/profile.js";
 class a {
   ee = new p();