@genai-fi/nanogpt 0.2.10 → 0.2.12

package/dist/main.js

@@ -1,21 +1,23 @@
-import { default as m } from "./NanoGPTModel.js";
-import { default as i } from "./TeachableLLM.js";
-import { default as l } from "./tokeniser/CharTokeniser.js";
-import { default as d } from "./utilities/waitForModel.js";
-import { default as x } from "./data/textLoader.js";
-import { estimateMemoryUsage as T, estimateParameterCount as g, estimateResources as M, estimateTrainingMemoryUsage as C, validateConfig as c } from "./utilities/parameters.js";
+import { default as s } from "./NanoGPTModel.js";
+import { default as p } from "./TeachableLLM.js";
+import { default as d } from "./tokeniser/CharTokeniser.js";
+import { default as x } from "./utilities/waitForModel.js";
+import { default as T } from "./data/textLoader.js";
+import { estimateMemoryUsage as M, estimateParameterCount as C, estimateResources as c, estimateTrainingMemoryUsage as h, validateConfig as y } from "./utilities/parameters.js";
 import "./ops/scatterSub.js";
 import "./ops/gatherSub.js";
 import "./ops/attentionMask.js";
+import "./ops/qkv.js";
+import "./ops/rope.js";
 export {
-  l as CharTokeniser,
-  m as NanoGPT,
-  i as TeachableLLM,
-  T as estimateMemoryUsage,
-  g as estimateParameterCount,
-  M as estimateResources,
-  C as estimateTrainingMemoryUsage,
-  x as loadTextData,
-  c as validateConfig,
-  d as waitForModel
+  d as CharTokeniser,
+  s as NanoGPT,
+  p as TeachableLLM,
+  M as estimateMemoryUsage,
+  C as estimateParameterCount,
+  c as estimateResources,
+  h as estimateTrainingMemoryUsage,
+  T as loadTextData,
+  y as validateConfig,
+  x as waitForModel
 };

package/dist/{mat_mul-4v7St11W.js → mat_mul-Bu7bhLms.js}

@@ -1,4 +1,4 @@
-import { o as m, d as s, f as c, E as M, B as f } from "./index-CWQLouWz.js";
+import { o as m, d as s, f as c, E as M, B as f } from "./index-YPKosni4.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.

package/dist/ops/attentionMask.js

@@ -1,6 +1,6 @@
 import { engine as k } from "@tensorflow/tfjs";
