@genai-fi/nanogpt 0.6.3 → 0.7.1

package/dist/shared-DS5waSIY.js

@@ -0,0 +1,69 @@
+import { s as l } from "./shared-Czipaeb6.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.
+ * =============================================================================
+ */
+const { addImpl: m, bincountImpl: s, bincountReduceImpl: a, bitwiseAndImpl: I, castImpl: e, ceilImpl: t, concatImpl: r, equalImpl: i, expImpl: C, expm1Impl: P, floorImpl: U, gatherNdImpl: n, gatherV2Impl: g, greaterImpl: o, greaterEqualImpl: c, lessImpl: u, lessEqualImpl: d, linSpaceImpl: q, logImpl: R, maxImpl: h, maximumImpl: x, minimumImpl: b, multiplyImpl: E, negImpl: S, notEqualImpl: T, prodImpl: w, raggedGatherImpl: y, raggedRangeImpl: A, raggedTensorToTensorImpl: F, rangeImpl: G, rsqrtImpl: N, scatterImpl: f, sigmoidImpl: k, simpleAbsImpl: B, sliceImpl: H, sparseFillEmptyRowsImpl: K, sparseReshapeImpl: V, sparseSegmentReductionImpl: j, sqrtImpl: v, staticRegexReplaceImpl: z, stridedSliceImpl: D, stringNGramsImpl: J, stringSplitImpl: L, stringToHashBucketFastImpl: M, subImpl: O, tileImpl: Q, topKImpl: W, transposeImpl: X, uniqueImpl: Y } = l;
+export {
+  b as A,
+  O as B,
+  S as C,
+  w as D,
+  y as E,
+  A as F,
+  F as G,
+  G as H,
+  N as I,
+  k as J,
+  K,
+  V as L,
+  j as M,
+  f as N,
+  v as O,
+  z as P,
+  D as Q,
+  J as R,
+  L as S,
+  M as T,
+  Q as U,
+  W as V,
+  Y as W,
+  E as a,
+  m as b,
+  H as c,
+  s as d,
+  I as e,
+  e as f,
+  t as g,
+  r as h,
+  a as i,
+  i as j,
+  C as k,
+  P as l,
+  h as m,
+  T as n,
+  U as o,
+  n as p,
+  g as q,
+  o as r,
+  B as s,
+  X as t,
+  c as u,
+  u as v,
+  d as w,
+  q as x,
+  R as y,
+  x as z
+};

package/dist/{sin-D_h-qCSx.js → sin-CiBxrDqX.js}

@@ -1,4 +1,4 @@
-import { o, q as t, E as c, a0 as a, a1 as e } from "./index-BAzbokzv.js";
+import { B as o, C as t, E as c, a8 as a, a9 as e } from "./index-UdZhlibC.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.

package/dist/slice-BHbDHObE.js

@@ -0,0 +1,28 @@
+import { B as c, C as e, E as a, W as l } from "./index-UdZhlibC.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 p(r, n, o) {
+  const s = e(r, "x", "slice", "string_or_numeric");
+  if (s.rank === 0)
+    throw new Error("Slicing scalar is not possible");
+  const t = { x: s }, i = { begin: n, size: o };
+  return a.runKernel(l, t, i);
+}
+const x = /* @__PURE__ */ c({ slice_: p });
+export {
+  x as s
+};

package/dist/{softmax-fsdtf6JC.js → softmax-JMEIUo2J.js}

@@ -1,4 +1,4 @@
-import { o as r, q as f, E as e, S as i } from "./index-BAzbokzv.js";
+import { B as r, C as f, E as e, S as i } from "./index-UdZhlibC.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.

package/dist/{split-eiktj-6L.js → split-CRU0PjVV.js}

@@ -1,4 +1,4 @@
-import { o as p, q as i, E as a, F as c } from "./index-BAzbokzv.js";
+import { B as p, C as i, E as a, L as c } from "./index-UdZhlibC.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -15,9 +15,9 @@ import { o as p, q as i, E as a, F as c } from "./index-BAzbokzv.js";
  * limitations under the License.
  * =============================================================================
  */
