@genai-fi/nanogpt 0.2.4 → 0.2.6

package/dist/ops/scatterSub.js

@@ -0,0 +1,150 @@
+import { engine as $ } from "@tensorflow/tfjs";
+import { i as u, j as S, k as h, E as f, l as E, o as N, c as l, n as y, r as p, a as D, m as x } from "../index-D1SlunD-.js";
+import { c as m } from "../complex-D6Bq1XDf.js";
+import { r as v, s as T } from "../stack-DB2YLlAs.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(e, t = "float32") {
+  if (u(e), t === "complex64") {
+    const a = i(e, "float32"), o = i(e, "float32");
+    return m(a, o);
+  }
+  const r = S(h(e), t);
+  return f.makeTensor(r, e, t);
+}
+/**
+ * @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 d(e, t = "float32") {
+  if (u(e), t === "complex64") {
+    const a = d(e, "float32"), o = i(e, "float32");
+    return m(a, o);
+  }
+  const r = E(h(e), t);
+  return f.makeTensor(r, e, t);
+}
+function C(e, t, r) {
+  const a = t.rank > 1 ? t.shape[t.rank - 1] : 1, o = t.rank > 1 ? t.rank - 1 : 1, s = `Must have updates.shape = indices.shape[:batchDim] + shape[sliceDim:], got updates.shape: ${r.shape}, indices.shape: ${t.shape}, shape: ${e}, sliceDim: ${a}, and batchDim: ${o}.`;
+  if (r.rank < o)
+    throw new Error(s + ` update.rank < ${o}. `);
+  if (e.length < a + (r.rank - o))
+    throw new Error(s + ` Output shape length < ${a + (r.rank - o)}`);
+  if (r.rank !== o + e.length - a)
+    throw new Error(s + ` update.rank != ${o + e.length - a}`);
+  for (let n = 0; n < o; ++n)
+    if (r.shape[n] !== t.shape[n])
+      throw new Error(s + ` updates.shape[${n}] (${r.shape[n]}) != indices.shape[${n}] (${t.shape[n]}).`);
+  for (let n = 0; n < r.rank - o; ++n)
+    if (r.shape[n + o] !== e[n + a])
+      throw new Error(s + ` updates.shape[${n + o}] (${r.shape[n + o]}) != shape[${n + o}] (${e[n + o]})`);
+}
+function O(e, t, r) {
+  if (t.rank < 1)
+    throw new Error(`tf.scatterND() expects the indices to be rank 1 or higher, but the rank was ${t.rank}.`);
+  if (e.rank < 1)
+    throw new Error(`tf.scatterND() expects the updates to be rank 1 or higher, but the rank was ${e.rank}.`);
+  if (t.dtype !== "int32")
+    throw new Error(`The dtype of 'indices' should be int32, but got dtype: ${t.dtype}`);
+  if (r.length < 1)
+    throw new Error(`Output rank must be greater or equal to 1, but got shape: ${r}`);
+  if (r.length === 0) {
+    if (t.size === 0)
+      throw new Error(`Indices specified for empty output. indices shape: ${t.shape}`);
+    if (e.size === 0)
+      throw new Error(`Updates specified for empty output. updates shape: ${e.shape}`);
+  }
+  C(r, t, e);
+}
+/**
+ * @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 z(e, t, r) {
+  u(r);
+  const a = l(e, "indices", "scatterND", "int32"), o = l(t, "updates", "scatterND");
+  O(o, a, r);
+  const s = { indices: a, updates: o }, n = { shape: r };
+  return f.runKernel(y, s, n);
+}
+const I = /* @__PURE__ */ N({ scatterND_: z });
+class L {
+  variableNames = ["labels", "softmaxProbs", "dy"];
+  outputShape;
+  userCode;
+  constructor(t, r) {
+    this.outputShape = [t, r], this.userCode = `
+      void main() {
+        ivec2 coords = getOutputCoords();
+        int index = int(getLabels(coords.x));
+        float prob = getSoftmaxProbsAtOutCoords();
+        float dy = getDy(coords.x);
+        setOutput((index == coords.y ? prob - 1.0 : prob) * dy);
+      }
+    `;
+  }
+}
+function P(e) {
+  const { logits: t, labels: r, dy: a } = e.inputs, o = e.backend, s = r.shape[0], n = t.shape[1], c = new L(s, n);
+  return o.runWebGLProgram(c, [r, t, a], "float32");
+}
+const K = {
+  kernelName: "EfficientScatterSub",
+  backendName: "webgl",
+  kernelFunc: P
+};
+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 = d([o]), g = I(c, b, [o, s]), k = D(t, g), w = a.reshape([o, 1]);
+  return x(k, w);
+}
+const F = {
+  kernelName: "EfficientScatterSub",
+  backendName: "cpu",
+  kernelFunc: A
+};
+p(F);
+function M(e, t, r) {
+  return $().runKernel("EfficientScatterSub", { logits: e, labels: t, dy: r }, {});
+}
+export {
+  M as scatterSub
+};