-import { r as m, c as d, s as p } from "../index-CWQLouWz.js";
-import { m as f } from "../mat_mul-4v7St11W.js";
+import { r as m, c as d, s as p } from "../index-YPKosni4.js";
+import { m as f } from "../mat_mul-Bu7bhLms.js";
 class h {
   variableNames = ["q", "k", "mask"];
   outputShape;

package/dist/ops/gatherSub.js

@@ -1,6 +1,6 @@
 import { engine as l } from "@tensorflow/tfjs";
-import { o as g, d as i, E as b, G as d, r as c, b as h } from "../index-CWQLouWz.js";
-import { r as p, s as f } from "../stack-CTdK-itU.js";
+import { o as g, d as i, E as b, G as d, r as c, b as h } from "../index-YPKosni4.js";
+import { r as p, s as f } from "../stack-BtKpB0Ry.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.

package/dist/ops/node/sparseCrossEntropy.js

@@ -1,4 +1,4 @@
-import { r as o } from "../../index-CWQLouWz.js";
+import { r as o } from "../../index-YPKosni4.js";
 function r(e) {
   const { logits: t, labels: n } = e.inputs;
   return e.backend.executeMultipleOutputs("SparseSoftmaxCrossEntropyWithLogits", [], [t, n], 2);

package/dist/ops/qkv.d.ts

@@ -0,0 +1,7 @@
+import { Tensor } from '@tensorflow/tfjs';
+import { TensorInfo, NamedTensorInfoMap, NamedAttrMap } from '@tensorflow/tfjs-core';
+export declare function qkvCPU(args: {
+    inputs: NamedTensorInfoMap;
+    attrs?: NamedAttrMap;
+}): TensorInfo[];
+export declare function qkv(x: Tensor, kernel: Tensor, heads: number): Tensor[];

package/dist/ops/qkv.js

@@ -0,0 +1,127 @@
+import { engine as C } from "@tensorflow/tfjs";
+import { o as N, d as V, E as Q, i as T, r as b, c as G } from "../index-YPKosni4.js";
+import { r as x } from "../reshape-DmnmKT6r.js";
+/**
+ * @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 M(r, e, o = 0) {
+  const t = { x: V(r, "x", "split") }, s = { numOrSizeSplits: e, axis: o };
+  return Q.runKernel(T, t, s);
+}
+const P = /* @__PURE__ */ N({ split_: M });
+class S {
+  variableNames = ["x", "kernel"];
+  outputShape;
+  userCode;
+  // enableShapeUniforms = true;
+  customUniforms = [{ name: "mode", type: "int" }];
+  constructor(e, o, a, t) {
+    const s = t / o;
+    this.outputShape = [e, o, a, s], this.userCode = `
+        void main() {
+            ivec4 coords = getOutputCoords(); // [b, h, t, d]
+            int b = coords.x;
+            int h = coords.y;
+            int t = coords.z;
+            int d = coords.w;
+
+            // Compute output channel index in fused kernel
+            int out_offset = mode * ${o} * ${s} + h * ${s} + d;
+
+            float sum = 0.0;
+            for (int c = 0; c < ${t}; ++c) {
+                float xval = getX(b, t, c); // fetch from x
+                float kval = getKernel(c, out_offset); // fetch from kernel
+                sum += xval * kval;
+            }
+
+            setOutput(sum);
+        }
+        `;
+  }
+}
+function $(r) {
+  const { x: e, kernel: o } = r.inputs, { heads: a } = r.attrs, t = r.backend, s = e.shape[0], l = e.shape[1], d = e.shape[2], c = new S(s, a, l, d);
+  return [
+    t.runWebGLProgram(c, [e, o], "float32", [[0]]),
+    t.runWebGLProgram(c, [e, o], "float32", [[1]]),
+    t.runWebGLProgram(c, [e, o], "float32", [[2]])
+  ];
+}
+const w = {
+  kernelName: "QKV",
+  backendName: "webgl",
+  kernelFunc: $
+};
+b(w);
+function q(r) {
+  const { x: e, kernel: o } = r.inputs, { heads: a } = r.attrs, [t, s, l] = e.shape, d = x(e, [t * s, l]), c = d.dot(o);
+  d.dispose();
+  const n = x(c, [t, s, 3 * l]);
+  c.dispose();
+  const [p, m, k] = P(n, 3, -1);
+  n.dispose();
+  const f = l / a, v = x(p, [t, s, a, f]);
+  p.dispose();
+  const g = v.transpose([0, 2, 1, 3]);
+  v.dispose();
+  const u = x(m, [t, s, a, f]);
+  m.dispose();
+  const h = u.transpose([0, 2, 1, 3]);
+  u.dispose();
+  const i = x(k, [t, s, a, f]);
+  k.dispose();
+  const K = i.transpose([0, 2, 1, 3]);
+  return i.dispose(), [g, h, K];
+}
+const F = {
+  kernelName: "QKV",
+  backendName: "cpu",
+  kernelFunc: q
+};
+b(F);
+const L = {
+  kernelName: "QKV",
+  backendName: "tensorflow",
+  kernelFunc: q
+};
+b(L);
+function y(r, e, o) {
+  return C().runKernel("QKV", { x: r, kernel: e }, { heads: o });
+}
+const _ = {
+  kernelName: "QKV",
+  inputsToSave: ["x", "kernel"],
+  outputsToSave: [],
+  gradFunc: (r, e) => {
+    const [o, a, t] = r, [s, l] = e, [d, c, n] = s.shape, p = o.transpose([0, 2, 1, 3]).reshape([d * c, n]), m = a.transpose([0, 2, 1, 3]).reshape([d * c, n]), k = t.transpose([0, 2, 1, 3]).reshape([d * c, n]), f = l.slice([0, 0], [n, n]), v = l.slice([0, n], [n, n]), g = l.slice([0, 2 * n], [n, n]);
+    return {
+      x: () => {
+        const u = p.matMul(f, !1, !0), h = m.matMul(v, !1, !0), i = k.matMul(g, !1, !0);
+        return u.add(h).add(i).reshape([d, c, n]);
+      },
+      kernel: () => {
+        const u = s.reshape([d * c, n]), h = u.matMul(p, !0, !1), i = u.matMul(m, !0, !1), K = u.matMul(k, !0, !1);
+        return h.concat(i, 1).concat(K, 1);
+      }
+    };
+  }
+};
+G(_);
+export {
+  y as qkv,
+  q as qkvCPU
+};

package/dist/ops/rope.d.ts

@@ -0,0 +1,8 @@
+import { default as RoPECache } from '../layers/RoPECache';
+import { Tensor } from '@tensorflow/tfjs';
+import { TensorInfo, NamedTensorInfoMap, NamedAttrMap } from '@tensorflow/tfjs-core';
+export declare function ropeCPU(args: {
+    inputs: NamedTensorInfoMap;
+    attrs?: NamedAttrMap;
+}): TensorInfo;
+export declare function rope(x: Tensor, cache: RoPECache, pastLength: number): Tensor;

package/dist/ops/rope.js

@@ -0,0 +1,154 @@
+import { engine as D } from "@tensorflow/tfjs";
+import { o as G, l as F, k as _, n as U, E as K, p as z, d as I, q as O, r as f, c as A } from "../index-YPKosni4.js";
+import { r as T, s as B } from "../stack-BtKpB0Ry.js";
+/**
+ * @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 W(o, t = 0) {
+  F(o.length >= 1, () => "Pass at least one tensor to concat");
+  const e = _(o, "tensors", "concat", "string_or_numeric");
+  if (e[0].dtype === "complex64" && e.forEach((r) => {
+    if (r.dtype !== "complex64")
+      throw new Error(`Cannot concatenate complex64 tensors with a tensor
+          with dtype ${r.dtype}. `);
+  }), e.length === 1)
+    return U(e[0]);
+  const n = e, s = { axis: t };
+  return K.runKernel(z, n, s);
+}
+const j = /* @__PURE__ */ G({ concat_: W });
+/**
+ * @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 H(o, t, e = 0, n = 0) {
+  const s = I(o, "x", "gather"), r = I(t, "indices", "gather", "int32"), c = { x: s, indices: r }, a = { axis: e, batchDims: n };
+  return K.runKernel(O, c, a);
+}
+const E = /* @__PURE__ */ G({ gather_: H });
+class J {
+  variableNames = ["x", "sin", "cos"];
+  outputShape;
+  userCode;
+  // enableShapeUniforms = true;
+  customUniforms = [{ name: "pastLen", type: "int" }];
+  constructor(t, e, n, s) {
+    this.outputShape = [t, e, n, s], this.userCode = `
+        void main() {
+            ivec4 coords = getOutputCoords(); // [b, h, t, d]
+            int b = coords.x;
+            int h = coords.y;
+            int t = coords.z;
+            int d = coords.w;
+
+            int rotaryDim = ${s};
+
+            float outVal = 0.0;
+
+            if (d < rotaryDim) {
+                int pairIdx = d / 2;
+                float cos = getCos(t + pastLen, pairIdx, 0);
+                float sin = getSin(t + pastLen, pairIdx, 0);
+
+                if (d % 2 == 0) {
+                    // even index
+                    float even = getX(b, h, t, d);
+                    float odd = getX(b, h, t, d + 1);
+                    outVal = even * cos - odd * sin;
+                } else {
+                    // odd index
+                    float even = getX(b, h, t, d - 1);
+                    float odd = getX(b, h, t, d);
+                    outVal = even * sin + odd * cos;
+                }
+            } else {
+                // pass through for non-rotary dims
+                outVal = getX(b, h, t, d);
+            }
+
+            setOutput(outVal);
+        }
+        `;
+  }
+}
+function M(o) {
+  const { x: t, sin: e, cos: n } = o.inputs, { pastLen: s } = o.attrs, r = o.backend, c = t.shape[0], a = t.shape[1], i = t.shape[2], d = t.shape[3], p = new J(c, a, i, d);
+  return r.runWebGLProgram(p, [t, e, n], "float32", [[s]]);
+}
+const Q = {
+  kernelName: "Rope",
+  backendName: "webgl",
+  kernelFunc: M
+};
+f(Q);
+function V(o, t, e, n, s) {
+  const r = n.shape[3], c = e;
+  if (c > r) return n;
+  const a = n.shape[2], i = c / 2, d = t.slice([s, 0, 0], [a, i, 1]).reshape([1, 1, a, i]), p = o.slice([s, 0, 0], [a, i, 1]).reshape([1, 1, a, i]), u = n.shape[0], l = n.shape[1], g = T(0, c, 2, "int32"), x = T(1, c, 2, "int32"), $ = ((b) => {
+    const v = b.slice([0, 0, 0, 0], [u, l, a, c]), k = c < r ? b.slice([0, 0, 0, c], [u, l, a, r - c]) : null, h = E(v, g, 3), m = E(v, x, 3), C = h.mul(d), y = m.mul(p), R = C.sub(y), N = m.mul(d), S = h.mul(p), w = N.add(S);
+    h.dispose(), m.dispose(), d.dispose(), p.dispose(), C.dispose(), y.dispose(), N.dispose(), S.dispose();
+    const L = B([R, w], -1);
+    R.dispose(), w.dispose();
+    const P = L.reshape([u, l, a, c]);
+    return L.dispose(), k ? j([P, k], 3) : P;
+  })(n);
+  return g.dispose(), x.dispose(), $;
+}
+function X(o) {
+  const { x: t, sin: e, cos: n } = o.inputs, { pastLen: s } = o.attrs, r = t.shape[3];
+  return V(e, n, r, t, s);
+}
+const Y = {
+  kernelName: "Rope",
+  backendName: "cpu",
+  kernelFunc: X
+};
+f(Y);
+const Z = {
+  kernelName: "Rope",
+  backendName: "tensorflow",
+  kernelFunc: X
+};
+f(Z);
+function st(o, t, e) {
+  return t.ensureRopeCache(o.shape[1] + e), D().runKernel("Rope", { x: o, sin: t.getSin(), cos: t.getCos() }, { pastLen: e });
+}
+const q = {
+  kernelName: "Rope",
+  inputsToSave: ["x", "sin", "cos"],
+  outputsToSave: [],
+  gradFunc: (o, t) => {
+    const [e, n, s] = t, r = n.neg(), c = e.shape[3], i = V(r, s, c, o, 0);
+    return r.dispose(), { x: () => i };
+  }
+};
+A(q);
+export {
+  st as rope,
+  X as ropeCPU
+};

package/dist/ops/scatterSub.js

@@ -1,7 +1,7 @@
 import { engine as $ } from "@tensorflow/tfjs";
-import { l as u, n as S, p, E as f, q as E, o as N, d as l, t as y, r as h, b as D, a as x } from "../index-CWQLouWz.js";
-import { c as d } from "../complex-x7w5HPOS.js";
-import { r as v, s as T } from "../stack-CTdK-itU.js";
+import { t as u, u as S, v as l, E as f, w as E, o as N, d as h, x as y, r as p, b as x, a as D } from "../index-YPKosni4.js";
+import { c as d } from "../complex-CJ-qCcLB.js";
+import { r as v, s as T } from "../stack-BtKpB0Ry.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -23,7 +23,7 @@ function i(e, t = "float32") {
     const a = i(e, "float32"), o = i(e, "float32");
     return d(a, o);
   }
-  const r = S(p(e), t);
+  const r = S(l(e), t);
   return f.makeTensor(r, e, t);
 }
 /**
@@ -47,7 +47,7 @@ function m(e, t = "float32") {
     const a = m(e, "float32"), o = i(e, "float32");
     return d(a, o);
   }
-  const r = E(p(e), t);
+  const r = E(l(e), t);
   return f.makeTensor(r, e, t);
 }
 function C(e, t, r) {
@@ -100,7 +100,7 @@ function O(e, t, r) {
  */
 function z(e, t, r) {
   u(r);
-  const a = l(e, "indices", "scatterND", "int32"), o = l(t, "updates", "scatterND");
+  const a = h(e, "indices", "scatterND", "int32"), o = h(t, "updates", "scatterND");
   O(o, a, r);
   const s = { indices: a, updates: o }, n = { shape: r };
   return f.runKernel(y, s, n);
@@ -131,17 +131,17 @@ const K = {
   backendName: "webgl",
   kernelFunc: P
 };
-h(K);
+p(K);
 function A(e) {
-  const { logits: t, labels: r, dy: a } = e.inputs, o = r.shape[0], s = t.shape[1], n = v(0, o, 1, "int32"), c = T([n, r], 1), b = m([o]), g = I(c, b, [o, s]), k = D(t, g), w = a.reshape([o, 1]);
-  return x(k, w);
+  const { logits: t, labels: r, dy: a } = e.inputs, o = r.shape[0], s = t.shape[1], n = v(0, o, 1, "int32"), c = T([n, r], 1), b = m([o]), g = I(c, b, [o, s]), k = x(t, g), w = a.reshape([o, 1]);
+  return D(k, w);
 }
 const F = {
   kernelName: "EfficientScatterSub",
   backendName: "cpu",
   kernelFunc: A
 };
-h(F);
+p(F);
 function R(e, t, r) {
   return $().runKernel("EfficientScatterSub", { logits: e, labels: t, dy: r }, {});
 }

package/dist/reshape-DmnmKT6r.js

@@ -0,0 +1,25 @@
+import { o, d as t, E as a, R as p } from "./index-YPKosni4.js";
+/**
+ * @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 c(r, s) {
+  const e = { x: t(r, "x", "reshape", "string_or_numeric") }, n = { shape: s };
+  return a.runKernel(p, e, n);
+}
+const x = /* @__PURE__ */ o({ reshape_: c });
+export {
+  x as r
+};

package/dist/{stack-CTdK-itU.js → stack-BtKpB0Ry.js}

@@ -1,4 +1,4 @@
-import { E as e, R as c, o as f, h as i, i as a, P as u } from "./index-CWQLouWz.js";
+import { E as e, j as c, o as f, k as u, l as a, P as i } from "./index-YPKosni4.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -15,7 +15,7 @@ import { E as e, R as c, o as f, h as i, i as a, P as u } from "./index-CWQLouWz
  * limitations under the License.
  * =============================================================================
  */
-function l(n, s, t = 1, r = "float32") {
+function g(n, s, t = 1, r = "float32") {
   if (t === 0)
     throw new Error("Cannot have a step of zero");
   const o = { start: n, stop: s, step: t, dtype: r };
@@ -38,13 +38,13 @@ function l(n, s, t = 1, r = "float32") {
  * =============================================================================
  */
 function k(n, s = 0) {
-  const t = i(n, "tensors", "stack", "string_or_numeric");
+  const t = u(n, "tensors", "stack", "string_or_numeric");
   a(t.length >= 1, () => "Pass at least one tensor to tf.stack"), t.length > 0 && a(s <= t[0].rank, () => "Axis must be <= rank of the tensor");
   const r = t, o = { axis: s };
-  return e.runKernel(u, r, o);
+  return e.runKernel(i, r, o);
 }
-const g = /* @__PURE__ */ f({ stack_: k });
+const h = /* @__PURE__ */ f({ stack_: k });
 export {
-  l as r,
-  g as s
+  g as r,
+  h as s
 };

package/dist/sum-D7fu15XL.js

@@ -0,0 +1,27 @@
+import { o as e, d as u, h as c, E as l, S as m } from "./index-YPKosni4.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 i(t, o = null, n = !1) {
+  let s = u(t, "x", "sum");
+  s.dtype === "bool" && (s = c(s, "int32"));
+  const r = { x: s }, a = { axis: o, keepDims: n };
+  return l.runKernel(m, r, a);
+}
+const f = /* @__PURE__ */ e({ sum_: i });
+export {
+  f as s
+};

package/dist/training/AdamExt.js

@@ -1,4 +1,4 @@
-import { A as r, a as c, s as h, b as g, e as o } from "../index-CWQLouWz.js";
+import { A as r, a as c, s as h, b as g, e as o } from "../index-YPKosni4.js";
 class u extends r {
   constructor(t, e, s, a, i) {
     super(t, e, s, a), this.config = i, this.startLearningRate = t;