-function e(t, s, o = 0) {
-  const n = { x: i(t, "x", "split") }, r = { numOrSizeSplits: s, axis: o };
-  return a.runKernel(c, n, r);
+function e(t, s, n = 0) {
+  const o = { x: i(t, "x", "split") }, r = { numOrSizeSplits: s, axis: n };
+  return a.runKernel(c, o, r);
 }
 const u = /* @__PURE__ */ p({ split_: e });
 export {

package/dist/{stack-dfEEz2OY.js → stack-ikk2Y8_P.js}

@@ -1,4 +1,4 @@
-import { o as e, w as c, i as n, E as i, P as k } from "./index-BAzbokzv.js";
+import { B as e, D as c, l as n, E as k, P as i } from "./index-UdZhlibC.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -19,7 +19,7 @@ function u(r, t = 0) {
   const s = c(r, "tensors", "stack", "string_or_numeric");
   n(s.length >= 1, () => "Pass at least one tensor to tf.stack"), s.length > 0 && n(t <= s[0].rank, () => "Axis must be <= rank of the tensor");
   const o = s, a = { axis: t };
-  return i.runKernel(k, o, a);
+  return k.runKernel(i, o, a);
 }
 const l = /* @__PURE__ */ e({ stack_: u });
 export {

package/dist/{sum-BE_Irnim.js → sum-NLYbiDag.js}

@@ -1,4 +1,4 @@
-import { o as e, q as u, L as c, E as l, N as m } from "./index-BAzbokzv.js";
+import { B as e, C as u, X as c, E as l, Y as m } from "./index-UdZhlibC.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -15,10 +15,10 @@ import { o as e, q as u, L as c, E as l, N as m } from "./index-BAzbokzv.js";
  * limitations under the License.
  * =============================================================================
  */
-function i(t, o = null, n = !1) {
+function i(t, n = null, o = !1) {
   let s = u(t, "x", "sum");
   s.dtype === "bool" && (s = c(s, "int32"));
-  const r = { x: s }, a = { axis: o, keepDims: n };
+  const r = { x: s }, a = { axis: n, keepDims: o };
   return l.runKernel(m, r, a);
 }
 const f = /* @__PURE__ */ e({ sum_: i });

package/dist/{tensor-Xyi595sG.js → tensor-Do9PKbIE.js}

@@ -1,4 +1,4 @@
-import { O as t, Q as a } from "./index-BAzbokzv.js";
+import { Z as t, _ as a } from "./index-UdZhlibC.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.

package/dist/{tensor2d-CPEkynbH.js → tensor2d-CWHxHpLh.js}

@@ -1,4 +1,4 @@
-import { V as t, O as s, Q as a } from "./index-BAzbokzv.js";
+import { $ as t, Z as s, _ as a } from "./index-UdZhlibC.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.

package/dist/training/Adam.d.ts

@@ -0,0 +1,22 @@
+import { Optimizer } from '@tensorflow/tfjs-core';
+import { ConfigDict, Serializable, SerializableConstructor } from '@tensorflow/tfjs-core/dist/serialization';
+import { NamedTensor, NamedVariableMap } from '@tensorflow/tfjs-core/dist/tensor_types';
+export declare class AdamOptimizer extends Optimizer {
+    protected learningRate: number;
+    protected beta1: number;
+    protected beta2: number;
+    protected epsilon: number | null;
+    /** @nocollapse */
+    static get className(): string;
+    private accBeta1;
+    private accBeta2;
+    private accumulatedMoments;
+    constructor(learningRate: number, beta1: number, beta2: number, epsilon?: number | null);
+    applyGradients(variableGradients: NamedVariableMap | NamedTensor[]): void;
+    dispose(): void;
+    getWeights(): Promise<NamedTensor[]>;
+    setWeights(weightValues: NamedTensor[]): Promise<void>;
+    getConfig(): ConfigDict;
+    /** @nocollapse */
+    static fromConfig<T extends Serializable>(cls: SerializableConstructor<T>, config: ConfigDict): T;
+}

package/dist/training/Adam.js

@@ -0,0 +1,93 @@
+import { adamAdjust as b } from "../ops/adamAdjust.js";
+import { adamMoments as d } from "../ops/adamMoments.js";
+import { O as g, e as h, t as o, d as B } from "../index-UdZhlibC.js";
+import { z as M } from "../zeros-DnPT2nD4.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.
+ * =============================================================================
+ */
+class R extends g {
+  constructor(t, a, e, s = null) {
+    super(), this.learningRate = t, this.beta1 = a, this.beta2 = e, this.epsilon = s, this.accBeta1 = a, this.accBeta2 = e, s === null && (this.epsilon = h().backend.epsilon());
+  }
+  /** @nocollapse */
+  static get className() {
+    return "Adam";
+  }
+  accBeta1 = 0;
+  accBeta2 = 0;
+  accumulatedMoments = [];
+  applyGradients(t) {
+    const a = Array.isArray(t) ? t.map((e) => e.name) : Object.keys(t);
+    o(() => {
+      const e = 1 - this.accBeta1, s = 1 - this.accBeta2;
+      a.forEach((n, i) => {
+        const c = h().registeredVariables[n], u = !1;
+        this.accumulatedMoments[i] == null && (this.accumulatedMoments[i] = {
+          originalName: `${n}/m`,
+          variable: o(() => M([...c.shape, 2]).variable(u))
+        });
+        const r = Array.isArray(t) ? t[i].tensor : t[n];
+        if (r == null)
+          return;
+        const m = this.accumulatedMoments[i].variable, l = d(m, r, this.beta1, this.beta2);
+        m.assign(l);
+        const p = b(
+          l,
+          c,
+          e,
+          s,
+          this.epsilon ?? 1e-8,
+          this.learningRate
+        );
+        c.assign(p);
+      }), this.accBeta1 = this.accBeta1 * this.beta1, this.accBeta2 = this.accBeta2 * this.beta2;
+    }), this.incrementIterations();
+  }
+  dispose() {
+    this.accumulatedMoments != null && B(this.accumulatedMoments.map((t) => t.variable));
+  }
+  async getWeights() {
+    const t = [...this.accumulatedMoments];
+    return [await this.saveIterations()].concat(
+      t.map((a) => ({ name: a.originalName, tensor: a.variable }))
+    );
+  }
+  async setWeights(t) {
+    t = await this.extractIterations(t), o(() => {
+      this.accBeta1 = Math.pow(this.beta1, this.iterations_ + 1), this.accBeta2 = Math.pow(this.beta2, this.iterations_ + 1);
+    });
+    const a = t.length / 2, e = !1;
+    this.accumulatedMoments = t.slice(0, a).map((s) => ({
+      originalName: s.name,
+      variable: s.tensor.variable(e)
+    }));
+  }
+  getConfig() {
+    return {
+      learningRate: this.learningRate,
+      beta1: this.beta1,
+      beta2: this.beta2,
+      epsilon: this.epsilon
+    };
+  }
+  /** @nocollapse */
+  static fromConfig(t, a) {
+    return new t(a.learningRate, a.beta1, a.beta2, a.epsilon);
+  }
+}
+export {
+  R as AdamOptimizer
+};

package/dist/training/AdamExt.d.ts

@@ -1,5 +1,5 @@
-import { AdamOptimizer } from '@tensorflow/tfjs-core';
 import { NamedTensor, NamedVariableMap } from '@tensorflow/tfjs-core/dist/tensor_types';
+import { AdamOptimizer } from './Adam';
 interface AdamExtConfig {
     warmupSteps: number;
     decaySteps: number;

package/dist/training/AdamExt.js

@@ -1,7 +1,8 @@
-import { A as r, a as c, s as h, b as g, e as o } from "../index-BAzbokzv.js";
-class u extends r {
-  constructor(t, e, s, a, i) {
-    super(t, e, s, a), this.config = i, this.startLearningRate = t;
+import { a as r, b as c, c as h, e as o } from "../index-UdZhlibC.js";
+import { AdamOptimizer as g } from "./Adam.js";
+class y extends g {
+  constructor(t, e, s, i, a) {
+    super(t, e, s, i), this.config = a, this.startLearningRate = t;
   }
   step = 0;
   startLearningRate;
@@ -23,21 +24,21 @@ class u extends r {
   }
   decayVariable(t, e, s) {
     if (t && t.shape.length >= 2) {
-      const a = c(t, h(s * e));
-      t.assign(g(t, a)), a.dispose();
+      const i = r(t, c(s * e));
+      t.assign(h(t, i)), i.dispose();
     }
   }
   applyWeightDecay(t) {
-    const e = this.config.weightDecay, s = this.learningRate, a = o().registeredVariables;
-    Array.isArray(t) ? t.forEach(({ name: i }) => {
-      const n = a[i];
+    const e = this.config.weightDecay, s = this.learningRate, i = o().registeredVariables;
+    Array.isArray(t) ? t.forEach(({ name: a }) => {
+      const n = i[a];
       this.decayVariable(n, e, s);
-    }) : Object.keys(t).forEach((i) => {
-      const n = a[i];
+    }) : Object.keys(t).forEach((a) => {
+      const n = i[a];
       this.decayVariable(n, e, s);
     });
   }
 }
 export {
-  u as default
+  y as default
 };

package/dist/training/DatasetBuilder.js

@@ -1,5 +1,5 @@
-import { t as u } from "../index-BAzbokzv.js";
-import { d as z, i as f } from "../dataset-pgqp-YfL.js";
+import { t as u } from "../index-UdZhlibC.js";
+import { d as z, i as f } from "../dataset-U3PrjwgU.js";
 import "../index-Tf7vU29b.js";
 /**
  * @license

package/dist/training/Evaluator.js

@@ -11,7 +11,7 @@ class p {
       if (e.done) break;
       const c = e.value, { xs: r, ys: i } = c, [u, l] = this.model.forward({ training: !1 }, r, i);
       u.dispose(), r.dispose(), i.dispose();
-      const d = l.arraySync();
+      const d = await l.array();
       l.dispose(), t += d, o++;
     }
     return t / o;

package/dist/training/FullTrainer.js

@@ -1,21 +1,21 @@
-import { generateText as T } from "../utilities/generate.js";
-import L from "./Trainer.js";
-import x from "./Evaluator.js";
-import { d as h } from "../index-BAzbokzv.js";
-import y from "../utilities/profile.js";
-const D = {
+import { generateText as w } from "../utilities/generate.js";
+import T from "./Trainer.js";
+import L from "./Evaluator.js";
+import { d as h } from "../index-UdZhlibC.js";
+import x from "../utilities/profile.js";
+const y = {
   desiredLoss: 0.01,
   logInterval: 1,
   maxSteps: 1e3
 };
-class I extends L {
-  constructor(i, e, o = 3e-4) {
-    super(i, e, o);
+class E extends T {
+  constructor(i, e, r = 3e-4) {
+    super(i, e, r);
   }
   // Train for multiple epochs using Dataset API - FIXED memory leaks
-  async trainOnDataset(i, e, o) {
-    const { desiredLoss: p, logInterval: g, onStep: l, prompt: c, maxSteps: u } = {
-      ...D,
+  async trainOnDataset(i, e, r) {
+    const { logInterval: g, onStep: l, prompt: c, maxSteps: u } = {
+      ...y,
       ...e
     }, n = Date.now(), t = {
       step: 0,
@@ -27,22 +27,22 @@ class I extends L {
       trainingDuration: 0,
       ...this.lastState || {}
     };
-    this.lastState = t, this.dummyPass(), this.model.trainable = !0, e?.advancedMetrics && (this.model.getProfiler() || (this.model.config.layerConfig.profiler = new y())), this.running = !0, t.logStartTime = n;
-    const m = o ? new x(this.model, o) : void 0, f = await i.iterator();
+    this.lastState = t, await this.dummyPass(), this.model.trainable = !0, e?.advancedMetrics && (this.model.getProfiler() || (this.model.config.layerConfig.profiler = new x())), this.running = !0, t.logStartTime = n;
+    const m = r ? new L(this.model, r) : void 0, f = await i.iterator();
     try {
-      for (; this.running && !(t.lastLoss < p); ) {
-        const r = await f.next();
-        if (r.done) break;
-        const d = r.value, v = this.trainBatch(t, d, e.advancedMetrics || !1), s = {
+      for (; this.running; ) {
+        const o = await f.next();
+        if (o.done) break;
+        const d = o.value, p = this.trainBatch(t, d), s = {
           loss: t.lastLoss,
           step: t.step,
           time: Date.now() - n,
           batchSize: d.xs.shape[0],
-          learningRate: e?.advancedMetrics ? this.optimizer.lr : void 0,
-          gradientNorm: e?.advancedMetrics ? t.gradientNorm : void 0
+          learningRate: e?.advancedMetrics ? this.optimizer.lr : void 0
+          //gradientNorm: options?.advancedMetrics ? await state.gradientNorm : undefined,
         };
         if (this.model.log.push(s), t.step % g === 0) {
-          await v;
+          await p.data();
           const S = Date.now();
           if (t.trainingDuration += S - t.logStartTime, m)
             try {
@@ -53,10 +53,10 @@ class I extends L {
             }
           if (l) {
             if (c) {
-              const w = await T(this.tokenizer, this.model, c, 100, {
+              const v = await w(this.tokenizer, this.model, c, 100, {
                 temperature: 0.8
               });
-              s.example = w;
+              s.example = v;
             }
             const a = {
               duration: t.trainingDuration,
@@ -68,14 +68,14 @@ class I extends L {
           }
           t.logStartTime = Date.now();
         }
-        t.step >= u && this.stop();
+        p.dispose(), t.step >= u && this.stop();
       }
-    } catch (r) {
-      throw console.error("Training error:", r), h(), r;
+    } catch (o) {
+      throw console.error("Training error:", o), h(), o;
     }
     return h(), this.running = !1, { losses: t.losses, validationLosses: t.validationLosses };
   }
 }
 export {
-  I as default
+  E as default
 };

package/dist/training/Trainer.d.ts

@@ -11,7 +11,7 @@ export interface TrainingState {
     totalSteps: number;
     losses: number[];
     validationLosses: number[];
-    gradientNorm?: number;
+    gradientNorm?: Promise<number>;
 }
 export interface TrainingProgress {
     duration: number;
@@ -47,16 +47,15 @@ export default abstract class GPTTrainer {
     stop(): void;
     getOptimizer(): AdamExt;
     resetOptimizer(config?: AdamConfig): void;
-    private maxGradNorm;
-    protected trainStep(state: Partial<TrainingState>, batch: {
+    protected trainStep(_state: Partial<TrainingState>, batch: {
         xs: Tensor;
         ys: Tensor;
-    }, dummy?: boolean, calcNorm?: boolean): Scalar;
-    protected dummyPass(): void;
+    }, dummy?: boolean): Scalar;
+    protected dummyPass(): Promise<void>;
     protected trainBatch(state: TrainingState, batch: {
         xs: Tensor;
         ys: Tensor;
-    }, calcNorm?: boolean): Promise<number>;
+    }): Scalar;
     abstract trainOnDataset(dataset: Dataset<{
         xs: Tensor;
         ys: Tensor;

package/dist/training/Trainer.js

@@ -1,11 +1,10 @@
-import { DatasetBuilder as g, flattenTokens as m, PAGE_FACTOR as u } from "./DatasetBuilder.js";
-import f from "./AdamExt.js";
-import { t as y, v as z, d as c } from "../index-BAzbokzv.js";
-import { n as S } from "../norm-CzltS9Fz.js";
-import { z as p } from "../zeros-YJDE7oRb.js";
-class R {
-  constructor(t, e, s = 1e-3) {
-    this.tokenizer = e, this.model = t, this.learningRate = s, this.resetOptimizer(), this.datasetBuilder = new g(e, t.config.gpt.blockSize);
+import { DatasetBuilder as h, flattenTokens as p, PAGE_FACTOR as g } from "./DatasetBuilder.js";
+import u from "./AdamExt.js";
+import { t as f, v as y, d as c } from "../index-UdZhlibC.js";
+import { z as m } from "../zeros-DnPT2nD4.js";
+class x {
+  constructor(t, e, a = 1e-3) {
+    this.tokenizer = e, this.model = t, this.learningRate = a, this.resetOptimizer(), this.datasetBuilder = new h(e, t.config.gpt.blockSize);
   }
   model;
   optimizer;
@@ -27,7 +26,7 @@ class R {
   }
   resetOptimizer(t = { learningRateFactor: 1, beta1: 0.9, beta2: 0.99, epsilon: 1e-8 }) {
     this.optimizer && this.optimizer.dispose();
-    const e = new f(
+    const e = new u(
       t.learningRateFactor * this.learningRate,
       t.beta1,
       t.beta2,
@@ -41,76 +40,76 @@ class R {
     );
     this.optimizer = e;
   }
-  maxGradNorm(t) {
-    let e = 0;
-    return Object.keys(t).forEach((s) => {
-      const a = t[s], r = S(a), i = r.dataSync()[0];
-      r.dispose(), i > e && (e = i);
-    }), e;
-  }
-  trainStep(t, e, s = !1, a = !1) {
-    return y(() => {
+  /*private async maxGradNorm(grads: NamedVariableMap): Promise<number> {
+      let maxNorm = 0;
+      // Print all gradients
+      await Promise.all(
+          Object.keys(grads).map(async (varName) => {
+              const grad = grads[varName];
+              const temp = norm(grad);
+              const gradNorm = (await temp.data())[0];
+              temp.dispose();
+              if (gradNorm > maxNorm) {
+                  maxNorm = gradNorm;
+              }
+          })
+      );
+      return maxNorm;
+  }*/
+  trainStep(t, e, a = !1) {
+    return f(() => {
       this.model.getProfiler()?.startMemory();
-      const { xs: r, ys: i } = e, d = () => {
-        const [n, h] = this.model.forward({ training: !0 }, r, i);
-        return n.dispose(), h;
-      }, { value: l, grads: o } = z(d);
-      if (s)
-        this.model.getProfiler()?.endMemory("Training");
-      else {
-        if (a) {
-          const n = this.maxGradNorm(o);
-          t.gradientNorm = n;
-        }
-        this.optimizer.applyGradients(o), this.model.getProfiler()?.endMemory("Training"), c(o);
-      }
-      return l;
+      const { xs: s, ys: i } = e, o = () => {
+        const [l, d] = this.model.forward({ training: !0 }, s, i);
+        return l.dispose(), d;
+      }, { value: n, grads: r } = y(o);
+      return a ? this.model.getProfiler()?.endMemory("Training") : (this.optimizer.applyGradients(r), this.model.getProfiler()?.endMemory("Training"), c(r)), n;
     });
   }
-  dummyPass() {
-    const t = p([1, this.model.config.gpt.blockSize], "int32"), e = p([1, this.model.config.gpt.blockSize], "int32");
+  async dummyPass() {
+    const t = m([1, this.model.config.gpt.blockSize], "int32"), e = m([1, this.model.config.gpt.blockSize], "int32");
     try {
-      const s = this.trainStep({}, { xs: t, ys: e }, !0);
-      s.dataSync(), s.dispose();
-    } catch (s) {
-      console.error("Error during dummy pass:", s);
+      const a = this.trainStep({}, { xs: t, ys: e }, !0);
+      await a.data(), a.dispose();
+    } catch (a) {
+      console.error("Error during dummy pass:", a);
     } finally {
       t.dispose(), e.dispose();
     }
   }
-  async trainBatch(t, e, s = !1) {
+  trainBatch(t, e) {
     try {
-      const a = this.trainStep(t, e, !1, s);
-      return e.xs.dispose(), e.ys.dispose(), t.step++, t.totalSteps++, a.array().then((r) => (t.lastLoss = r, t.losses.push(t.lastLoss), a.dispose(), t.lastLoss));
+      const a = this.trainStep(t, e, !1);
+      return e.xs.dispose(), e.ys.dispose(), t.step++, t.totalSteps++, a;
     } catch (a) {
       throw console.error(`Error processing batch at step ${t.step}:`, a), c(), a;
     }
   }
-  async createTrainValidationSplit(t, e = 32, s = 0.1) {
-    const a = await m(t, this.tokenizer), r = /* @__PURE__ */ new Set();
-    if (s > 0) {
-      const l = Math.floor(a.length / (this.datasetBuilder.blockSize * u)), o = Math.max(1, Math.floor(l * s));
-      for (; r.size < o; ) {
-        const n = Math.floor(Math.random() * l);
-        r.add(n);
+  async createTrainValidationSplit(t, e = 32, a = 0.1) {
+    const s = await p(t, this.tokenizer), i = /* @__PURE__ */ new Set();
+    if (a > 0) {
+      const r = Math.floor(s.length / (this.datasetBuilder.blockSize * g)), l = Math.max(1, Math.floor(r * a));
+      for (; i.size < l; ) {
+        const d = Math.floor(Math.random() * r);
+        i.add(d);
       }
     }
-    const i = await this.datasetBuilder.createTextDataset(a, e, r, !1), d = await this.datasetBuilder.createTextDataset(
-      a,
+    const o = await this.datasetBuilder.createTextDataset(s, e, i, !1), n = await this.datasetBuilder.createTextDataset(
+      s,
       e,
-      r,
+      i,
       !0
     );
-    return { trainDataset: i, validationDataset: d };
+    return { trainDataset: o, validationDataset: n };
   }
   async createDataset(t, e = 32) {
-    const s = await m(t, this.tokenizer);
-    return await this.datasetBuilder.createTextDataset(s, e);
+    const a = await p(t, this.tokenizer);
+    return await this.datasetBuilder.createTextDataset(a, e);
   }
   dispose() {
     this.optimizer && this.optimizer.dispose();
   }
 }
 export {
-  R as default
+  x as default
 };

package/dist/training/sparseCrossEntropy.d.ts

@@ -4,8 +4,4 @@ import * as tf from '@tensorflow/tfjs-core';
  * 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>;