@genai-fi/nanogpt 0.18.0 → 0.19.0

package/dist/training/BasicTrainer.js

@@ -1,10 +1,10 @@
-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 S from "./Evaluator.js";
+import { t as k, v as x, k as y, d as u, a as w } from "../index-CUXkjxiT.js";
+import v from "../utilities/profile.js";
+import { createTensorStatistics as N } from "../checks/weights.js";
+import { calculateLoss as b, calculateAccuracy as P } from "./loss.js";
 import { AdamWOptimizer as T } from "./AdamW.js";
-import { z as v } from "../zeros-Bw0puq_w.js";
+import { z as L } from "../zeros-DvZpK8s6.js";
 const z = {
   logInterval: 1,
   maxEpochs: 100,
@@ -22,15 +22,15 @@ const z = {
   minLearningRate: 1e-5,
   lossScaling: 1
 };
-class G {
-  constructor(s, e, n, l) {
-    this.tokenizer = e, this.model = s, this.optimizerConfig = {
+class B {
+  constructor(s, i, n, l) {
+    this.tokenizer = i, this.model = s, this.optimizerConfig = {
       ...D,
       ...n,
       lossScaling: s.lossScaling
     };
-    const m = l || new T(this.optimizerConfig);
-    l && l.updateConfig(this.optimizerConfig), this.optimizer = m;
+    const d = l || new T(this.optimizerConfig);
+    l && l.updateConfig(this.optimizerConfig), this.optimizer = d;
   }
   model;
   optimizer;
@@ -44,6 +44,9 @@ class G {
   _labelSmoothing = 0;
   _layerDrop = 0;
   _dropout = 0;
+  setLossMasking() {
+    this.maskedLoss = !0;
+  }
   setGradientCheckpointing(s) {
     this._gradientCheckpointing = s;
   }
@@ -92,44 +95,45 @@ class G {
     });
   }
   // A single forward pass, backward pass, and optimizer step
-  trainStep(s, e, n = !1, l = !1) {
-    return L(() => {
+  trainStep(s, i, n = !1, l = !1) {
+    return k(() => {
       this.model.getProfiler()?.startMemory();
-      const { xs: m, ys: i } = e, d = () => {
-        const r = this.model.forward(
+      const { xs: d, ys: r } = i, m = () => {
+        const a = this.model.forward(
           {
             training: !0,
             checkpointing: this._gradientCheckpointing,
             mixedPrecision: this._mixedPrecision,
             dropout: this._dropout,
-            layerDrop: this._layerDrop
+            layerDrop: this._layerDrop,
+            ropePositionOffset: 0
           },
-          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);
+          d
+        ), o = b(a, r, this.maskedLoss, !1, this._labelSmoothing);
+        this.metrics.has("accuracy") && (s.accuracy = P(a, r), y(s.accuracy)), a.dispose();
+        const e = o.mul(w(this.optimizerConfig.lossScaling));
+        return o.dispose(), e;
+      }, { value: t, grads: c } = x(m);
       if (n)
         this.model.getProfiler()?.endMemory("Training");
       else {
-        const r = this.optimizer.applyGradients(c);
-        this.metrics.has("gradientNorm") ? (s.gradientNorm = r, u(r)) : (s.gradientNorm = void 0, r.dispose());
+        const a = this.optimizer.applyGradients(c);
+        this.metrics.has("gradientNorm") ? (s.gradientNorm = a, y(a)) : (s.gradientNorm = void 0, a.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);
+        this.model.weightStore.touchVariables(o), this.model.getProfiler()?.endMemory("Training"), l ? (s.gradients = c, Object.values(c).forEach((e) => y(e))) : u(c);
       }
-      return t.mul(S(1 / this.optimizerConfig.lossScaling));
+      return t.mul(w(1 / this.optimizerConfig.lossScaling));
     });
   }
   async dummyPass() {
-    const s = v([1, this.model.config.blockSize], "int32"), e = v([1, this.model.config.blockSize], "int32");
+    const s = L([1, this.model.config.blockSize], "int32"), i = L([1, this.model.config.blockSize], "int32");
     try {
-      const n = this.trainStep({}, { xs: s, ys: e }, !0);
+      const n = this.trainStep({}, { xs: s, ys: i }, !0);
       await n.data(), n.dispose();
     } catch (n) {
       console.error("Error during dummy pass:", n);
     } finally {
-      s.dispose(), e.dispose();
+      s.dispose(), i.dispose();
     }
   }
   dispose() {
@@ -147,40 +151,46 @@ class G {
       ...this.lastState || {}
     };
   }
-  async stepDataset(s, e, n) {
+  async stepDataset(s, i, n) {
     const { logInterval: l = 10 } = {
       ...z,
-      ...e
+      ...i
     };
-    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();
+    i.metrics && this.setMetrics(i.metrics);
+    const d = Date.now(), r = this.createEmptyState();
+    this.lastState = r, await this.dummyPass(), this.metrics.has("memoryUsage") && (this.model.getProfiler() || this.model.setProfiler(new v())), this.running = !0, r.logStartTime = d;
+    const m = n ? new S(this.model, n, this.maskedLoss) : void 0, t = await s.iterator();
     try {
       for (; this.running; ) {
         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();
+        const a = c.value, o = this.trainStep(r, a, !1);
+        if (i.debug) {
+          const e = (await o.data())[0];
+          if (isNaN(e) || !isFinite(e))
+            throw console.error("Invalid loss value:", e), console.error("Batch xs:", a.xs.toString()), console.error("Batch ys:", a.ys.toString()), console.error("State:", r), new Error("Loss is NaN or Infinity");
+          console.log(`Step ${r.step}: Loss = ${e}`);
+        }
+        a.xs.dispose(), a.ys.dispose(), r.step++, r.totalSteps++, r.step % l === 0 ? await this.performLogging(o, a.xs.shape[0], i, m) : (r.gradientNorm && (r.gradientNorm.dispose(), r.gradientNorm = void 0), r.accuracy && (r.accuracy.dispose(), r.accuracy = void 0)), o.dispose();
       }
     } catch (c) {
       throw console.error("Training error:", c), c;
     }
     throw this.model.trainingState = {
-      steps: i.totalSteps,
+      steps: r.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;
+      batchSize: i.batchSize || 32,
+      loss: r.lastLoss,
+      tokensProcessed: r.totalSteps * (i.batchSize || 32) * this.model.config.blockSize,
+      duration: r.trainingDuration
+    }, u(), this.running = !1, new Error("No log returned before training stopped.");
+  }
+  async performLogging(s, i, n, l) {
+    const d = n?.onStep, r = this.metrics.has("gradientStatistics"), m = (await s.data())[0], t = this.lastState;
+    t.lastLoss = m;
     const c = Date.now();
     t.trainingDuration += c - t.logStartTime;
-    const r = t.totalSteps * e * this.model.config.blockSize, o = {
+    const a = t.totalSteps * i * this.model.config.blockSize, o = {
       trainingMetrics: {
         loss: t.lastLoss,
         perplexity: this.metrics.has("perplexity") ? Math.exp(t.lastLoss) : void 0,
@@ -189,61 +199,67 @@ class G {
       step: t.step,
       time: Date.now() - t.logStartTime,
       gradientNorm: t.gradientNorm ? (await t.gradientNorm.data())[1] : void 0,
-      batchSize: e,
+      batchSize: i,
       learningRate: this.metrics.has("learningRate") ? this.optimizer.lr : void 0,
       duration: t.trainingDuration,
-      totalTokens: r,
-      tokensPerSecond: r / (t.trainingDuration / 1e3),
+      totalTokens: a,
+      tokensPerSecond: a / (t.trainingDuration / 1e3),
       memoryUsage: this.metrics.has("memoryUsage") ? this.model.getProfiler()?.getPeakMemory() || 0 : void 0
     };
     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: e,
+      batchSize: i,
       loss: t.lastLoss,
-      tokensProcessed: r,
+      tokensProcessed: a,
       duration: t.trainingDuration
-    }, i && t.gradients) {
-      const a = /* @__PURE__ */ new Map();
+    }, r && t.gradients) {
+      const e = /* @__PURE__ */ new Map();
       for (const [h, g] of Object.entries(t.gradients))
-        a.set(h, await b(g)), g.dispose();
-      o.gradientMetrics = a;
+        e.set(h, await N(g)), g.dispose();
+      o.gradientMetrics = e;
     }
     if (l)
       try {
-        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
+        const e = await l.evaluate(5);
+        Array.isArray(e) ? o.validationMetrics = { loss: e[0].loss, accuracy: e[0].accuracy } : (t.validationLosses.push(e.loss), o.validationMetrics = {
+          accuracy: e.accuracy,
+          loss: e.loss,
+          perplexity: this.metrics.has("perplexity") ? Math.exp(e.loss) : void 0
         });
-      } catch (a) {
-        console.error("Validation error:", a);
+      } catch (e) {
+        console.error("Validation error:", e);
       }
-    m && await m(o), t.logStartTime = Date.now();
+    d && await d(o), t.logStartTime = Date.now();
   }
-  async trainOnDataset(s, e, n) {
-    const { logInterval: l = 10, maxEpochs: m = 1 / 0 } = {
+  async trainOnDataset(s, i, n) {
+    const { logInterval: l = 10, maxEpochs: d = 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(), 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();
+      ...i
+    }, r = d * (i?.epochSteps || 1e3);
+    i.metrics && this.setMetrics(i.metrics);
+    const m = Date.now(), t = this.createEmptyState();
+    this.lastState = t, await this.dummyPass(), i?.metrics?.includes("memoryUsage") && (this.model.getProfiler() || this.model.setProfiler(new v())), this.running = !0, t.logStartTime = m;
+    const c = n ? new S(this.model, n, this.maskedLoss) : void 0, a = await s.iterator();
     try {
       for (; this.running; ) {
-        const o = await r.next();
+        const o = await a.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();
+        const e = o.value, h = t.step % l === 0, g = (i?.metrics?.includes("gradientStatistics") || !1) && h, f = this.trainStep(t, e, !1, g);
+        if (i.debug) {
+          const p = (await f.data())[0];
+          if (isNaN(p) || !isFinite(p))
+            throw console.error("Invalid loss value:", p), console.error("Batch xs:", await e.xs.array()), console.error("Batch ys:", await e.ys.array()), console.error("State:", t), new Error("Loss is NaN or Infinity");
+          console.log(`Step ${t.step}: Loss = ${p}`);
+        }
+        e.xs.dispose(), e.ys.dispose(), t.step++, t.totalSteps++, h ? await this.performLogging(f, e.xs.shape[0], i, c) : (t.gradientNorm && (t.gradientNorm.dispose(), t.gradientNorm = void 0), t.accuracy && (t.accuracy.dispose(), t.accuracy = void 0)), f.dispose(), t.step >= r && this.stop();
       }
     } catch (o) {
-      throw console.error("Training error:", o), p(), o;
+      throw console.error("Training error:", o), u(), o;
     }
-    return p(), this.running = !1, { losses: t.losses, validationLosses: t.validationLosses };
+    return u(), this.running = !1, { losses: t.losses, validationLosses: t.validationLosses };
   }
 }
 export {
-  G as default
+  B as default
 };

package/dist/training/DatasetBuilder.d.ts

@@ -2,7 +2,11 @@ import { Tensor } from '@tensorflow/tfjs-core';
 import { Conversation, ITokeniser } from '../tokeniser/type';
 import { Dataset } from '@tensorflow/tfjs-data';
 export declare const PAGE_FACTOR = 8;
-export declare function flattenTokens(textData: Conversation[][], tokenizer: ITokeniser): Promise<Uint16Array>;
+export declare function flattenTokens(textData: Conversation[][], tokenizer: ITokeniser): Uint16Array;
+export declare function flattenTokensWithMask(textData: Conversation[][], tokenizer: ITokeniser): {
+    tokens: Uint16Array;
+    mask: Uint8Array;
+};
 export declare function shuffle(array: Uint32Array): Uint32Array;
 export interface DatasetState {
     shuffledIndexes: Uint32Array;
@@ -12,7 +16,7 @@ export declare class DatasetBuilder {
     tokenizer: ITokeniser;
     blockSize: number;
     constructor(tokenizer: ITokeniser, blockSize?: number);
-    createTextDataset(flatTokens: Uint16Array, batchSize?: number, indexes?: Uint32Array): Promise<{
+    createTextDataset(flatTokens: Uint16Array, batchSize?: number, indexes?: Uint32Array, mask?: Uint8Array, ignoreIndex?: number): Promise<{
         dataset: Dataset<{
             xs: Tensor;
             ys: Tensor;

package/dist/training/DatasetBuilder.js

@@ -1,67 +1,86 @@
-import { t as f } from "../index-DSGwv2Yx.js";
-import "../dataset-DlqAN81i.js";
-import { g as a } from "../readers-17HLdxVM.js";
+import { t as x } from "../index-CUXkjxiT.js";
+import { d as g, i as m } from "../dataset-CGGp1z9P.js";
 import "../index-Cp39cXWe.js";
-const g = 8;
-async function p(n, e) {
-  const o = (await Promise.all(n.map((t) => e.encodeConversation(t)))).flat();
-  return new Uint16Array(o);
+function p(e) {
+  return g(async () => {
+    const t = await e();
+    return m(() => t.next());
+  });
 }
-function c(n) {
-  for (let e = n.length - 1; e > 0; e--) {
-    const i = Math.floor(Math.random() * (e + 1));
-    [n[e], n[i]] = [n[i], n[e]];
+const I = 8;
+function z(e, t) {
+  const r = e.map((c) => t.encodeConversation(c)).flat();
+  return new Uint16Array(r);
+}
+function A(e, t) {
+  const s = e.map((i) => t.encodeConversation(i, !1, !0));
+  console.log("Tokenised Texts with Mask:", s);
+  const r = s.map((i) => i.tokens).flat(), c = s.map((i) => i.mask).flat();
+  return { tokens: new Uint16Array(r), mask: new Uint8Array(c.map((i) => i ? 1 : 0)) };
+}
+function u(e) {
+  for (let t = e.length - 1; t > 0; t--) {
+    const s = Math.floor(Math.random() * (t + 1));
+    [e[t], e[s]] = [e[s], e[t]];
   }
-  return n;
+  return e;
 }
-class m {
+class S {
   tokenizer;
   blockSize;
-  constructor(e, i = 128) {
-    this.tokenizer = e, this.blockSize = i;
+  constructor(t, s = 128) {
+    this.tokenizer = t, this.blockSize = s;
   }
   // Create dataset from text files
-  async createTextDataset(e, i = 32, o) {
-    if (e.length < this.blockSize + 1)
-      throw new Error(`Not enough tokens (${e.length}) for block size ${this.blockSize}`);
-    const t = {
-      shuffledIndexes: new Uint32Array(e.length),
+  async createTextDataset(t, s = 32, r, c, i = 65535) {
+    if (t.length < this.blockSize + 1)
+      throw new Error(`Not enough tokens (${t.length}) for block size ${this.blockSize}`);
+    const o = {
+      shuffledIndexes: new Uint32Array(t.length),
       step: 0
     };
-    if (o)
-      t.shuffledIndexes = o;
+    if (r)
+      o.shuffledIndexes = r;
     else {
-      t.shuffledIndexes = new Uint32Array(e.length);
-      for (let s = 0; s < e.length; s++)
-        t.shuffledIndexes[s] = s;
-      c(t.shuffledIndexes);
+      o.shuffledIndexes = new Uint32Array(t.length);
+      for (let n = 0; n < t.length; n++)
+        o.shuffledIndexes[n] = n;
+      u(o.shuffledIndexes);
     }
-    const h = (function* () {
+    const d = (function* () {
       for (; ; ) {
-        const s = t.shuffledIndexes[t.step++];
-        if (t.step >= t.shuffledIndexes.length && (t.step = 0, c(t.shuffledIndexes)), s + this.blockSize + 1 > e.length)
+        const n = o.shuffledIndexes[o.step++];
+        if (o.step >= o.shuffledIndexes.length && (o.step = 0, u(o.shuffledIndexes)), n + this.blockSize + 1 > t.length)
           continue;
-        const r = new Int32Array(e.subarray(s, s + this.blockSize)), l = new Int32Array(e.subarray(s + 1, s + this.blockSize + 1));
-        yield { xs: r, ys: l };
+        const a = new Int32Array(t.subarray(n, n + this.blockSize)), k = t.subarray(n + 1, n + this.blockSize + 1), l = new Int32Array(k);
+        if (c) {
+          let h = 0;
+          for (let f = 0; f < l.length; f++)
+            c[n + 1 + f] === 0 && (l[f] = i, h++);
+          if (h === l.length)
+            continue;
+        }
+        yield { xs: a, ys: l };
       }
     }).bind(this);
     return {
-      dataset: a(h).batch(i).map((s) => {
-        const r = s;
-        return f(() => ({
-          xs: r.xs.cast("int32"),
-          ys: r.ys.cast("int32")
+      dataset: p(d).batch(s).map((n) => {
+        const a = n;
+        return x(() => ({
+          xs: a.xs.cast("int32"),
+          ys: a.ys.cast("int32")
           // this.tf.oneHot(batchData.ys.cast('int32'), this.tokenizer.vocabSize),
         }));
       }).prefetch(2),
       // Smaller prefetch to reduce memory pressure
-      state: t
+      state: o
     };
   }
 }
 export {
-  m as DatasetBuilder,
-  g as PAGE_FACTOR,
-  p as flattenTokens,
-  c as shuffle
+  S as DatasetBuilder,
+  I as PAGE_FACTOR,
+  z as flattenTokens,
+  A as flattenTokensWithMask,
+  u as shuffle
 };

package/dist/training/Evaluator.d.ts

@@ -1,7 +1,6 @@
 import { Dataset } from '@tensorflow/tfjs-data';
 import { TensorContainer } from '@tensorflow/tfjs-core';
 import { default as Model, ModelForwardAttributes } from '../models/model';
-import { Conversation, ITokeniser } from '../main';
 interface Result {
     loss: number;
     accuracy: number;
@@ -12,7 +11,7 @@ export default class Evaluator {
     private xs?;
     private ys?;
     private masked;
-    constructor(model: Model<ModelForwardAttributes>, dataset: Dataset<TensorContainer> | Conversation[][], tokeniser?: ITokeniser, masked?: boolean);
+    constructor(model: Model<ModelForwardAttributes>, dataset: Dataset<TensorContainer>, masked?: boolean);
     dispose(): void;
     private calculateBatchLoss;
     evaluate(maxBatches?: number): Promise<Result | Result[]>;

package/dist/training/Evaluator.js

@@ -1,18 +1,8 @@
-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-D8e0Gd7c.js";
-class k {
-  constructor(i, t, o, c) {
-    if (this.model = i, this.masked = !!c, Array.isArray(t)) {
-      if (!o)
-        throw new Error("Tokeniser is required when dataset is an array of conversations");
-      const a = t.map((s) => x(s, -100, o, i.config.blockSize)).filter((s) => s !== null);
-      if (a.length === 0)
-        return;
-      this.xs = h(a.map((s) => s.xs)), this.ys = h(a.map((s) => s.ys));
-    } else
-      this.iterator = t.iterator();
+import { t as d } from "../index-CUXkjxiT.js";
+import { calculateLoss as f, calculateAccuracy as p } from "./loss.js";
+class b {
+  constructor(o, t, a) {
+    this.model = o, this.masked = !!a, this.iterator = t.iterator();
   }
   iterator;
   xs;
@@ -21,29 +11,29 @@ class k {
   dispose() {
     this.xs && this.xs.dispose(), this.ys && this.ys.dispose();
   }
-  async calculateBatchLoss(i, t, o, c) {
-    const [a, s] = p(() => {
-      const r = this.model.forward({ training: !1 }, i), y = d(r, t, c, o), f = m(r, t);
-      return r.dispose(), [y, f];
-    }), n = await a.array(), u = await s.array(), e = n, l = u;
-    return s.dispose(), a.dispose(), Array.isArray(e) ? e.map((r) => ({ loss: r, accuracy: l })) : { loss: e, accuracy: l };
+  async calculateBatchLoss(o, t, a, r) {
+    const [l, e] = d(() => {
+      const s = this.model.forward({ training: !1 }, o), h = f(s, t, r, a), y = p(s, t);
+      return s.dispose(), [h, y];
+    }), u = await l.array(), n = await e.array(), c = u, i = n;
+    return e.dispose(), l.dispose(), Array.isArray(c) ? c.map((s) => ({ loss: s, accuracy: i })) : { loss: c, accuracy: i };
   }
-  async evaluate(i = 100) {
-    let t = 0, o = 0, c = 0;
+  async evaluate(o = 100) {
+    let t = 0, a = 0, r = 0;
     if (this.iterator) {
-      const a = await this.iterator;
-      for (let s = 0; s < i; s++) {
-        const n = await a.next();
-        if (n.done) break;
-        const u = n.value, { xs: e, ys: l } = u, r = await this.calculateBatchLoss(e, l, !1, this.masked);
-        e.dispose(), l.dispose(), t += r.loss, o += r.accuracy, c++;
+      const l = await this.iterator;
+      for (let e = 0; e < o; e++) {
+        const u = await l.next();
+        if (u.done) break;
+        const n = u.value, { xs: c, ys: i } = n, s = await this.calculateBatchLoss(c, i, !1, this.masked);
+        c.dispose(), i.dispose(), t += s.loss, a += s.accuracy, r++;
       }
-      return { loss: t / c, accuracy: o / c };
+      return { loss: t / r, accuracy: a / r };
     } else if (this.xs && this.ys)
       return this.calculateBatchLoss(this.xs, this.ys, !0, !0);
     throw new Error("No data available for evaluation");
   }
 }
 export {
-  k as default
+  b as default
 };

package/dist/training/SFTTrainer.d.ts

@@ -1,11 +1,12 @@
 import { default as Model, ModelForwardAttributes } from '../models/model';
 import { default as BasicTrainer } from './BasicTrainer';
 import { ITokeniser } from '../tokeniser/type';
-import { SFTDatasetBuilder } from './SFTDatasetBuilder';
 import { AdamWOptimizer } from './AdamW';
 import { AdamWOptimizerConfig } from './types';
+import { DatasetBuilder } from './DatasetBuilder';
 export default class SFTTrainer extends BasicTrainer {
     tokenizer: ITokeniser;
-    datasetBuilder: SFTDatasetBuilder;
+    datasetBuilder: DatasetBuilder;
+    loraName?: string;
     constructor(model: Model<ModelForwardAttributes>, tokenizer: ITokeniser, optConfig?: Partial<AdamWOptimizerConfig>, optimizer?: AdamWOptimizer);
 }

package/dist/training/SFTTrainer.js

@@ -1,5 +1,5 @@
 import r from "./BasicTrainer.js";
-import { SFTDatasetBuilder as s } from "./SFTDatasetBuilder.js";
+import { DatasetBuilder as s } from "./DatasetBuilder.js";
 const n = {
   decayEpochs: 100,
   epochSteps: 1e4,
@@ -10,12 +10,13 @@ const n = {
   learningRate: 3e-4
   // clipNorm: 1.0,
 };
-class m extends r {
+class c extends r {
   constructor(t, e, i, a) {
     super(t, e, { ...n, ...i }, a), this.tokenizer = e, this.optimizerConfig.minLearningRate = i?.minLearningRate ?? this.optimizerConfig.learningRate / 20, this.updateOptimizer(), this.datasetBuilder = new s(e, t.config.blockSize), this.maskedLoss = !0;
   }
   datasetBuilder;
+  loraName;
 }
 export {
-  m as default
+  c as default
 };

package/dist/training/orthoGrad.js

@@ -1,4 +1,4 @@
-import { t as d } from "../index-DSGwv2Yx.js";
+import { t as d } from "../index-CUXkjxiT.js";
 function l(m, t, s) {
   return d(() => {
     const r = m.reshape([-1]), o = t.reshape([-1]), e = r.mul(r).sum().add(s), c = r.mul(o).sum().div(e), n = o.sub(r.mul(c)), a = o.norm(), u = n.norm().add(s);