@genai-fi/nanogpt 0.1.8 → 0.1.9

package/dist/Trainer.d.ts

@@ -12,7 +12,9 @@ export interface ITrainerOptions {
 }
 export default class Trainer extends EE<'start' | 'stop' | 'log'> {
     private trainer;
+    private hasTrained;
     constructor(model: NanoGPT, tokeniser: ITokeniser);
     stop(): void;
+    reset(): void;
     train(text: string[], options?: ITrainerOptions): Promise<void>;
 }

package/dist/Trainer.js

@@ -1,11 +1,16 @@
 import { E as l } from "./index-SOhdqzHq.js";
-import o from "./training/FullTrainer.js";
-class d extends l {
+import h from "./training/FullTrainer.js";
+class m extends l {
   trainer;
+  hasTrained = !1;
   constructor(a, t) {
-    super(), this.trainer = new o(a.tf, a, t, 1e-3);
+    super(), this.trainer = new h(a.tf, a, t, 1e-3);
   }
   stop() {
+    this.trainer.stop();
+  }
+  reset() {
+    this.hasTrained = !1, this.trainer.reset();
   }
   async train(a, t) {
     const { trainDataset: e, validationDataset: r } = await this.trainer.createTrainValidationSplit(
@@ -13,7 +18,7 @@ class d extends l {
       t?.batchSize || 32,
       t?.validationSplit || 0.1
     );
-    this.trainer.setLearningRate(t?.learningRate || 1e-3), this.emit("start"), await this.trainer.trainOnDataset(
+    this.hasTrained || this.trainer.setLearningRate(t?.learningRate || 1e-3), this.hasTrained = !0, this.emit("start"), await this.trainer.trainOnDataset(
       e,
       {
         prompt: t?.prompt,
@@ -31,5 +36,5 @@ class d extends l {
   }
 }
 export {
-  d as default
+  m as default
 };

package/dist/training/FullTrainer.js

@@ -1,70 +1,68 @@
 import { generateText as L } from "../utilities/generate.js";
 import w from "./Trainer.js";
-import g from "./Evaluator.js";
-const x = {
+import x from "./Evaluator.js";
+const g = {
   desiredLoss: 0.01,
   logInterval: 1,
   maxSteps: 1e3
 };
-class D extends w {
+class P extends w {
   constructor(r, i, o, n = 3e-4) {
     super(r, i, o, n);
   }
   // Train for multiple epochs using Dataset API - FIXED memory leaks
   async trainOnDataset(r, i, o) {
-    const { desiredLoss: n, logInterval: d, onStep: l, prompt: p, maxSteps: m } = {
-      ...x,
+    const { desiredLoss: n, logInterval: m, onStep: l, prompt: c, maxSteps: d } = {
+      ...g,
       ...i
-    }, s = {
-      pass: 0,
-      depth: 1,
+    }, t = {
       step: 0,
-      stepSinceDepthChange: 0,
       lastLoss: 1e6,
       totalSteps: 0,
       losses: [],
-      validationLosses: []
+      validationLosses: [],
+      ...this.lastState || {}
     };
-    this.dummyPass(), this.model.trainable = !0;
+    this.lastState = t, this.dummyPass(), this.model.trainable = !0;
     const u = Date.now();
     this.running = !0;
-    const c = o ? new g(this.model, o) : void 0, f = await r.iterator();
+    const h = o ? new x(this.model, o) : void 0, f = await r.iterator();
     try {
-      for (; this.running && !(s.lastLoss < n); ) {
+      for (; this.running && !(t.lastLoss < n); ) {
         const e = await f.next();
         if (e.done) break;
-        const h = e.value, v = this.trainBatch(s, h), a = {
-          loss: s.lastLoss,
-          step: s.step,
+        const p = e.value, v = this.trainBatch(t, p), a = {
+          loss: t.lastLoss,
+          step: t.step,
           time: Date.now() - u,
-          batchSize: h.xs.shape[0]
+          batchSize: p.xs.shape[0]
         };
-        if (this.model.log.push(a), s.step % d === 0) {
-          if (await v, c)
+        if (this.model.log.push(a), t.step % m === 0) {
+          if (await v, h)
             try {
-              const t = await c.evaluate(5);
-              s.validationLosses.push(t), a.valLoss = t;
-            } catch (t) {
-              console.error("Validation error:", t);
+              const s = await h.evaluate(5);
+              t.validationLosses.push(s), a.valLoss = s;
+            } catch (s) {
+              console.error("Validation error:", s);
             }
           if (l) {
-            if (p) {
-              const t = await L(this.tokenizer, this.model, p, 100, {
+            if (c) {
+              const s = await L(this.tokenizer, this.model, c, 100, {
                 temperature: 0.8
               });
-              a.example = t;
+              a.example = s;
             }
             await l(a);
           }
         }
-        s.step >= m && this.stop();
+        t.step >= d && this.stop();
       }
     } catch (e) {
       throw console.error("Training error:", e), this.tf.dispose(), e;
     }
-    return this.tf.dispose(), this.running = !1, { losses: s.losses, validationLosses: s.validationLosses };
+    return this.tf.dispose(), this.running = !1, { losses: t.losses, validationLosses: t.validationLosses };
   }
 }
 export {
-  D as default
+  P as default
 };

package/dist/training/Trainer.d.ts

@@ -31,8 +31,10 @@ export default abstract class GPTTrainer {
     protected tf: typeof TF;
     protected learningRate: number;
     protected running: boolean;
+    protected lastState?: TrainingState;
     constructor(tf: typeof TF, model: NanoGPT, tokenizer: ITokeniser, learningRate?: number);
     setLearningRate(learningRate: number): void;
+    reset(): void;
     stop(): void;
     getOptimizer(): AdamExt;
     resetOptimizer(config?: AdamConfig): void;

package/dist/training/Trainer.js

@@ -1,8 +1,8 @@
 import { DatasetBuilder as d } from "./DatasetBuilder.js";
-import p from "./AdamExt.js";
+import h from "./AdamExt.js";
 class u {
-  constructor(t, s, e, i = 1e-3) {
-    this.tokenizer = e, this.tf = t, this.model = s, this.learningRate = i, this.resetOptimizer(), this.datasetBuilder = new d(this.tf, e, s.config.blockSize);
+  constructor(t, e, s, i = 1e-3) {
+    this.tokenizer = s, this.tf = t, this.model = e, this.learningRate = i, this.resetOptimizer(), this.datasetBuilder = new d(this.tf, s, e.config.blockSize);
   }
   model;
   optimizer;
@@ -10,9 +10,13 @@ class u {
   tf;
   learningRate;
   running = !1;
+  lastState;
   setLearningRate(t) {
     this.learningRate = t, this.resetOptimizer({ learningRateFactor: 1, beta1: 0.9, beta2: 0.99, epsilon: 1e-8 });
   }
+  reset() {
+    this.lastState = void 0, this.running = !1;
+  }
   stop() {
     this.running = !1;
   }
@@ -21,7 +25,7 @@ class u {
   }
   resetOptimizer(t = { learningRateFactor: 1, beta1: 0.9, beta2: 0.99, epsilon: 1e-8 }) {
     this.optimizer && this.optimizer.dispose();
-    const s = new p(
+    const e = new h(
       t.learningRateFactor * this.learningRate,
       t.beta1,
       t.beta2,
@@ -33,53 +37,53 @@ class u {
         weightDecay: 0
       }
     );
-    this.optimizer = s;
+    this.optimizer = e;
   }
   printGradients(t) {
-    Object.keys(t).forEach((s) => {
-      const e = t[s];
-      console.log(`${s}:`), console.log(`  Shape: ${e.shape}`), console.log(`  Mean: ${this.tf.mean(e).dataSync()[0]}`), console.log(`  Std: ${this.tf.moments(e).variance.sqrt().dataSync()[0]}`), console.log(`  Min: ${this.tf.min(e).dataSync()[0]}`), console.log(`  Max: ${this.tf.max(e).dataSync()[0]}`), console.log(`  Norm: ${this.tf.norm(e).dataSync()[0]}`);
+    Object.keys(t).forEach((e) => {
+      const s = t[e];
+      console.log(`${e}:`), console.log(`  Shape: ${s.shape}`), console.log(`  Mean: ${this.tf.mean(s).dataSync()[0]}`), console.log(`  Std: ${this.tf.moments(s).variance.sqrt().dataSync()[0]}`), console.log(`  Min: ${this.tf.min(s).dataSync()[0]}`), console.log(`  Max: ${this.tf.max(s).dataSync()[0]}`), console.log(`  Norm: ${this.tf.norm(s).dataSync()[0]}`);
     });
   }
-  trainStep(t, s = !1, e = !1) {
+  trainStep(t, e = !1, s = !1) {
     return this.tf.tidy(() => {
       const { xs: i, ys: a } = t, o = () => {
         const { loss: l, logits: c } = this.model.forward(i, a, !0);
         return c.dispose(), l;
       }, { value: n, grads: r } = this.tf.variableGrads(o);
-      return s || (e && (console.log("-------"), this.printGradients(r), console.log("-------")), this.optimizer.applyGradients(r), this.tf.dispose(r)), n;
+      return e || (s && (console.log("-------"), this.printGradients(r), console.log("-------")), this.optimizer.applyGradients(r), this.tf.dispose(r)), n;
     });
   }
   dummyPass() {
-    const t = this.tf.zeros([1, this.model.config.blockSize], "int32"), s = this.tf.zeros([1, this.model.config.blockSize, this.model.config.vocabSize]);
+    const t = this.tf.zeros([1, this.model.config.blockSize], "int32"), e = this.tf.zeros([1, this.model.config.blockSize, this.model.config.vocabSize]);
     try {
-      const e = this.trainStep({ xs: t, ys: s }, !0);
-      e.dataSync(), e.dispose();
-    } catch (e) {
-      console.error("Error during dummy pass:", e);
+      const s = this.trainStep({ xs: t, ys: e }, !0);
+      s.dataSync(), s.dispose();
+    } catch (s) {
+      console.error("Error during dummy pass:", s);
     } finally {
-      t.dispose(), s.dispose();
+      t.dispose(), e.dispose();
     }
   }
-  async trainBatch(t, s) {
+  async trainBatch(t, e) {
     try {
-      const e = this.trainStep(s, !1, !1);
-      return s.xs.dispose(), s.ys.dispose(), t.step++, t.totalSteps++, e.array().then((i) => (t.lastLoss = i, t.losses.push(t.lastLoss), e.dispose(), t.lastLoss));
-    } catch (e) {
-      throw console.error(`Error processing batch at step ${t.step}:`, e), this.tf.dispose(), e;
+      const s = this.trainStep(e, !1, !1);
+      return e.xs.dispose(), e.ys.dispose(), t.step++, t.totalSteps++, s.array().then((i) => (t.lastLoss = i, t.losses.push(t.lastLoss), s.dispose(), t.lastLoss));
+    } catch (s) {
+      throw console.error(`Error processing batch at step ${t.step}:`, s), this.tf.dispose(), s;
     }
   }
-  async createTrainValidationSplit(t, s = 32, e = 0.1) {
-    const i = await this.datasetBuilder.createTextDataset(t, s, 0, 1 - e), a = await this.datasetBuilder.createTextDataset(
+  async createTrainValidationSplit(t, e = 32, s = 0.1) {
+    const i = await this.datasetBuilder.createTextDataset(t, e, 0, 1 - s), a = await this.datasetBuilder.createTextDataset(
       t,
-      s,
-      1 - e,
+      e,
+      1 - s,
       1
     );
     return { trainDataset: i, validationDataset: a };
   }
-  async createDataset(t, s = 32) {
-    return await this.datasetBuilder.createTextDataset(t, s);
+  async createDataset(t, e = 32) {
+    return await this.datasetBuilder.createTextDataset(t, e);
   }
   dispose() {
     this.optimizer && this.optimizer.dispose();

package/package.json

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