@genai-fi/nanogpt 0.7.1 → 0.7.3

package/dist/training/DatasetBuilder.js

@@ -1,5 +1,5 @@
-import { t as u } from "../index-UdZhlibC.js";
-import { d as z, i as f } from "../dataset-U3PrjwgU.js";
+import { t as g } from "../index-BoWRt-10.js";
+import { d as u, i as d } from "../dataset-CtdBYwjo.js";
 import "../index-Tf7vU29b.js";
 /**
  * @license
@@ -18,57 +18,60 @@ import "../index-Tf7vU29b.js";
  *
  * =============================================================================
  */
-function S(c) {
-  return z(async () => {
-    const t = await c();
-    return f(() => t.next());
+function z(r) {
+  return u(async () => {
+    const t = await r();
+    return d(() => t.next());
   });
 }
-const p = 8;
-async function y(c, t) {
-  const s = await Promise.all(c.map((n) => t.encode(n))), i = t.eosToken >= 0;
-  return s.map((n) => i ? [...n, t.eosToken] : n).flat();
+const S = 8;
+async function y(r, t) {
+  const s = await Promise.all(r.map((e) => t.encode(e))), o = t.eosToken >= 0, a = s.map((e) => o ? [...e, t.eosToken] : e).flat();
+  for (const e of a)
+    if (e < 0 || e >= t.vocabSize)
+      throw new Error(`Invalid token index ${e} found in tokenised data`);
+  return a;
 }
 class w {
   tokenizer;
   blockSize;
   pageSize;
   constructor(t, s = 128) {
-    this.tokenizer = t, this.blockSize = s, this.pageSize = s * p;
+    this.tokenizer = t, this.blockSize = s, this.pageSize = s * S;
   }
   // Create dataset from text files
-  async createTextDataset(t, s = 32, i, r) {
+  async createTextDataset(t, s = 32, o, a) {
     if (t.length < this.blockSize + 1)
       throw new Error(`Not enough tokens (${t.length}) for block size ${this.blockSize}`);
-    if (i && i.size > t.length / this.pageSize / 2)
+    if (o && o.size > t.length / this.pageSize / 2)
       throw new Error("Too many masked pages - would leave insufficient training data");
-    const n = (function* () {
-      if (i && r) {
-        const e = Array.from(i);
+    const e = (function* () {
+      if (o && a) {
+        const i = Array.from(o);
         for (; ; ) {
-          const a = Math.floor(Math.random() * e.length), l = Math.floor(Math.random() * this.pageSize), o = e[a] * this.pageSize + l;
-          if (o + this.blockSize + 1 > t.length)
+          const c = Math.floor(Math.random() * i.length), l = Math.floor(Math.random() * this.pageSize), n = i[c] * this.pageSize + l;
+          if (n + this.blockSize + 1 > t.length)
             continue;
-          const h = t.slice(o, o + this.blockSize), g = t.slice(o + 1, o + this.blockSize + 1);
-          yield { xs: h, ys: g };
+          const h = t.slice(n, n + this.blockSize), f = t.slice(n + 1, n + this.blockSize + 1);
+          yield { xs: h, ys: f };
         }
       } else
         for (; ; ) {
-          const e = Math.floor(Math.random() * (t.length - this.blockSize - 1));
-          if (i) {
-            const o = Math.floor(e / this.pageSize), h = i.has(o);
-            if (h && !r || !h && r)
+          const i = Math.floor(Math.random() * (t.length - this.blockSize - 1));
+          if (o) {
+            const n = Math.floor(i / this.pageSize), h = o.has(n);
+            if (h && !a || !h && a)
               continue;
           }
-          const a = t.slice(e, e + this.blockSize), l = t.slice(e + 1, e + this.blockSize + 1);
-          yield { xs: a, ys: l };
+          const c = t.slice(i, i + this.blockSize), l = t.slice(i + 1, i + this.blockSize + 1);
+          yield { xs: c, ys: l };
         }
     }).bind(this);
-    return S(n).batch(s).map((e) => {
-      const a = e;
-      return u(() => ({
-        xs: a.xs.cast("int32"),
-        ys: a.ys.cast("int32")
+    return z(e).batch(s).map((i) => {
+      const c = i;
+      return g(() => ({
+        xs: c.xs.cast("int32"),
+        ys: c.ys.cast("int32")
         // this.tf.oneHot(batchData.ys.cast('int32'), this.tokenizer.vocabSize),
       }));
     }).prefetch(2);
@@ -76,6 +79,6 @@ class w {
 }
 export {
   w as DatasetBuilder,
-  p as PAGE_FACTOR,
+  S as PAGE_FACTOR,
   y as flattenTokens
 };

package/dist/training/FullTrainer.d.ts

@@ -1,10 +1,23 @@
 import { ITokeniser } from '../tokeniser/type';
-import { default as NanoGPT } from '../NanoGPTModel';
-import { default as GPTTrainer, TrainingOptions } from './Trainer';
+import { default as NanoGPT, TrainingLogEntry } from '../NanoGPTModel';
+import { default as GPTTrainer, TrainingOptions, TrainingProgress } from './Trainer';
 import { Tensor } from '@tensorflow/tfjs-core';
 import { Dataset } from '@tensorflow/tfjs-data';
 export default class FullTrainer extends GPTTrainer {
     constructor(model: NanoGPT, tokenizer: ITokeniser, learningRate?: number);
+    private createEmptyState;
+    private createLogEntry;
+    private createProgress;
+    stepDataset(dataset: Dataset<{
+        xs: Tensor;
+        ys: Tensor;
+    }>, options: Partial<TrainingOptions>, validationDataset?: Dataset<{
+        xs: Tensor;
+        ys: Tensor;
+    }>): Promise<{
+        log: TrainingLogEntry;
+        progress: TrainingProgress;
+    }>;
     trainOnDataset(dataset: Dataset<{
         xs: Tensor;
         ys: Tensor;

package/dist/training/FullTrainer.js

@@ -1,81 +1,127 @@
-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 = {
+import { generateText as v } from "../utilities/generate.js";
+import x from "./Trainer.js";
+import S from "./Evaluator.js";
+import { d as w } from "../index-BoWRt-10.js";
+import y from "../utilities/profile.js";
+const T = {
   desiredLoss: 0.01,
   logInterval: 1,
   maxSteps: 1e3
 };
-class E extends T {
-  constructor(i, e, r = 3e-4) {
-    super(i, e, r);
+class z extends x {
+  constructor(r, t, s = 3e-4) {
+    super(r, t, s);
   }
-  // Train for multiple epochs using Dataset API - FIXED memory leaks
-  async trainOnDataset(i, e, r) {
-    const { logInterval: g, onStep: l, prompt: c, maxSteps: u } = {
-      ...y,
-      ...e
-    }, n = Date.now(), t = {
+  createEmptyState() {
+    return {
       step: 0,
       lastLoss: 1e6,
       totalSteps: 0,
       losses: [],
       validationLosses: [],
-      logStartTime: n,
+      logStartTime: 0,
       trainingDuration: 0,
       ...this.lastState || {}
     };
-    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();
+  }
+  createLogEntry(r, t, s, h) {
+    return {
+      loss: r.lastLoss,
+      step: r.step,
+      time: Date.now() - t,
+      batchSize: s,
+      learningRate: h ? this.optimizer.lr : void 0
+    };
+  }
+  createProgress(r, t, s) {
+    return {
+      duration: r.trainingDuration,
+      totalSamples: r.totalSteps * t.batchSize,
+      samplesPerSecond: r.totalSteps * t.batchSize / (r.trainingDuration / 1e3),
+      memory: s ? this.model.getProfiler()?.getPeakMemory() || 0 : void 0
+    };
+  }
+  async stepDataset(r, t, s) {
+    const { logInterval: h, prompt: m } = {
+      ...T,
+      ...t
+    }, g = Date.now(), a = this.createEmptyState();
+    this.lastState = a, await this.dummyPass(), this.model.trainable = !0, t?.advancedMetrics && (this.model.getProfiler() || (this.model.config.layerConfig.profiler = new y())), this.running = !0, a.logStartTime = g;
+    const p = s ? new S(this.model, s) : void 0, e = await r.iterator();
+    try {
+      for (; this.running; ) {
+        const i = await e.next();
+        if (i.done) break;
+        const u = i.value, o = this.trainBatch(a, u), n = this.createLogEntry(a, g, u.xs.shape[0], t?.advancedMetrics);
+        if (this.model.log.push(n), a.step % h === 0) {
+          await o.data();
+          const f = Date.now();
+          if (a.trainingDuration += f - a.logStartTime, p)
+            try {
+              const l = await p.evaluate(5);
+              a.validationLosses.push(l), n.valLoss = l;
+            } catch (l) {
+              console.error("Validation error:", l);
+            }
+          if (m) {
+            const l = await v(this.tokenizer, this.model, m, 100, {
+              temperature: 0.8
+            });
+            n.example = l;
+          }
+          const c = this.createProgress(a, n, t?.advancedMetrics);
+          return o.dispose(), this.stop(), { log: n, progress: c };
+        }
+        o.dispose();
+      }
+    } catch (i) {
+      throw console.error("Training error:", i), w(), i;
+    }
+    throw w(), this.running = !1, new Error("No log returned before training stopped.");
+  }
+  // Train for multiple epochs using Dataset API - FIXED memory leaks
+  async trainOnDataset(r, t, s) {
+    const { logInterval: h, onStep: m, prompt: g, maxSteps: a } = {
+      ...T,
+      ...t
+    }, p = Date.now(), e = this.createEmptyState();
+    this.lastState = e, await this.dummyPass(), this.model.trainable = !0, t?.advancedMetrics && (this.model.getProfiler() || (this.model.config.layerConfig.profiler = new y())), this.running = !0, e.logStartTime = p;
+    const i = s ? new S(this.model, s) : void 0, u = await r.iterator();
     try {
       for (; this.running; ) {
-        const o = await f.next();
+        const o = await u.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: options?.advancedMetrics ? await state.gradientNorm : undefined,
-        };
-        if (this.model.log.push(s), t.step % g === 0) {
-          await p.data();
-          const S = Date.now();
-          if (t.trainingDuration += S - t.logStartTime, m)
+        const n = o.value, f = this.trainBatch(e, n), c = this.createLogEntry(e, p, n.xs.shape[0], t?.advancedMetrics);
+        if (this.model.log.push(c), e.step % h === 0) {
+          await f.data();
+          const l = Date.now();
+          if (e.trainingDuration += l - e.logStartTime, i)
             try {
-              const a = await m.evaluate(5);
-              t.validationLosses.push(a), s.valLoss = a;
-            } catch (a) {
-              console.error("Validation error:", a);
+              const d = await i.evaluate(5);
+              e.validationLosses.push(d), c.valLoss = d;
+            } catch (d) {
+              console.error("Validation error:", d);
             }
-          if (l) {
-            if (c) {
-              const v = await w(this.tokenizer, this.model, c, 100, {
+          if (m) {
+            if (g) {
+              const L = await v(this.tokenizer, this.model, g, 100, {
                 temperature: 0.8
               });
-              s.example = v;
+              c.example = L;
             }
-            const a = {
-              duration: t.trainingDuration,
-              totalSamples: t.totalSteps * s.batchSize,
-              samplesPerSecond: t.totalSteps * s.batchSize / (t.trainingDuration / 1e3),
-              memory: e.advancedMetrics ? this.model.getProfiler()?.getPeakMemory() || 0 : void 0
-            };
-            await l(s, a);
+            const d = this.createProgress(e, c, t?.advancedMetrics);
+            await m(c, d);
           }
-          t.logStartTime = Date.now();
+          e.logStartTime = Date.now();
         }
-        p.dispose(), t.step >= u && this.stop();
+        f.dispose(), e.step >= a && this.stop();
       }
     } catch (o) {
-      throw console.error("Training error:", o), h(), o;
+      throw console.error("Training error:", o), w(), o;
     }
-    return h(), this.running = !1, { losses: t.losses, validationLosses: t.validationLosses };
+    return w(), this.running = !1, { losses: e.losses, validationLosses: e.validationLosses };
   }
 }
 export {
-  E as default
+  z as default
 };

package/dist/training/Trainer.d.ts

@@ -66,6 +66,16 @@ export default abstract class GPTTrainer {
         losses: number[];
         validationLosses: number[];
     }>;
+    abstract stepDataset(dataset: Dataset<{
+        xs: Tensor;
+        ys: Tensor;
+    }>, options: Partial<TrainingOptions>, validationDataset?: Dataset<{
+        xs: Tensor;
+        ys: Tensor;
+    }>): Promise<{
+        log: TrainingLogEntry;
+        progress: TrainingProgress;
+    }>;
     createTrainValidationSplit(textData: string[], batchSize?: number, validationSplit?: number): Promise<{
         trainDataset: Dataset<{
             xs: Tensor;

package/dist/training/Trainer.js

@@ -1,7 +1,7 @@
 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";
+import { t as f, v as y, d as c } from "../index-BoWRt-10.js";
+import { z as m } from "../zeros--BdLQ3oG.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);

package/dist/training/sparseCrossEntropy.js

@@ -1,9 +1,9 @@
 import { gatherSub as x } from "../ops/gatherSub.js";
 import { scatterSub as L } from "../ops/scatterSub.js";
-import { y, t as u, z as C, c as E } from "../index-UdZhlibC.js";
-import { s as G } from "../softmax-JMEIUo2J.js";
-import { m as z } from "../max-ByjEGoFx.js";
-import { l as v } from "../log_sum_exp-BnmCkHWl.js";
+import { y, t as u, z as C, c as E } from "../index-BoWRt-10.js";
+import { s as G } from "../softmax-Bv_6lyMX.js";
+import { m as z } from "../max-Ddnnb5xe.js";
+import { l as v } from "../log_sum_exp-DbjkV734.js";
 function k(t, s) {
   return u(() => {
     const n = t.shape[t.shape.length - 1], c = t.shape.slice(0, -1).reduce((o, e) => o * e, 1), h = t.shape.length > 2 ? t.reshape([c, n]) : t, p = s.shape.length > 1 ? s.reshape([c]).cast("int32") : s.cast("int32"), r = z(h, -1, !0), a = E(h, r), m = v(a, -1);

package/dist/utilities/dummy.js

@@ -1,5 +1,5 @@
-import { m as y, v as P, e as S } from "../index-UdZhlibC.js";
-import { z as i } from "../zeros-DnPT2nD4.js";
+import { m as y, v as P, e as S } from "../index-BoWRt-10.js";
+import { z as i } from "../zeros--BdLQ3oG.js";
 async function w(s) {
   const t = i([1, s.config.gpt.blockSize], "int32"), [e, n] = s.forward({ training: !1 }, t);
   await e.data(), e.dispose(), n && n.dispose(), t.dispose();

package/dist/utilities/generate.js

@@ -1,6 +1,6 @@
-import "../index-UdZhlibC.js";
-import { t as m } from "../tensor2d-CWHxHpLh.js";
-import { c as u } from "../concat-CbXTetof.js";
+import "../index-BoWRt-10.js";
+import { t as m } from "../tensor2d-wxPAnDQy.js";
+import { c as u } from "../concat-CsxrgovM.js";
 async function v(o, r, a, c, f) {
   if (c <= 0)
     throw new Error("Length must be a positive integer");

package/dist/utilities/multinomialCPU.js

@@ -1,5 +1,5 @@
-import "../index-UdZhlibC.js";
-import { t as e } from "../tensor2d-CWHxHpLh.js";
+import "../index-BoWRt-10.js";
+import { t as e } from "../tensor2d-wxPAnDQy.js";
 function l(n) {
   let r = 0;
   const i = Math.random();

package/dist/utilities/performance.js

@@ -1,4 +1,4 @@
-import { t as s } from "../index-UdZhlibC.js";
+import { t as s } from "../index-BoWRt-10.js";
 async function f(e, o = 10, r = !1) {
   for (let t = 0; t < 100; t++) {
     const a = r ? await e() : s(e);

package/dist/utilities/profile.js

@@ -1,4 +1,4 @@
-import { m as a } from "../index-UdZhlibC.js";
+import { m as a } from "../index-BoWRt-10.js";
 const s = 1024 * 1024;
 class l {
   log = /* @__PURE__ */ new Map();

package/dist/utilities/safetensors.js

@@ -1,5 +1,5 @@
-import "../index-UdZhlibC.js";
-import { t as y } from "../tensor-Do9PKbIE.js";
+import "../index-BoWRt-10.js";
+import { t as y } from "../tensor-JwS7ZYY6.js";
 function l(t) {
   if (t === "float32") return "F32";
   if (t === "int32") return "I32";

package/dist/utilities/weights.js

@@ -1,5 +1,5 @@
-import "../index-UdZhlibC.js";
-import { t as p } from "../tensor-Do9PKbIE.js";
+import "../index-BoWRt-10.js";
+import { t as p } from "../tensor-JwS7ZYY6.js";
 function h(n) {
   const e = n.reduce((s, o) => s + o.length, 0), a = new Float32Array(e);
   let t = 0;

package/dist/{variable-BTBkayv_.js → variable-BuddVFLa.js}

@@ -1,4 +1,4 @@
-import { E as i } from "./index-UdZhlibC.js";
+import { E as i } from "./index-BoWRt-10.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.

package/dist/{webgpu_program-WaoMq-WD.js → webgpu_program-PFzf1hAQ.js}

@@ -1,4 +1,4 @@
-import { aa as k, ab as z, ac as E, a1 as j, l as A } from "./index-UdZhlibC.js";
+import { aa as k, ab as z, ac as E, a1 as j, l as A } from "./index-BoWRt-10.js";
 /**
  * @license
  * Copyright 2019 Google LLC. All Rights Reserved.

package/dist/{webgpu_util-DhSeP4b6.js → webgpu_util-D____QpY.js}

@@ -1,4 +1,4 @@
-import { l as u } from "./index-UdZhlibC.js";
+import { l as u } from "./index-BoWRt-10.js";
 /**
  * @license
  * Copyright 2019 Google LLC. All Rights Reserved.

package/dist/{zeros-DnPT2nD4.js → zeros--BdLQ3oG.js}

@@ -1,4 +1,4 @@
-import { B as m, C as r, a5 as l, E as c, a6 as i, F as p, a7 as u, j as f } from "./index-UdZhlibC.js";
+import { B as m, C as r, a2 as l, E as c, a6 as i, F as p, a7 as u, j as f } from "./index-BoWRt-10.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@genai-fi/nanogpt",
-    "version": "0.7.1",
+    "version": "0.7.3",
     "type": "module",
     "main": "dist/main.js",
     "types": "dist/main.d.ts",