@genai-fi/nanogpt 0.15.13 → 0.16.1

package/dist/tokeniser/bpe.js

@@ -1,6 +1,6 @@
 import { yieldIfNeeded as f } from "../utilities/yielder.js";
-import k from "../utilities/tokenParse.js";
-import z, { SPECIALS as m } from "./BaseTokeniser.js";
+import m from "../utilities/tokenParse.js";
+import z, { SPECIALS as k } from "./BaseTokeniser.js";
 function p(o, e) {
   return `${o}-::-${e}`;
 }
@@ -8,25 +8,25 @@ function w(o) {
   const e = /* @__PURE__ */ new Map();
   for (let s = 0; s < o.length; s++) {
     const t = o[s];
-    for (let r = 0; r < t.length - 1; r++) {
-      const n = p(t[r], t[r + 1]), a = e.get(n) || {
-        a: t[r],
-        b: t[r + 1],
+    for (let n = 0; n < t.length - 1; n++) {
+      const r = p(t[n], t[n + 1]), a = e.get(r) || {
+        a: t[n],
+        b: t[n + 1],
         count: 0,
         instances: /* @__PURE__ */ new Set()
       };
-      a.count += 1, a.instances.add(s), e.set(n, a);
+      a.count += 1, a.instances.add(s), e.set(r, a);
     }
   }
   return { pairs: e, tokens: o };
 }
-function d(o, e, s, t, r) {
-  const n = p(e, s);
-  if (o.pairs.has(n)) {
-    const a = o.pairs.get(n);
-    a.count += r, r > 0 ? a.instances.add(t) : a.count <= 0 ? o.pairs.delete(n) : a.instances.delete(t);
+function d(o, e, s, t, n) {
+  const r = p(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(n, { a: e, b: s, count: r, instances: /* @__PURE__ */ new Set([t]) });
+    o.pairs.set(r, { a: e, b: s, count: n, instances: /* @__PURE__ */ new Set([t]) });
 }
 function T(o) {
   let e = null, s = 0;
@@ -37,21 +37,21 @@ function T(o) {
 function y(o, e) {
   return o.map((s) => {
     const t = [];
-    for (let r = 0; r < s.length; r++)
-      r < s.length - 1 && s[r] === e[0] && s[r + 1] === e[1] ? (t.push(e[0] + e[1]), r++) : t.push(s[r]);
+    for (let n = 0; n < s.length; n++)
+      n < s.length - 1 && s[n] === e[0] && s[n + 1] === e[1] ? (t.push(e[0] + e[1]), n++) : t.push(s[n]);
     return t;
   });
 }
 function I(o, e) {
   e.instances.forEach((s) => {
-    const t = o.tokens[s], r = [];
-    for (let n = 0; n < t.length; n++)
-      if (n < t.length - 1 && t[n] === e.a && t[n + 1] === e.b) {
+    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;
-        r.push(a), n > 0 && (d(o, t[n - 1], e.a, s, -1), d(o, t[n - 1], a, s, 1)), n++, n < t.length - 1 && (d(o, e.b, t[n + 1], s, -1), d(o, a, t[n + 1], s, 1));
+        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));
       } else
-        r.push(t[n]);
-    o.tokens[s] = r;
+        n.push(t[r]);
+    o.tokens[s] = n;
   }), o.pairs.delete(p(e.a, e.b));
 }
 class E extends z {
@@ -61,11 +61,11 @@ class E extends z {
   merges = [];
   pretokenMap = /* @__PURE__ */ new Map();
   constructor(e, s) {
-    super(), Array.isArray(e) ? (e.forEach((t, r) => {
-      this.vocab.add(t), this.vocabIndex.set(t, r);
-    }), s && (this.merges = s), this.targetSize = e.length, m.forEach((t) => {
-      const r = e.indexOf(t);
-      r !== -1 && this.addSpecialToken(t, r);
+    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) => {
+      const n = e.indexOf(t);
+      n !== -1 && this.addSpecialToken(t, n);
     })) : (this.addSpecialTokens(), this.targetSize = e);
   }
   addToken(e, s) {
@@ -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 > m.length && this.vocab.size <= this.targetSize;
+    return this.vocab.size > k.length && this.vocab.size <= this.targetSize;
   }
   get vocabSize() {
     return this.vocab.size;
@@ -95,23 +95,23 @@ class E extends z {
   get unkToken() {
     return this.vocabIndex.get("") ?? 1;
   }
-  async train(e, s) {
+  async train(e = [], s) {
     let t = performance.now();
-    const r = e.map((i) => k(i)).flat(1);
+    const n = e.map((i) => i.map((h) => m(h.content))).flat(2);
     t = await f(t, s, this.vocab.size);
-    const n = new Set(r);
+    const r = new Set(n);
     this.vocab = /* @__PURE__ */ new Set(), this.pretokenMap.clear(), this.merges = [], this.addSpecialTokens();
-    const a = Array.from(n), b = a.map((i) => Array.from(i).map((h) => (this.vocab.add(h), h))), g = w(b);
+    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) {
       console.warn("Initial vocab size is greater than or equal to target size. No merges will be performed.");
       const i = /* @__PURE__ */ new Map();
-      r.forEach((c) => {
+      n.forEach((c) => {
         Array.from(c).forEach((u) => {
           i.set(u, (i.get(u) || 0) + 1);
         });
       });
-      const l = Array.from(i.entries()).sort((c, u) => u[1] - c[1]);
-      this.vocab = /* @__PURE__ */ new Set(), this.addSpecialTokens(), l.slice(0, this.targetSize - this.vocab.size).map(([c]) => c).forEach((c) => this.vocab.add(c)), this.vocabIndex.clear();
+      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;
       for (const c of this.vocab.keys())
         this.vocabIndex.set(c, S++);
@@ -123,9 +123,9 @@ class E extends z {
         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);
     }
-    a.forEach((i, l) => {
-      const h = b[l];
-      this.pretokenMap.set(i, h);
+    a.forEach((i, h) => {
+      const l = b[h];
+      this.pretokenMap.set(i, l);
     }), this.vocabIndex.clear();
     let v = 0;
     for (const i of this.vocab.keys())
@@ -145,15 +145,15 @@ class E extends z {
     }), this.pretokenMap.set(e, s), s;
   }
   tokeniseStrings(e) {
-    return e.map((s) => k(s).map((n) => this.pretokenMap.has(n) ? this.pretokenMap.get(n) : this.tokeniseWord(n)).flat(1));
+    return e.map((s) => m(s).map((r) => this.pretokenMap.has(r) ? this.pretokenMap.get(r) : this.tokeniseWord(r)).flat(1));
   }
   tokenise(e, s) {
     const t = this.tokeniseStrings(e);
-    return s ? t.map((r) => r.map((n) => this.vocabIndex.get(n) ?? this.unkToken)) : t.map((r) => r.map((n) => this.vocab.has(n) ? n : ""));
+    return s ? t.map((n) => n.map((r) => this.vocabIndex.get(r) ?? this.unkToken)) : t.map((n) => n.map((r) => this.vocab.has(r) ? r : ""));
   }
   detokenise(e) {
     const s = this.getVocab();
-    return e.map((r) => r.map((n) => s[n]).join(""));
+    return e.map((n) => n.map((r) => s[r]).join(""));
   }
   encode(e) {
     return this.tokenise([e], !0)[0];

package/dist/tokeniser/type.d.ts

@@ -1,11 +1,11 @@
 import { default as EE } from 'eventemitter3';
-export type Roles = 'user' | 'assistant' | 'system';
+export type Roles = 'user' | 'assistant' | 'system' | 'text';
 export interface Conversation {
     role: Roles;
     content: string;
 }
 export interface ITokeniser extends EE<'trainStatus'> {
-    train(text: string[], cb?: (vocab: number) => void): Promise<number>;
+    train(text: Conversation[][], cb?: (vocab: number) => void): Promise<number>;
     getVocab(): string[];
     getMerges(): [string, string][];
     destroy(): void;

package/dist/training/BasicTrainer.js

@@ -1,8 +1,8 @@
 import u from "./Evaluator.js";
-import { t as L, v as P, k as h, d as g, a as y } from "../index-CUXkjxiT.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 x } from "../checks/weights.js";
-import { calculateLoss as k, calculateAccuracy as T } from "./loss.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 = {
@@ -23,11 +23,11 @@ const v = {
   lossScaling: 1
 };
 class G {
-  constructor(e, i, o, c) {
-    this.tokenizer = i, this.model = e, this.optimizerConfig = {
+  constructor(s, i, o, c) {
+    this.tokenizer = i, this.model = s, this.optimizerConfig = {
       ...b,
       ...o,
-      lossScaling: e.lossScaling
+      lossScaling: s.lossScaling
     };
     const l = c || new N(this.optimizerConfig);
     c && c.updateConfig(this.optimizerConfig), this.optimizer = l;
@@ -44,26 +44,26 @@ class G {
   _labelSmoothing = 0;
   _layerDrop = 0;
   _dropout = 0;
-  setGradientCheckpointing(e) {
-    this._gradientCheckpointing = e;
+  setGradientCheckpointing(s) {
+    this._gradientCheckpointing = s;
   }
-  setMixedPrecision(e) {
-    this._mixedPrecision = e;
+  setMixedPrecision(s) {
+    this._mixedPrecision = s;
   }
-  setLabelSmoothing(e) {
-    this._labelSmoothing = e;
+  setLabelSmoothing(s) {
+    this._labelSmoothing = s;
   }
-  setDropout(e) {
-    this._dropout = e;
+  setDropout(s) {
+    this._dropout = s;
   }
-  setLayerDrop(e) {
-    this._layerDrop = e;
+  setLayerDrop(s) {
+    this._layerDrop = s;
   }
-  setLearningRate(e) {
-    this.optimizerConfig.learningRate = e, this.updateOptimizer();
+  setLearningRate(s) {
+    this.optimizerConfig.learningRate = s, this.updateOptimizer();
   }
-  setMetrics(e) {
-    this.metrics = new Set(e);
+  setMetrics(s) {
+    this.metrics = new Set(s);
   }
   reset() {
     this.lastState = void 0, this.running = !1;
@@ -77,12 +77,12 @@ class G {
   getOptimizer() {
     return this.optimizer;
   }
-  updateOptimizer(e) {
-    e && (this.optimizerConfig = { ...this.optimizerConfig, ...e }), this.optimizer.updateConfig(this.optimizerConfig);
+  updateOptimizer(s) {
+    s && (this.optimizerConfig = { ...this.optimizerConfig, ...s }), this.optimizer.updateConfig(this.optimizerConfig);
   }
   // A single forward pass, backward pass, and optimizer step
-  trainStep(e, i, o = !1, c = !1) {
-    return L(() => {
+  trainStep(s, i, o = !1, c = !1) {
+    return z(() => {
       this.model.getProfiler()?.startMemory();
       const { xs: l, ys: a } = i, d = () => {
         const r = this.model.forward(
@@ -94,31 +94,31 @@ class G {
             layerDrop: this._layerDrop
           },
           l
-        ), s = k(r, a, this.maskedLoss, !1, this._labelSmoothing);
-        this.metrics.has("accuracy") && (e.accuracy = T(r, a), h(e.accuracy)), r.dispose();
-        const m = s.mul(y(this.optimizerConfig.lossScaling));
-        return s.dispose(), m;
+        ), 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)
         this.model.getProfiler()?.endMemory("Training");
       else {
         const r = this.optimizer.applyGradients(n);
-        this.metrics.has("gradientNorm") ? (e.gradientNorm = r, h(r)) : (e.gradientNorm = void 0, r.dispose());
-        const s = Object.keys(n);
-        this.model.weightStore.touchVariables(s), this.model.getProfiler()?.endMemory("Training"), c ? (e.gradients = n, Object.values(n).forEach((m) => h(m))) : g(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);
       }
       return t.mul(y(1 / this.optimizerConfig.lossScaling));
     });
   }
   async dummyPass() {
-    const e = w([1, this.model.config.blockSize], "int32"), i = w([1, this.model.config.blockSize], "int32");
+    const s = w([1, this.model.config.blockSize], "int32"), i = w([1, this.model.config.blockSize], "int32");
     try {
-      const o = this.trainStep({}, { xs: e, ys: i }, !0);
+      const o = this.trainStep({}, { xs: s, ys: i }, !0);
       await o.data(), o.dispose();
     } catch (o) {
       console.error("Error during dummy pass:", o);
     } finally {
-      e.dispose(), i.dispose();
+      s.dispose(), i.dispose();
     }
   }
   dispose() {
@@ -136,7 +136,7 @@ class G {
       ...this.lastState || {}
     };
   }
-  async stepDataset(e, i, o) {
+  async stepDataset(s, i, o) {
     const { logInterval: c = 10 } = {
       ...v,
       ...i
@@ -144,21 +144,21 @@ class G {
     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, t = await e.iterator();
+    const d = o ? new u(this.model, o, 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, s = this.trainStep(a, r, !1);
-        r.xs.dispose(), r.ys.dispose(), a.step++, a.totalSteps++, a.step % c === 0 ? await this.performLogging(s, r.xs.shape[0], i, d) : (a.gradientNorm && (a.gradientNorm.dispose(), a.gradientNorm = void 0), a.accuracy && (a.accuracy.dispose(), a.accuracy = void 0)), s.dispose();
+        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();
       }
     } catch (n) {
-      throw console.error("Training error:", n), g(), n;
+      throw console.error("Training error:", n), p(), n;
     }
-    throw g(), this.running = !1, new Error("No log returned before training stopped.");
+    throw p(), this.running = !1, new Error("No log returned before training stopped.");
   }
-  async performLogging(e, i, o, c) {
-    const l = o?.onStep, a = this.metrics.has("gradientStatistics"), d = (await e.data())[0], t = this.lastState;
+  async performLogging(s, i, o, c) {
+    const l = o?.onStep, a = this.metrics.has("gradientStatistics"), d = (await s.data())[0], t = this.lastState;
     t.lastLoss = d;
     const n = Date.now();
     t.trainingDuration += n - t.logStartTime;
@@ -184,25 +184,25 @@ class G {
       batchSize: i,
       loss: t.lastLoss
     }, a && t.gradients) {
-      const s = /* @__PURE__ */ new Map();
-      for (const [m, p] of Object.entries(t.gradients))
-        s.set(m, await x(p)), p.dispose();
-      r.gradientMetrics = s;
+      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;
     }
     if (c)
       try {
-        const s = await c.evaluate(5);
-        Array.isArray(s) ? r.validationMetrics = { loss: s[0].loss, accuracy: s[0].accuracy } : (t.validationLosses.push(s.loss), r.validationMetrics = {
-          accuracy: s.accuracy,
-          loss: s.loss,
-          perplexity: this.metrics.has("perplexity") ? Math.exp(s.loss) : void 0
+        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
         });
-      } catch (s) {
-        console.error("Validation error:", s);
+      } catch (e) {
+        console.error("Validation error:", e);
       }
     l && await l(r), t.logStartTime = Date.now();
   }
-  async trainOnDataset(e, i, o) {
+  async trainOnDataset(s, i, o) {
     const { logInterval: c = 10, maxEpochs: l = 1 / 0 } = {
       ...v,
       ...i
@@ -210,18 +210,18 @@ class G {
     i.metrics && this.setMetrics(i.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, r = await e.iterator();
+    const n = o ? new u(this.model, o, void 0, this.maskedLoss) : void 0, r = await s.iterator();
     try {
       for (; this.running; ) {
-        const s = await r.next();
-        if (s.done) break;
-        const m = s.value, p = t.step % c === 0, z = (i?.metrics?.includes("gradientStatistics") || !1) && p, f = this.trainStep(t, m, !1, z);
-        m.xs.dispose(), m.ys.dispose(), t.step++, t.totalSteps++, p ? 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 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();
       }
-    } catch (s) {
-      throw console.error("Training error:", s), g(), s;
+    } catch (e) {
+      throw console.error("Training error:", e), p(), e;
     }
-    return g(), this.running = !1, { losses: t.losses, validationLosses: t.validationLosses };
+    return p(), this.running = !1, { losses: t.losses, validationLosses: t.validationLosses };
   }
 }
 export {

package/dist/training/Evaluator.d.ts

@@ -11,7 +11,8 @@ export default class Evaluator {
     private iterator?;
     private xs?;
     private ys?;
-    constructor(model: Model<ModelForwardAttributes>, dataset: Dataset<TensorContainer> | Conversation[][], tokeniser?: ITokeniser);
+    private masked;
+    constructor(model: Model<ModelForwardAttributes>, dataset: Dataset<TensorContainer> | Conversation[][], tokeniser?: ITokeniser, masked?: boolean);
     dispose(): void;
     private calculateBatchLoss;
     evaluate(maxBatches?: number): Promise<Result | Result[]>;

package/dist/training/Evaluator.js

@@ -2,12 +2,12 @@ import { t as p } from "../index-CUXkjxiT.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";
-class g {
-  constructor(c, t, o) {
-    if (this.model = c, Array.isArray(t)) {
+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, c.config.blockSize)).filter((s) => s !== null);
+      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));
@@ -17,32 +17,33 @@ class g {
   iterator;
   xs;
   ys;
+  masked = !1;
   dispose() {
     this.xs && this.xs.dispose(), this.ys && this.ys.dispose();
   }
-  async calculateBatchLoss(c, t, o, a) {
-    const [s, e] = p(() => {
-      const r = this.model.forward({ training: !1 }, c), y = d(r, t, a, o), f = m(r, t);
+  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 s.array(), u = await e.array(), i = n, l = u;
-    return e.dispose(), s.dispose(), Array.isArray(i) ? i.map((r) => ({ loss: r, accuracy: l })) : { loss: i, accuracy: l };
+    }), 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 evaluate(c = 100) {
-    let t = 0, o = 0, a = 0;
+  async evaluate(i = 100) {
+    let t = 0, o = 0, c = 0;
     if (this.iterator) {
-      const s = await this.iterator;
-      for (let e = 0; e < c; e++) {
-        const n = await s.next();
+      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: i, ys: l } = u, r = await this.calculateBatchLoss(i, l, !1, !1);
-        i.dispose(), l.dispose(), t += r.loss, o += r.accuracy, a++;
+        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++;
       }
-      return { loss: t / a, accuracy: o / a };
+      return { loss: t / c, accuracy: o / c };
     } else if (this.xs && this.ys)
       return this.calculateBatchLoss(this.xs, this.ys, !0, !0);
     throw new Error("No data available for evaluation");
   }
 }
 export {
-  g as default
+  k as default
 };

package/dist/training/SFTDatasetBuilder.js

@@ -1,50 +1,52 @@
-import { t as x } from "../index-CUXkjxiT.js";
+import { t as y } from "../index-CUXkjxiT.js";
 import "../dataset-CGGp1z9P.js";
 import { g as I } from "../readers-iz5u3HBo.js";
 import "../index-Cp39cXWe.js";
-function w(u, a, t, r) {
-  const s = [t.bosToken], n = [!1], f = {
+function w(p, o, t, l) {
+  const s = [t.bosToken], a = [!1], u = {
     user: t.getSpecialTokenIndex("<|user_start|>"),
     assistant: t.getSpecialTokenIndex("<|assistant_start|>"),
-    system: t.getSpecialTokenIndex("<|system_start|>")
-  }, i = {
+    system: t.getSpecialTokenIndex("<|system_start|>"),
+    text: void 0
+  }, c = {
     user: t.getSpecialTokenIndex("<|user_end|>"),
     assistant: t.getSpecialTokenIndex("<|assistant_end|>"),
-    system: t.getSpecialTokenIndex("<|system_end|>")
+    system: t.getSpecialTokenIndex("<|system_end|>"),
+    text: void 0
   };
-  for (const e of u) {
-    const c = f[e.role], h = i[e.role];
-    if (!c || !h)
+  for (const e of p) {
+    const r = u[e.role], h = c[e.role];
+    if (!r || !h)
       throw new Error(`Missing special tokens for role: ${e.role}`);
-    s.push(c), n.push(!1);
-    const m = e.role === "assistant", S = t.encode(e.content);
-    for (const T of S) {
+    s.push(r), a.push(!1);
+    const m = e.role === "assistant", x = t.encode(e.content);
+    for (const T of x) {
       s.push(T);
-      const y = t.isSpecialToken(T);
-      n.push(m && !y);
+      const S = t.isSpecialToken(T);
+      a.push(m && !S);
     }
-    s.push(h), n.push(m);
+    s.push(h), a.push(m);
   }
-  s.push(t.eosToken), n.push(!1);
-  const o = r + 1;
-  if (s.length < o) {
-    const e = o - s.length, c = t.getSpecialTokenIndex("<pad>");
+  s.push(t.eosToken), a.push(!1);
+  const n = l + 1;
+  if (s.length < n) {
+    const e = n - s.length, r = t.getSpecialTokenIndex("<pad>");
     for (let h = 0; h < e; h++)
-      s.push(c), n.push(!1);
-  } else s.length > o && (s.length = o, n.length = o);
-  const p = new Int32Array(s.slice(0, r)), l = s.slice(1, r + 1), k = n.slice(1, r + 1), d = new Int32Array(l.length);
+      s.push(r), a.push(!1);
+  } else s.length > n && (s.length = n, a.length = n);
+  const f = new Int32Array(s.slice(0, l)), i = s.slice(1, l + 1), k = a.slice(1, l + 1), d = new Int32Array(i.length);
   let g = !1;
-  for (let e = 0; e < l.length; e++) {
-    const c = k[e] ? l[e] : a;
-    d[e] = c, c !== a && (g = !0);
+  for (let e = 0; e < i.length; e++) {
+    const r = k[e] ? i[e] : o;
+    d[e] = r, r !== o && (g = !0);
   }
-  return g ? { xs: p, ys: d } : null;
+  return g ? { xs: f, ys: d } : null;
 }
 class A {
   tokenizer;
   blockSize;
-  constructor(a, t = 128) {
-    this.tokenizer = a, this.blockSize = t;
+  constructor(o, t = 128) {
+    this.tokenizer = o, this.blockSize = t;
   }
   /**
    * Create SFT dataset from structured conversations.
@@ -52,20 +54,27 @@ class A {
    * - Pads with eosToken and masks padding.
    * - Masks non-assistant tokens in labels with ignoreIndex (default -100).
    */
-  async createSFTDataset(a, t = 32, r = -100) {
-    if (!a.length)
+  async createSFTDataset(o, t = 32, l = -100) {
+    if (!o.length)
       throw new Error("No conversations provided.");
-    const s = this.tokenizer, n = this.blockSize;
+    const s = this.tokenizer, a = this.blockSize;
+    for (const c of o)
+      c.shuffle();
     return I(function* () {
       for (; ; ) {
-        const i = Math.floor(Math.random() * a.length), p = a[i].getRandomConversation(), l = w(p, r, s, n);
-        l && (yield l);
+        const c = Math.floor(Math.random() * o.length), n = o[c], f = n.nextConversation();
+        if (!f) {
+          n.shuffle();
+          continue;
+        }
+        const i = w(f, l, s, a);
+        i && (yield i);
       }
-    }).batch(t).map((i) => {
-      const o = i;
-      return x(() => ({
-        xs: o.xs.cast("int32"),
-        ys: o.ys.cast("int32")
+    }).batch(t).map((c) => {
+      const n = c;
+      return y(() => ({
+        xs: n.xs.cast("int32"),
+        ys: n.ys.cast("int32")
       }));
     }).prefetch(2);
   }

package/dist/training/tasks/ConversationTask.d.ts

@@ -2,13 +2,13 @@ import { Conversation, ITokeniser } from '../../main';
 import { Task } from './Task';
 export default class ConversationTask extends Task {
     private rawConvo;
+    private shuffledIndices;
     private index;
     get length(): number;
     constructor(conversations: Conversation[][]);
     hasMoreConversations(): boolean;
     nextConversation(): Conversation[] | null;
     nextTokens(tokeniser: ITokeniser): number[] | null;
-    getRandomConversation(): Conversation[];
-    getRandomTokens(tokeniser: ITokeniser): number[];
+    shuffle(): void;
     estimateTokens(tokeniser: ITokeniser): Promise<number>;
 }