package/dist/stack-DB2YLlAs.js

@@ -0,0 +1,50 @@
+import { E as e, R as c, o as f, d as u, f as a, P as i } from "./index-D1SlunD-.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 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 };
+  return e.runKernel(c, {}, o);
+}
+/**
+ * @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(n, s = 0) {
+  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(i, r, o);
+}
+const h = /* @__PURE__ */ f({ stack_: k });
+export {
+  g as r,
+  h as s
+};

package/dist/sum-02UQ5Eaq.js

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

package/dist/tokeniser/CharTokeniser.d.ts

@@ -6,6 +6,7 @@ export default class CharTokeniser extends EE<'trainStatus'> implements ITokenis
     unkToken: number;
     vocab: string[];
     private cache;
+    private _trained;
     constructor(vocabSize: number);
     constructor(vocab: string[]);
     get trained(): boolean;

package/dist/tokeniser/CharTokeniser.js

@@ -1,57 +1,66 @@
-import { E as r } from "../index-Dwqa6Zy2.js";
-const h = ["<eos>", "<unk>"];
-class l extends r {
+import { E as k } from "../index-Dwqa6Zy2.js";
+const u = ["<eos>", "<unk>"];
+class b extends k {
   vocabSize = 0;
   eosToken = 0;
   unkToken = 0;
   vocab = [];
   cache = /* @__PURE__ */ new Map();
-  constructor(s) {
-    if (super(), Array.isArray(s))
-      if (this.vocab = s, this.vocab.length > 0)
-        this.vocabSize = this.vocab.length, this.eosToken = this.vocab.indexOf("<eos>"), this.unkToken = this.vocab.indexOf("<unk>"), this.unkToken === -1 && (this.unkToken = this.vocab.indexOf("<pad>")), this.unkToken === -1 && (this.unkToken = this.vocab.indexOf("_")), this.unkToken === -1 && (this.unkToken = this.vocab.indexOf(" ")), this.unkToken === -1 && (this.unkToken = this.eosToken), this.vocab.forEach((i, o) => {
-          this.cache.set(i, o);
+  _trained = !1;
+  constructor(t) {
+    if (super(), Array.isArray(t)) {
+      if (this.vocab = t, this.vocab.length > 0)
+        this.vocabSize = this.vocab.length, this.eosToken = this.vocab.indexOf("<eos>"), this.unkToken = this.vocab.indexOf("<unk>"), this.unkToken === -1 && (this.unkToken = this.vocab.indexOf("<pad>")), this.unkToken === -1 && (this.unkToken = this.vocab.indexOf("_")), this.unkToken === -1 && (this.unkToken = this.vocab.indexOf(" ")), this.unkToken === -1 && (this.unkToken = this.eosToken), this.vocab.forEach((i, n) => {
+          this.cache.set(i, n);
         });
       else
         throw new Error("Vocab cannot be empty");
-    else
-      this.vocabSize = s;
+      this._trained = !0;
+    } else
+      this.vocabSize = t, this.vocab = new Array(this.vocabSize).fill("<pad>"), this.vocab[0] = "<eos>", this.vocab[1] = "<unk>", this.eosToken = 0, this.unkToken = 1, this.cache.set("<eos>", 0), this.cache.set("<unk>", 1);
   }
   get trained() {
-    return this.vocab.length === this.vocabSize;
+    return this.vocab.length === this.vocabSize && this._trained;
   }
   destroy() {
   }
-  async train(s) {
-    const i = s.map((t) => t.split("")).flat(), o = new Set(i), e = Array.from(o), n = this.vocabSize - h.length;
-    if (e.length > n) {
-      const t = /* @__PURE__ */ new Map();
-      i.forEach((a) => {
-        t.set(a, (t.get(a) || 0) + 1);
-      }), e.sort((a, c) => (t.get(a) || 0) - (t.get(c) || 0)), e.splice(0, e.length - n);
-    } else if (e.length < n)
-      for (; e.length < n; )
-        e.push("<pad>");
-    return e.sort((t, a) => t.charCodeAt(0) - a.charCodeAt(0)), this.vocab = [...e, ...h], this.eosToken = this.vocab.indexOf("<eos>"), this.unkToken = this.vocab.indexOf("<unk>"), this.vocabSize = this.vocab.length, this.cache.clear(), this.vocab.forEach((t, a) => {
-      this.cache.set(t, a);
+  async train(t) {
+    const i = t.map((s) => s.split("")).flat(), n = new Set(i), e = Array.from(n), c = this.vocab.indexOf("<pad>"), a = this.vocabSize - u.length;
+    if (c === -1)
+      return this.vocabSize;
+    if (this._trained = !0, e.length > a) {
+      const s = /* @__PURE__ */ new Map();
+      i.forEach((o) => {
+        s.set(o, (s.get(o) || 0) + 1);
+      }), e.sort((o, r) => (s.get(o) || 0) - (s.get(r) || 0)), e.splice(0, e.length - a);
+    }
+    let h = c;
+    if (h !== -1) {
+      const s = new Set(this.vocab);
+      for (const o of e)
+        if (!s.has(o) && (this.vocab[h] = o, s.add(o), h = this.vocab.indexOf("<pad>", h + 1), h === -1))
+          break;
+    }
+    return this.cache.clear(), this.vocab.forEach((s, o) => {
+      this.cache.set(s, o);
     }), this.emit("trainStatus", "trained"), this.vocabSize;
   }
-  async tokenise(s, i) {
+  async tokenise(t, i) {
     if (!this.trained)
       throw new Error("Tokeniser not trained");
-    return s.map((e) => i ? e.split("").map((n) => this.cache.get(n) ?? this.unkToken) : e.split("").map((n) => {
-      const t = this.cache.get(n);
-      return t !== void 0 ? this.vocab[t] : "<unk>";
+    return t.map((e) => i ? e.split("").map((c) => this.cache.get(c) ?? this.unkToken) : e.split("").map((c) => {
+      const a = this.cache.get(c);
+      return a !== void 0 ? this.vocab[a] : "<unk>";
     }));
   }
-  async detokenise(s) {
-    return s.map((o) => o.map((e) => this.vocab[e]).join(""));
+  async detokenise(t) {
+    return t.map((n) => n.map((e) => this.vocab[e]).join(""));
   }
-  async encode(s) {
-    return (await this.tokenise([s], !0))[0];
+  async encode(t) {
+    return (await this.tokenise([t], !0))[0];
   }
-  async decode(s) {
-    return (await this.detokenise([s]))[0];
+  async decode(t) {
+    return (await this.detokenise([t]))[0];
   }
   getVocab() {
     return this.vocab;
@@ -59,13 +68,13 @@ class l extends r {
   async getMerges() {
     return [];
   }
-  async createTrainingData(s, i = 5) {
-    const o = await this.tokenise(s, !0), e = [], n = [];
-    for (let t = 0; t < o.length - i; t++)
-      e.push(...o[t].slice(0, i)), n.push(o[t + 1][0]);
-    return [e, n];
+  async createTrainingData(t, i = 5) {
+    const n = await this.tokenise(t, !0), e = [], c = [];
+    for (let a = 0; a < n.length - i; a++)
+      e.push(...n[a].slice(0, i)), c.push(n[a + 1][0]);
+    return [e, c];
   }
 }
 export {
-  l as default
+  b as default
 };

package/dist/training/AdamExt.js

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

package/dist/training/DatasetBuilder.js

@@ -7,7 +7,7 @@ class z {
   }
   // Create dataset from text files
   async createTextDataset(s, o = 32, i = 0, c = 1) {
-    const h = await Promise.all(s.map((t) => this.tokenizer.encode(t))), r = this.tokenizer.eosToken >= 0, n = h.map((t) => r ? [...t, this.tokenizer.eosToken] : t).flat(), a = n.slice(
+    const r = await Promise.all(s.map((t) => this.tokenizer.encode(t))), h = this.tokenizer.eosToken >= 0, n = r.map((t) => h ? [...t, this.tokenizer.eosToken] : t).flat(), a = n.slice(
       Math.floor(i * n.length),
       c === 1 ? void 0 : Math.floor(c * n.length)
     ), l = (function* () {
@@ -20,7 +20,8 @@ class z {
       const e = t;
       return this.tf.tidy(() => ({
         xs: e.xs.cast("int32"),
-        ys: this.tf.oneHot(e.ys.cast("int32"), this.tokenizer.vocabSize)
+        ys: e.ys.cast("int32")
+        // this.tf.oneHot(batchData.ys.cast('int32'), this.tokenizer.vocabSize),
       }));
     }).prefetch(2);
   }

package/dist/training/Trainer.js

@@ -1,5 +1,5 @@
 import { DatasetBuilder as d } from "./DatasetBuilder.js";
-import h from "./AdamExt.js";
+import p from "./AdamExt.js";
 class u {
   constructor(t, e, s, i = 1e-3) {
     this.tokenizer = s, this.tf = t, this.model = e, this.learningRate = i, this.resetOptimizer(), this.datasetBuilder = new d(this.tf, s, e.config.blockSize);
@@ -25,7 +25,7 @@ class u {
   }
   resetOptimizer(t = { learningRateFactor: 1, beta1: 0.9, beta2: 0.99, epsilon: 1e-8 }) {
     this.optimizer && this.optimizer.dispose();
-    const e = new h(
+    const e = new p(
       t.learningRateFactor * this.learningRate,
       t.beta1,
       t.beta2,
@@ -55,7 +55,7 @@ class u {
     });
   }
   dummyPass() {
-    const t = this.tf.zeros([1, this.model.config.blockSize], "int32"), e = this.tf.zeros([1, this.model.config.blockSize, this.model.config.vocabSize]);
+    const t = this.tf.zeros([1, this.model.config.blockSize], "int32"), e = this.tf.zeros([1, this.model.config.blockSize], "int32");
     try {
       const s = this.trainStep({ xs: t, ys: e }, !0);
       s.dataSync(), s.dispose();

package/dist/training/sparseCrossEntropy.d.ts

@@ -0,0 +1,11 @@
+import * as tf from '@tensorflow/tfjs-core';
+/**
+ * Numerically stable sparse cross-entropy with gradient support
+ * This version handles potential numerical issues better
+ */
+export declare function sparseSoftmaxCrossEntropy(logits: tf.Tensor, labels: tf.Tensor): tf.Tensor;
+/**
+ * Custom gradient implementation for sparse cross-entropy
+ * This ensures proper backpropagation
+ */
+export declare function createSoftmaxCrossEntropyWithGrad(): (...args: tf.Tensor[]) => tf.Tensor<tf.Rank>;

package/dist/training/sparseCrossEntropy.js

@@ -0,0 +1,177 @@
+import { gatherSub as K } from "../ops/gatherSub.js";
+import { scatterSub as _ } from "../ops/scatterSub.js";
+import { o as l, c as d, E as f, M as G, p as z, L as I, q as N, a as E, t as M, u as T, e as m, v as S, w as $, z as g } from "../index-D1SlunD-.js";
+import { s as F, r as b } from "../sum-02UQ5Eaq.js";
+/**
+ * @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 P(n, s, t) {
+  const a = n.length + s.length, e = [];
+  let r = 0, c = 0;
+  for (let o = 0; o < a; o++)
+    t.indexOf(o) === -1 ? e.push(n[r++]) : e.push(s[c++]);
+  return e;
+}
+function q(n, s) {
+  const t = s.map((a) => 1);
+  return P(n, t, s);
+}
+/**
+ * @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 A(n, s = null, t = !1) {
+  const e = { x: d(n, "x", "max") }, r = { reductionIndices: s, keepDims: t };
+  return f.runKernel(G, e, r);
+}
+const k = /* @__PURE__ */ l({ max_: A });
+/**
+ * @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 D(n) {
+  const t = { x: d(n, "x", "exp") };
+  return f.runKernel(z, t);
+}
+const O = /* @__PURE__ */ l({ exp_: D });
+/**
+ * @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 W(n) {
+  const t = { x: d(n, "x", "log", "float32") };
+  return f.runKernel(I, t);
+}
+const j = /* @__PURE__ */ l({ log_: W });
+/**
+ * @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 B(n, s = null, t = !1) {
+  const a = d(n, "x", "logSumExp"), e = N(s, a.shape), r = k(
+    a,
+    e,
+    !0
+    /* keepDims */
+  ), c = E(a, r), o = O(c), p = F(o, e), u = j(p), i = M(b(r, u.shape), u);
+  if (t) {
+    const h = q(i.shape, e);
+    return b(i, h);
+  }
+  return i;
+}
+const H = /* @__PURE__ */ l({ logSumExp_: B });
+/**
+ * @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 J(n, s = -1) {
+  const t = d(n, "logits", "softmax", "float32");
+  if (s === -1 && (s = t.rank - 1), s !== t.rank - 1)
+    throw Error(`Softmax along a non-last dimension is not yet supported. Logits was rank ${t.rank} and dim was ${s}`);
+  const a = { logits: t }, e = { dim: s };
+  return f.runKernel(T, a, e);
+}
+const Q = /* @__PURE__ */ l({ softmax_: J });
+function R(n, s) {
+  return $(() => {
+    const t = n.shape[n.shape.length - 1], e = n.shape.slice(0, -1).reduce((h, x) => h * x, 1), r = n.shape.length > 2 ? n.reshape([e, t]) : n, c = s.shape.length > 1 ? s.reshape([e]).cast("int32") : s.cast("int32"), o = k(r, -1, !0), p = E(r, o), u = H(p, -1);
+    return K(u, c, p);
+  });
+}
+function Z() {
+  return m().backendName === "tensorflow" ? S((s, t, a) => {
+    const e = s.shape.length > 2 ? s.reshape([-1, s.shape[s.shape.length - 1]]) : s, r = t.shape.length > 1 ? t.reshape([-1]).cast("int32") : t.cast("int32"), [c, o] = m().runKernel(
+      "NativeSparseSoftmaxCrossEntropy",
+      { logits: e, labels: r },
+      {}
+    );
+    return a([o.reshape(s.shape)]), { value: c, gradFunc: (p, u) => [u[0], g(t)] };
+  }) : S(
+    // @ts-expect-error Invalid params
+    (s, t, a) => {
+      const e = s.shape[s.shape.length - 1], c = s.shape.slice(0, -1).reduce((h, x) => h * x, 1), o = s.reshape([c, e]), p = t.reshape([c]).cast("int32"), u = R(o, p);
+      return a([o, p]), o.dispose(), p.dispose(), { value: u, gradFunc: (h, x) => $(() => {
+        const y = x[0], C = x[1], L = Q(y), v = _(L, C, h), w = g(t);
+        return [v, w];
+      }) };
+    }
+  );
+}
+export {
+  Z as createSoftmaxCrossEntropyWithGrad,
+  R as sparseSoftmaxCrossEntropy
+};

package/dist/utilities/load.js

@@ -1,16 +1,16 @@
-import { z as u } from "../jszip.min-CAxN99oA.js";
+import { j as u } from "../jszip.min-CjP2V1VV.js";
 import { importWeights as F } from "./weights.js";
 import k from "../tokeniser/CharTokeniser.js";
 import j from "../NanoGPTModel.js";
-import { dummyPassAsync as z } from "./dummy.js";
-async function A(o) {
+import { dummyPassAsync as A } from "./dummy.js";
+async function E(o) {
   const e = await fetch(o);
   if (!e.ok)
     throw new Error(`Failed to fetch ${o}: ${e.statusText}`);
   return e.arrayBuffer();
 }
 async function T(o, e) {
-  const m = typeof e == "string" ? await A(e) : e, n = await u.loadAsync(m), s = /* @__PURE__ */ new Map(), c = await n.file("manifest.json")?.async("string");
+  const m = typeof e == "string" ? await E(e) : e, n = await u.loadAsync(m), s = /* @__PURE__ */ new Map(), c = await n.file("manifest.json")?.async("string");
   if (!c)
     throw new Error("Manifest file not found in the zip archive");
   const f = JSON.parse(c);
@@ -29,7 +29,7 @@ async function T(o, e) {
     }
   o.disposeVariables();
   const a = new j(o, f.config);
-  await z(a), a.loadWeights(p);
+  await A(a), a.loadWeights(p);
   const w = await n.file("log.json")?.async("string");
   if (w)
     try {

package/dist/utilities/parameters.d.ts

@@ -0,0 +1,10 @@
+import { GPTConfig } from '../config';
+export declare function estimateParameterCount(config: GPTConfig): number;
+export declare function estimateMemoryUsage(config: GPTConfig): number;
+export declare function estimateTrainingMemoryUsage(config: GPTConfig, batchSize: number): number;
+export declare function estimateResources(config: GPTConfig, batchSize: number): {
+    numParams: number;
+    modelMemoryMB: number;
+    trainingMemoryMB: number;
+};
+export declare function validateConfig(config: GPTConfig): void;