@genai-fi/nanogpt 0.17.4 → 0.18.0

package/dist/tokeniser/bpe.js

@@ -1,15 +1,15 @@
-import { yieldIfNeeded as f } from "../utilities/yielder.js";
+import { yieldIfNeeded as p } from "../utilities/yielder.js";
 import m from "../utilities/tokenParse.js";
-import z, { SPECIALS as k } from "./BaseTokeniser.js";
-function p(o, e) {
+import T, { SPECIALS as S } from "./BaseTokeniser.js";
+function g(o, e) {
   return `${o}-::-${e}`;
 }
-function w(o) {
+function y(o) {
   const e = /* @__PURE__ */ new Map();
   for (let s = 0; s < o.length; s++) {
     const t = o[s];
     for (let n = 0; n < t.length - 1; n++) {
-      const r = p(t[n], t[n + 1]), a = e.get(r) || {
+      const r = g(t[n], t[n + 1]), a = e.get(r) || {
         a: t[n],
         b: t[n + 1],
         count: 0,
@@ -20,21 +20,21 @@ function w(o) {
   }
   return { pairs: e, tokens: o };
 }
-function d(o, e, s, t, n) {
-  const r = p(e, s);
+function f(o, e, s, t, n) {
+  const r = g(e, s);
   if (o.pairs.has(r)) {
     const a = o.pairs.get(r);
     a.count += n, n > 0 ? a.instances.add(t) : a.count <= 0 ? o.pairs.delete(r) : a.instances.delete(t);
   } else
     o.pairs.set(r, { a: e, b: s, count: n, instances: /* @__PURE__ */ new Set([t]) });
 }
-function T(o) {
+function I(o) {
   let e = null, s = 0;
   for (const t of o.pairs.values())
     t.count > s && (s = t.count, e = t);
   return e;
 }
-function y(o, e) {
+function x(o, e) {
   return o.map((s) => {
     const t = [];
     for (let n = 0; n < s.length; n++)
@@ -42,19 +42,19 @@ function y(o, e) {
     return t;
   });
 }
-function I(o, e) {
+function A(o, e) {
   e.instances.forEach((s) => {
     const t = o.tokens[s], n = [];
     for (let r = 0; r < t.length; r++)
       if (r < t.length - 1 && t[r] === e.a && t[r + 1] === e.b) {
         const a = e.a + e.b;
-        n.push(a), r > 0 && (d(o, t[r - 1], e.a, s, -1), d(o, t[r - 1], a, s, 1)), r++, r < t.length - 1 && (d(o, e.b, t[r + 1], s, -1), d(o, a, t[r + 1], s, 1));
+        n.push(a), r > 0 && (f(o, t[r - 1], e.a, s, -1), f(o, t[r - 1], a, s, 1)), r++, r < t.length - 1 && (f(o, e.b, t[r + 1], s, -1), f(o, a, t[r + 1], s, 1));
       } else
         n.push(t[r]);
     o.tokens[s] = n;
-  }), o.pairs.delete(p(e.a, e.b));
+  }), o.pairs.delete(g(e.a, e.b));
 }
-class E extends z {
+class P extends T {
   targetSize;
   vocab = /* @__PURE__ */ new Set();
   vocabIndex = /* @__PURE__ */ new Map();
@@ -63,7 +63,7 @@ class E extends z {
   constructor(e, s) {
     super(), Array.isArray(e) ? (e.forEach((t, n) => {
       this.vocab.add(t), this.vocabIndex.set(t, n);
-    }), s && (this.merges = s), this.targetSize = e.length, k.forEach((t) => {
+    }), s && (this.merges = s), this.targetSize = e.length, S.forEach((t) => {
       const n = e.indexOf(t);
       n !== -1 && this.addSpecialToken(t, n);
     })) : (this.addSpecialTokens(), this.targetSize = e);
@@ -81,7 +81,7 @@ class E extends z {
     this.vocab.clear(), this.vocabIndex.clear(), this.merges = [], this.pretokenMap.clear();
   }
   get trained() {
-    return this.vocab.size > k.length && this.vocab.size <= this.targetSize;
+    return this.vocab.size > S.length && this.vocab.size <= this.targetSize;
   }
   get vocabSize() {
     return this.vocab.size;
@@ -95,42 +95,48 @@ class E extends z {
   get unkToken() {
     return this.vocabIndex.get("") ?? 1;
   }
-  async train(e = [], s) {
-    let t = performance.now();
-    const n = e.map((i) => i.map((h) => m(h.content))).flat(2);
-    t = await f(t, s, this.vocab.size);
-    const r = new Set(n);
+  async train(e = [], s, t) {
+    this.datasetID = t;
+    let n = performance.now();
+    const r = new Array(e.length);
+    for (let i = 0; i < e.length; i++) {
+      const h = e[i], l = new Array(h.length);
+      for (let d = 0; d < h.length; d++)
+        l[d] = m(h[d].content);
+      n = await p(n, s, this.vocab.size), r[i] = l;
+    }
+    const a = r.flat(2), z = new Set(a);
     this.vocab = /* @__PURE__ */ new Set(), this.pretokenMap.clear(), this.merges = [], this.addSpecialTokens();
-    const a = Array.from(r), b = a.map((i) => Array.from(i).map((l) => (this.vocab.add(l), l))), g = w(b);
-    if (t = await f(t, s, this.vocab.size), this.vocab.size >= this.targetSize) {
+    const b = Array.from(z), v = b.map((i) => Array.from(i).map((l) => (this.vocab.add(l), l))), k = y(v);
+    if (n = await p(n, s, this.vocab.size), this.vocab.size >= this.targetSize) {
       console.warn("Initial vocab size is greater than or equal to target size. No merges will be performed.");
       const i = /* @__PURE__ */ new Map();
-      n.forEach((c) => {
+      a.forEach((c) => {
         Array.from(c).forEach((u) => {
           i.set(u, (i.get(u) || 0) + 1);
         });
       });
       const h = Array.from(i.entries()).sort((c, u) => u[1] - c[1]);
       this.vocab = /* @__PURE__ */ new Set(), this.addSpecialTokens(), h.slice(0, this.targetSize - this.vocab.size).map(([c]) => c).forEach((c) => this.vocab.add(c)), this.vocabIndex.clear();
-      let S = 0;
+      let d = 0;
       for (const c of this.vocab.keys())
-        this.vocabIndex.set(c, S++);
-      return this.emit("trainStatus", "trained"), this.vocab.size;
+        this.vocabIndex.set(c, d++);
+      return this.generateID(), this.emit("trainStatus", "trained"), this.vocab.size;
     }
     for (; this.vocab.size < this.targetSize && this.merges.length < this.targetSize; ) {
-      const i = T(g);
+      const i = I(k);
       if (!i)
         break;
-      this.merges.push([i.a, i.b]), this.vocab.add(i.a + i.b), I(g, i), t = await f(t, s, this.vocab.size);
+      this.merges.push([i.a, i.b]), this.vocab.add(i.a + i.b), A(k, i), n = await p(n, s, this.vocab.size);
     }
-    a.forEach((i, h) => {
-      const l = b[h];
+    b.forEach((i, h) => {
+      const l = v[h];
       this.pretokenMap.set(i, l);
     }), this.vocabIndex.clear();
-    let v = 0;
+    let w = 0;
     for (const i of this.vocab.keys())
-      this.vocabIndex.set(i, v++);
-    return this.emit("trainStatus", "trained"), this.vocab.size;
+      this.vocabIndex.set(i, w++);
+    return this.generateID(), this.emit("trainStatus", "trained"), this.vocab.size;
   }
   getVocab() {
     return Array.from(this.vocab);
@@ -141,7 +147,7 @@ class E extends z {
   tokeniseWord(e) {
     let s = Array.from(e);
     return this.merges.forEach((t) => {
-      s = y([s], t)[0];
+      s = x([s], t)[0];
     }), this.pretokenMap.set(e, s), s;
   }
   tokeniseStrings(e) {
@@ -163,5 +169,5 @@ class E extends z {
   }
 }
 export {
-  E as default
+  P as default
 };

package/dist/tokeniser/type.d.ts

@@ -5,7 +5,9 @@ export interface Conversation {
     content: string;
 }
 export interface ITokeniser extends EE<'trainStatus'> {
-    train(text: Conversation[][], cb?: (vocab: number) => void): Promise<number>;
+    id: string;
+    datasetID?: string;
+    train(text: Conversation[][], cb?: (vocab: number) => void, datasetID?: string): Promise<number>;
     getVocab(): string[];
     getMerges(): [string, string][];
     destroy(): void;

package/dist/training/AdamW.d.ts

@@ -21,6 +21,9 @@ export declare class AdamWOptimizer extends Optimizer {
     protected orthGrad: boolean;
     constructor(config: AdamWOptimizerConfig);
     get lr(): number;
+    saveMoments(): Promise<ArrayBuffer>;
+    loadMoments(momentData: ArrayBuffer): Promise<void>;
+    serializeConfig(): AdamWOptimizerConfig;
     private orthogonalizeGradient;
     updateConfig(newConfig: Partial<AdamWOptimizerConfig>): void;
     applyGradients(variableGradients: NamedVariableMap | NamedTensor[]): Tensor;

package/dist/training/AdamW.js

@@ -1,12 +1,13 @@
-import { adamAdjust as N } from "../ops/adamAdjust.js";
-import { adamMoments as S } from "../ops/adamMoments.js";
-import { O as R, h as b, t as h, a as w, d as B } from "../index-CUXkjxiT.js";
-import M from "./LRScheduler.js";
-import { clipScale as A } from "../ops/globalNorm.js";
-import { z as O } from "../zeros-DvZpK8s6.js";
-class G extends R {
+import { adamAdjust as B } from "../ops/adamAdjust.js";
+import { adamMoments as N } from "../ops/adamMoments.js";
+import { O as S, e as b, t as c, b as M, l as w } from "../index-DSGwv2Yx.js";
+import R from "./LRScheduler.js";
+import { clipScale as f } from "../ops/globalNorm.js";
+import { save_safetensors as v, load_safetensors as A } from "../utilities/safetensors.js";
+import { z as O } from "../zeros-Bw0puq_w.js";
+class _ extends S {
   constructor(t) {
-    super(), this.config = t, this.accBeta1 = t.beta1, this.accBeta2 = t.beta2, this.learningRate = t.learningRate, this.beta1 = t.beta1, this.beta2 = t.beta2, this.weightDecay = t.weightDecay, this.lossScaling = t.lossScaling, this.clipNorm = t.clipNorm, this.orthGrad = t.orthoGrad ?? !1, t.epsilon === null || t.epsilon === void 0 ? this.epsilon = b().backend.epsilon() : this.epsilon = t.epsilon, this.lrScheduler = new M(t.learningRate, t);
+    super(), this.config = t, this.accBeta1 = t.accBeta1 ?? t.beta1, this.accBeta2 = t.accBeta2 ?? t.beta2, this.learningRate = t.learningRate, this.beta1 = t.beta1, this.beta2 = t.beta2, this.weightDecay = t.weightDecay, this.lossScaling = t.lossScaling, this.clipNorm = t.clipNorm, this.orthGrad = t.orthoGrad ?? !1, t.epsilon === null || t.epsilon === void 0 ? this.epsilon = b().backend.epsilon() : this.epsilon = t.epsilon, this.lrScheduler = new R(t.learningRate, t);
   }
   className = "AdamW";
   accBeta1 = 0;
@@ -25,10 +26,38 @@ class G extends R {
   get lr() {
     return this.learningRate;
   }
+  saveMoments() {
+    const t = {};
+    return this.accumulatedMoments.forEach((e) => {
+      t[e.originalName] = e.variable;
+    }), v(t);
+  }
+  async loadMoments(t) {
+    const e = await A(t);
+    Object.entries(e).forEach(([a, s]) => {
+      const n = s.variable(!1);
+      this.accumulatedMoments.push({ originalName: a, variable: n });
+    });
+  }
+  serializeConfig() {
+    return {
+      learningRate: this.learningRate,
+      beta1: this.beta1,
+      beta2: this.beta2,
+      accBeta1: this.accBeta1,
+      accBeta2: this.accBeta2,
+      epsilon: this.epsilon ?? void 0,
+      weightDecay: this.weightDecay,
+      lossScaling: this.lossScaling,
+      clipNorm: this.clipNorm,
+      orthoGrad: this.orthGrad,
+      ...this.lrScheduler.serializeConfig()
+    };
+  }
   orthogonalizeGradient(t, e) {
-    return h(() => {
-      const a = t.reshape([-1]), s = e.reshape([-1]), l = a.mul(a).sum().add(this.orthGradEpsilon), c = a.mul(s).sum().div(l), n = s.sub(a.mul(c)), o = s.norm(), i = n.norm().add(this.orthGradEpsilon);
-      return n.mul(o.div(i)).reshape(e.shape);
+    return c(() => {
+      const a = t.reshape([-1]), s = e.reshape([-1]), n = a.mul(a).sum().add(this.orthGradEpsilon), h = a.mul(s).sum().div(n), o = s.sub(a.mul(h)), l = s.norm(), i = o.norm().add(this.orthGradEpsilon);
+      return o.mul(l.div(i)).reshape(e.shape);
     });
   }
   updateConfig(t) {
@@ -38,42 +67,42 @@ class G extends R {
   applyGradients(t) {
     const e = this.lrScheduler.getNextLR();
     this.learningRate = e;
-    const a = Array.isArray(t) ? t.map((l) => l.name) : Object.keys(t), s = h(() => {
-      const l = 1 - this.accBeta1, c = 1 - this.accBeta2;
-      let n;
+    const a = Array.isArray(t) ? t.map((n) => n.name) : Object.keys(t), s = c(() => {
+      const n = 1 - this.accBeta1, h = 1 - this.accBeta2;
+      let o;
       if (this.clipNorm !== void 0) {
-        const o = a.map((i, r) => Array.isArray(t) ? t[r].tensor : t[i]);
-        n = A(o, 1 / this.lossScaling, this.clipNorm);
+        const l = a.map((i, r) => Array.isArray(t) ? t[r].tensor : t[i]);
+        o = f(l, 1 / this.lossScaling, this.clipNorm);
       } else
-        n = w(1 / this.lossScaling);
-      return a.forEach((o, i) => {
-        const r = b().registeredVariables[o], p = !1;
+        o = M(1 / this.lossScaling);
+      return a.forEach((l, i) => {
+        const r = b().registeredVariables[l], p = !1;
         this.accumulatedMoments[i] == null && (this.accumulatedMoments[i] = {
-          originalName: `${o}/m`,
-          variable: h(() => O([...r.shape, 2]).variable(p))
+          originalName: `${l}/m`,
+          variable: c(() => O([...r.shape, 2]).variable(p))
         });
-        const m = Array.isArray(t) ? t[i].tensor : t[o];
+        const m = Array.isArray(t) ? t[i].tensor : t[l];
         if (m == null)
           return;
-        const u = this.orthGrad ? this.orthogonalizeGradient(r, m) : m, d = this.accumulatedMoments[i].variable, g = S(d, u, this.beta1, this.beta2, n);
+        const u = this.orthGrad ? this.orthogonalizeGradient(r, m) : m, d = this.accumulatedMoments[i].variable, g = N(d, u, this.beta1, this.beta2, o);
         d.assign(g), this.orthGrad && u.dispose();
-        const y = N(
+        const y = B(
           g,
           r,
-          l,
-          c,
+          n,
+          h,
           this.epsilon ?? 1e-8,
           this.learningRate,
           // Only apply weight decay if the variable is multi-dimensional (e.g. weights, not biases)
           r.shape.length > 1 ? this.weightDecay : 0
         );
         r.assign(y);
-      }), this.accBeta1 = this.accBeta1 * this.beta1, this.accBeta2 = this.accBeta2 * this.beta2, n;
+      }), this.accBeta1 = this.accBeta1 * this.beta1, this.accBeta2 = this.accBeta2 * this.beta2, o;
     });
     return this.incrementIterations(), s;
   }
   dispose() {
-    this.accumulatedMoments != null && B(this.accumulatedMoments.map((t) => t.variable));
+    this.accumulatedMoments != null && w(this.accumulatedMoments.map((t) => t.variable));
   }
   async getWeights() {
     const t = [...this.accumulatedMoments];
@@ -82,7 +111,7 @@ class G extends R {
     );
   }
   async setWeights(t) {
-    t = await this.extractIterations(t), h(() => {
+    t = await this.extractIterations(t), c(() => {
       this.accBeta1 = Math.pow(this.beta1, this.iterations_ + 1), this.accBeta2 = Math.pow(this.beta2, this.iterations_ + 1);
     });
     const e = t.length / 2, a = !1;
@@ -105,5 +134,5 @@ class G extends R {
   }
 }
 export {
-  G as AdamWOptimizer
+  _ as AdamWOptimizer
 };

package/dist/training/BasicTrainer.d.ts

@@ -31,6 +31,7 @@ export default class BasicTrainer {
     get isRunning(): boolean;
     getOptimizer(): AdamWOptimizer;
     updateOptimizer(config?: Partial<AdamWOptimizerConfig>): void;
+    resumeFromLog(log: TrainingLogEntry): void;
     protected trainStep(state: Partial<TrainingState>, batch: {
         xs: Tensor;
         ys: Tensor;

package/dist/training/BasicTrainer.js

@@ -1,16 +1,16 @@
-import u from "./Evaluator.js";
-import { t as z, v as P, k as g, d as p, a as y } from "../index-CUXkjxiT.js";
-import S from "../utilities/profile.js";
-import { createTensorStatistics as k } from "../checks/weights.js";
-import { calculateLoss as x, calculateAccuracy as T } from "./loss.js";
-import { AdamWOptimizer as N } from "./AdamW.js";
-import { z as w } from "../zeros-DvZpK8s6.js";
-const v = {
+import y from "./Evaluator.js";
+import { t as L, Z as k, k as u, l as p, b as S } from "../index-DSGwv2Yx.js";
+import w from "../utilities/profile.js";
+import { createTensorStatistics as b } from "../checks/weights.js";
+import { calculateLoss as x, calculateAccuracy as P } from "./loss.js";
+import { AdamWOptimizer as T } from "./AdamW.js";
+import { z as v } from "../zeros-Bw0puq_w.js";
+const z = {
   logInterval: 1,
   maxEpochs: 100,
   sftMode: "full",
   batchSize: 32
-}, b = {
+}, D = {
   learningRate: 3e-4,
   beta1: 0.9,
   beta2: 0.99,
@@ -23,14 +23,14 @@ const v = {
   lossScaling: 1
 };
 class G {
-  constructor(s, i, o, c) {
-    this.tokenizer = i, this.model = s, this.optimizerConfig = {
-      ...b,
-      ...o,
+  constructor(s, e, n, l) {
+    this.tokenizer = e, this.model = s, this.optimizerConfig = {
+      ...D,
+      ...n,
       lossScaling: s.lossScaling
     };
-    const l = c || new N(this.optimizerConfig);
-    c && c.updateConfig(this.optimizerConfig), this.optimizer = l;
+    const m = l || new T(this.optimizerConfig);
+    l && l.updateConfig(this.optimizerConfig), this.optimizer = m;
   }
   model;
   optimizer;
@@ -80,11 +80,22 @@ class G {
   updateOptimizer(s) {
     s && (this.optimizerConfig = { ...this.optimizerConfig, ...s }), this.optimizer.updateConfig(this.optimizerConfig);
   }
+  resumeFromLog(s) {
+    (!this.lastState || this.lastState.step === 0) && (this.lastState = {
+      losses: [],
+      validationLosses: [],
+      logStartTime: 0,
+      step: s.step,
+      lastLoss: s.trainingMetrics.loss,
+      totalSteps: s.step,
+      trainingDuration: s.duration
+    });
+  }
   // A single forward pass, backward pass, and optimizer step
-  trainStep(s, i, o = !1, c = !1) {
-    return z(() => {
+  trainStep(s, e, n = !1, l = !1) {
+    return L(() => {
       this.model.getProfiler()?.startMemory();
-      const { xs: l, ys: a } = i, d = () => {
+      const { xs: m, ys: i } = e, d = () => {
         const r = this.model.forward(
           {
             training: !0,
@@ -93,32 +104,32 @@ class G {
             dropout: this._dropout,
             layerDrop: this._layerDrop
           },
-          l
-        ), e = x(r, a, this.maskedLoss, !1, this._labelSmoothing);
-        this.metrics.has("accuracy") && (s.accuracy = T(r, a), g(s.accuracy)), r.dispose();
-        const m = e.mul(y(this.optimizerConfig.lossScaling));
-        return e.dispose(), m;
-      }, { value: t, grads: n } = P(d);
-      if (o)
+          m
+        ), o = x(r, i, this.maskedLoss, !1, this._labelSmoothing);
+        this.metrics.has("accuracy") && (s.accuracy = P(r, i), u(s.accuracy)), r.dispose();
+        const a = o.mul(S(this.optimizerConfig.lossScaling));
+        return o.dispose(), a;
+      }, { value: t, grads: c } = k(d);
+      if (n)
         this.model.getProfiler()?.endMemory("Training");
       else {
-        const r = this.optimizer.applyGradients(n);
-        this.metrics.has("gradientNorm") ? (s.gradientNorm = r, g(r)) : (s.gradientNorm = void 0, r.dispose());
-        const e = Object.keys(n);
-        this.model.weightStore.touchVariables(e), this.model.getProfiler()?.endMemory("Training"), c ? (s.gradients = n, Object.values(n).forEach((m) => g(m))) : p(n);
+        const r = this.optimizer.applyGradients(c);
+        this.metrics.has("gradientNorm") ? (s.gradientNorm = r, u(r)) : (s.gradientNorm = void 0, r.dispose());
+        const o = Object.keys(c);
+        this.model.weightStore.touchVariables(o), this.model.getProfiler()?.endMemory("Training"), l ? (s.gradients = c, Object.values(c).forEach((a) => u(a))) : p(c);
       }
-      return t.mul(y(1 / this.optimizerConfig.lossScaling));
+      return t.mul(S(1 / this.optimizerConfig.lossScaling));
     });
   }
   async dummyPass() {
-    const s = w([1, this.model.config.blockSize], "int32"), i = w([1, this.model.config.blockSize], "int32");
+    const s = v([1, this.model.config.blockSize], "int32"), e = v([1, this.model.config.blockSize], "int32");
     try {
-      const o = this.trainStep({}, { xs: s, ys: i }, !0);
-      await o.data(), o.dispose();
-    } catch (o) {
-      console.error("Error during dummy pass:", o);
+      const n = this.trainStep({}, { xs: s, ys: e }, !0);
+      await n.data(), n.dispose();
+    } catch (n) {
+      console.error("Error during dummy pass:", n);
     } finally {
-      s.dispose(), i.dispose();
+      s.dispose(), e.dispose();
     }
   }
   dispose() {
@@ -136,33 +147,40 @@ class G {
       ...this.lastState || {}
     };
   }
-  async stepDataset(s, i, o) {
-    const { logInterval: c = 10 } = {
-      ...v,
-      ...i
+  async stepDataset(s, e, n) {
+    const { logInterval: l = 10 } = {
+      ...z,
+      ...e
     };
-    i.metrics && this.setMetrics(i.metrics);
-    const l = Date.now(), a = this.createEmptyState();
-    this.lastState = a, await this.dummyPass(), this.metrics.has("memoryUsage") && (this.model.getProfiler() || this.model.setProfiler(new S())), this.running = !0, a.logStartTime = l;
-    const d = o ? new u(this.model, o, void 0, this.maskedLoss) : void 0, t = await s.iterator();
+    e.metrics && this.setMetrics(e.metrics);
+    const m = Date.now(), i = this.createEmptyState();
+    this.lastState = i, await this.dummyPass(), this.metrics.has("memoryUsage") && (this.model.getProfiler() || this.model.setProfiler(new w())), this.running = !0, i.logStartTime = m;
+    const d = n ? new y(this.model, n, void 0, this.maskedLoss) : void 0, t = await s.iterator();
     try {
       for (; this.running; ) {
-        const n = await t.next();
-        if (n.done) break;
-        const r = n.value, e = this.trainStep(a, r, !1);
-        r.xs.dispose(), r.ys.dispose(), a.step++, a.totalSteps++, a.step % c === 0 ? await this.performLogging(e, r.xs.shape[0], i, d) : (a.gradientNorm && (a.gradientNorm.dispose(), a.gradientNorm = void 0), a.accuracy && (a.accuracy.dispose(), a.accuracy = void 0)), e.dispose();
+        const c = await t.next();
+        if (c.done) break;
+        const r = c.value, o = this.trainStep(i, r, !1);
+        r.xs.dispose(), r.ys.dispose(), i.step++, i.totalSteps++, i.step % l === 0 ? await this.performLogging(o, r.xs.shape[0], e, d) : (i.gradientNorm && (i.gradientNorm.dispose(), i.gradientNorm = void 0), i.accuracy && (i.accuracy.dispose(), i.accuracy = void 0)), o.dispose();
       }
-    } catch (n) {
-      throw console.error("Training error:", n), p(), n;
+    } catch (c) {
+      throw console.error("Training error:", c), c;
     }
-    throw p(), this.running = !1, new Error("No log returned before training stopped.");
-  }
-  async performLogging(s, i, o, c) {
-    const l = o?.onStep, a = this.metrics.has("gradientStatistics"), d = (await s.data())[0], t = this.lastState;
+    throw this.model.trainingState = {
+      steps: i.totalSteps,
+      learningRate: this.optimizer.lr,
+      batchSize: e.batchSize || 32,
+      loss: i.lastLoss,
+      tokensProcessed: i.totalSteps * (e.batchSize || 32) * this.model.config.blockSize,
+      duration: i.trainingDuration
+    }, p(), this.running = !1, new Error("No log returned before training stopped.");
+  }
+  async performLogging(s, e, n, l) {
+    const m = n?.onStep, i = this.metrics.has("gradientStatistics"), d = (await s.data())[0], t = this.lastState;
     t.lastLoss = d;
-    const n = Date.now();
-    t.trainingDuration += n - t.logStartTime;
-    const r = {
+    const c = Date.now();
+    t.trainingDuration += c - t.logStartTime;
+    const r = t.totalSteps * e * this.model.config.blockSize, o = {
       trainingMetrics: {
         loss: t.lastLoss,
         perplexity: this.metrics.has("perplexity") ? Math.exp(t.lastLoss) : void 0,
@@ -171,55 +189,57 @@ class G {
       step: t.step,
       time: Date.now() - t.logStartTime,
       gradientNorm: t.gradientNorm ? (await t.gradientNorm.data())[1] : void 0,
-      batchSize: i,
+      batchSize: e,
       learningRate: this.metrics.has("learningRate") ? this.optimizer.lr : void 0,
       duration: t.trainingDuration,
-      totalSamples: t.totalSteps * i,
-      samplesPerSecond: t.totalSteps * i / (t.trainingDuration / 1e3),
+      totalTokens: r,
+      tokensPerSecond: r / (t.trainingDuration / 1e3),
       memoryUsage: this.metrics.has("memoryUsage") ? this.model.getProfiler()?.getPeakMemory() || 0 : void 0
     };
-    if (this.metrics.has("tokensPerSecond") && (r.tokensPerSecond = r.samplesPerSecond * this.model.config.blockSize), t.gradientNorm && (t.gradientNorm.dispose(), t.gradientNorm = void 0), t.accuracy && (t.accuracy.dispose(), t.accuracy = void 0), this.model.trainingState = {
+    if (t.gradientNorm && (t.gradientNorm.dispose(), t.gradientNorm = void 0), t.accuracy && (t.accuracy.dispose(), t.accuracy = void 0), this.model.trainingState = {
       steps: t.totalSteps,
       learningRate: this.optimizer.lr,
-      batchSize: i,
-      loss: t.lastLoss
-    }, a && t.gradients) {
-      const e = /* @__PURE__ */ new Map();
-      for (const [m, h] of Object.entries(t.gradients))
-        e.set(m, await k(h)), h.dispose();
-      r.gradientMetrics = e;
+      batchSize: e,
+      loss: t.lastLoss,
+      tokensProcessed: r,
+      duration: t.trainingDuration
+    }, i && t.gradients) {
+      const a = /* @__PURE__ */ new Map();
+      for (const [h, g] of Object.entries(t.gradients))
+        a.set(h, await b(g)), g.dispose();
+      o.gradientMetrics = a;
     }
-    if (c)
+    if (l)
       try {
-        const e = await c.evaluate(5);
-        Array.isArray(e) ? r.validationMetrics = { loss: e[0].loss, accuracy: e[0].accuracy } : (t.validationLosses.push(e.loss), r.validationMetrics = {
-          accuracy: e.accuracy,
-          loss: e.loss,
-          perplexity: this.metrics.has("perplexity") ? Math.exp(e.loss) : void 0
+        const a = await l.evaluate(5);
+        Array.isArray(a) ? o.validationMetrics = { loss: a[0].loss, accuracy: a[0].accuracy } : (t.validationLosses.push(a.loss), o.validationMetrics = {
+          accuracy: a.accuracy,
+          loss: a.loss,
+          perplexity: this.metrics.has("perplexity") ? Math.exp(a.loss) : void 0
         });
-      } catch (e) {
-        console.error("Validation error:", e);
+      } catch (a) {
+        console.error("Validation error:", a);
       }
-    l && await l(r), t.logStartTime = Date.now();
-  }
-  async trainOnDataset(s, i, o) {
-    const { logInterval: c = 10, maxEpochs: l = 1 / 0 } = {
-      ...v,
-      ...i
-    }, a = l * (i?.epochSteps || 1e3);
-    i.metrics && this.setMetrics(i.metrics);
+    m && await m(o), t.logStartTime = Date.now();
+  }
+  async trainOnDataset(s, e, n) {
+    const { logInterval: l = 10, maxEpochs: m = 1 / 0 } = {
+      ...z,
+      ...e
+    }, i = m * (e?.epochSteps || 1e3);
+    e.metrics && this.setMetrics(e.metrics);
     const d = Date.now(), t = this.createEmptyState();
-    this.lastState = t, await this.dummyPass(), i?.metrics?.includes("memoryUsage") && (this.model.getProfiler() || this.model.setProfiler(new S())), this.running = !0, t.logStartTime = d;
-    const n = o ? new u(this.model, o, void 0, this.maskedLoss) : void 0, r = await s.iterator();
+    this.lastState = t, await this.dummyPass(), e?.metrics?.includes("memoryUsage") && (this.model.getProfiler() || this.model.setProfiler(new w())), this.running = !0, t.logStartTime = d;
+    const c = n ? new y(this.model, n, void 0, this.maskedLoss) : void 0, r = await s.iterator();
     try {
       for (; this.running; ) {
-        const e = await r.next();
-        if (e.done) break;
-        const m = e.value, h = t.step % c === 0, L = (i?.metrics?.includes("gradientStatistics") || !1) && h, f = this.trainStep(t, m, !1, L);
-        m.xs.dispose(), m.ys.dispose(), t.step++, t.totalSteps++, h ? await this.performLogging(f, m.xs.shape[0], i, n) : (t.gradientNorm && (t.gradientNorm.dispose(), t.gradientNorm = void 0), t.accuracy && (t.accuracy.dispose(), t.accuracy = void 0)), f.dispose(), t.step >= a && this.stop();
+        const o = await r.next();
+        if (o.done) break;
+        const a = o.value, h = t.step % l === 0, g = (e?.metrics?.includes("gradientStatistics") || !1) && h, f = this.trainStep(t, a, !1, g);
+        a.xs.dispose(), a.ys.dispose(), t.step++, t.totalSteps++, h ? await this.performLogging(f, a.xs.shape[0], e, c) : (t.gradientNorm && (t.gradientNorm.dispose(), t.gradientNorm = void 0), t.accuracy && (t.accuracy.dispose(), t.accuracy = void 0)), f.dispose(), t.step >= i && this.stop();
       }
-    } catch (e) {
-      throw console.error("Training error:", e), p(), e;
+    } catch (o) {
+      throw console.error("Training error:", o), p(), o;
     }
     return p(), this.running = !1, { losses: t.losses, validationLosses: t.validationLosses };
   }

package/dist/training/DatasetBuilder.js

@@ -1,6 +1,6 @@
-import { t as f } from "../index-CUXkjxiT.js";
-import "../dataset-CGGp1z9P.js";
-import { g as a } from "../readers-iz5u3HBo.js";
+import { t as f } from "../index-DSGwv2Yx.js";
+import "../dataset-DlqAN81i.js";
+import { g as a } from "../readers-17HLdxVM.js";
 import "../index-Cp39cXWe.js";
 const g = 8;
 async function p(n, e) {

package/dist/training/Evaluator.js

@@ -1,7 +1,7 @@
-import { t as p } from "../index-CUXkjxiT.js";
+import { t as p } from "../index-DSGwv2Yx.js";
 import { calculateLoss as d, calculateAccuracy as m } from "./loss.js";
 import { buildSFTExample as x } from "./SFTDatasetBuilder.js";
-import { t as h } from "../tensor-BWFldCso.js";
+import { t as h } from "../tensor-D8e0Gd7c.js";
 class k {
   constructor(i, t, o, c) {
     if (this.model = i, this.masked = !!c, Array.isArray(t)) {