@genai-fi/nanogpt 0.1.1 → 0.1.3

package/dist/NanoGPTModel.d.ts

@@ -1,7 +1,6 @@
 import { default as TF } from '@tensorflow/tfjs';
 import { GPTConfig } from './config';
 export interface TrainingLogEntry {
-    epoch: number;
     loss: number;
     valLoss?: number;
     step: number;

package/dist/TeachableLLM.d.ts

@@ -6,7 +6,7 @@ import { default as Generator, IGenerateOptions } from './Generator';
 import { default as Trainer, ITrainerOptions } from './Trainer';
 import { default as EE } from 'eventemitter3';
 type TeachableLLMStatus = 'warmup' | 'ready' | 'training' | 'loading' | 'busy' | 'error';
-export default class TeachableLLM extends EE<'status' | 'error'> {
+export default class TeachableLLM extends EE<'status' | 'error' | 'trainStep'> {
     private _config?;
     private _model?;
     readonly tf: typeof TF;

package/dist/TeachableLLM.js

@@ -1,13 +1,13 @@
-import a from "./NanoGPTModel.js";
-import { defaultConfig as h } from "./config.js";
-import { saveModel as d } from "./utilities/save.js";
-import { loadModel as m } from "./utilities/load.js";
-import l from "./Generator.js";
-import u from "./Trainer.js";
-import { E as _ } from "./index-SOhdqzHq.js";
-import { dummyPassAsync as f } from "./utilities/dummy.js";
-import c from "./tokeniser/CharTokeniser.js";
-class s extends _ {
+import d from "./NanoGPTModel.js";
+import { defaultConfig as m } from "./config.js";
+import { saveModel as u } from "./utilities/save.js";
+import { loadModel as l } from "./utilities/load.js";
+import f from "./Generator.js";
+import _ from "./Trainer.js";
+import { E as c } from "./index-SOhdqzHq.js";
+import { dummyPassAsync as a } from "./utilities/dummy.js";
+import g from "./tokeniser/CharTokeniser.js";
+class n extends c {
   _config;
   _model;
   tf;
@@ -43,23 +43,27 @@ class s extends _ {
   saveModel() {
     if (!this._model || !this._tokeniser)
       throw new Error("Model or tokeniser is not initialized.");
-    return d(this._model, this._tokeniser);
+    return u(this._model, this._tokeniser);
   }
   static loadModel(t, r) {
-    const e = new s(t);
-    return m(t, r).then(({ model: i, tokeniser: o }) => {
-      e._model = i, e._tokeniser = o, e._config = i.config, e.setStatus("warmup"), f(i).then(() => {
+    const e = new n(t);
+    return l(t, r).then(({ model: i, tokeniser: o }) => {
+      e._model = i, e._tokeniser = o, e._config = i.config, e.setStatus("warmup"), a(i).then(() => {
         e.setStatus("ready");
-      }).catch((n) => {
-        e.setStatus("error"), e.emit("error", n);
+      }).catch((s) => {
+        e.setStatus("error"), e.emit("error", s);
       });
     }).catch((i) => {
       e.setStatus("error"), e.emit("error", i);
     }), e;
   }
   static create(t, r = {}) {
-    const e = { ...h, ...r }, i = new c(e.vocabSize), o = new a(t, e);
-    return new s(t, i, o);
+    const e = { ...m, ...r }, i = new g(e.vocabSize), o = new d(t, e), s = new n(t, i, o);
+    return s.setStatus("warmup"), a(o).then(() => {
+      s.setStatus("ready");
+    }).catch((h) => {
+      s.setStatus("error"), s.emit("error", h);
+    }), s;
   }
   getNumParams() {
     if (!this._model)
@@ -69,8 +73,12 @@ class s extends _ {
   trainer() {
     if (!this._model || !this._tokeniser)
       throw new Error("Model or tokeniser is not initialized.");
-    const t = new u(this._model, this._tokeniser);
-    return t.on("start", () => this.setStatus("training")), t.on("stop", () => this.setStatus("ready")), t;
+    const t = new _(this._model, this._tokeniser);
+    return t.on("start", () => this.setStatus("training")), t.on("stop", () => this.setStatus("ready")), t.on("log", async (r) => {
+      const e = this.listeners("trainStep");
+      for (const i of e)
+        await i(r);
+    }), t;
   }
   train(t, r) {
     return this.trainer().train(t, r);
@@ -78,13 +86,17 @@ class s extends _ {
   generator() {
     if (!this._model || !this._tokeniser)
       throw new Error("Model or tokeniser is not initialized.");
-    const t = new l(this._model, this._tokeniser);
-    return t.on("start", () => this.setStatus("busy")), t.on("stop", () => this.setStatus("ready")), t;
+    const t = new f(this._model, this._tokeniser);
+    return t.on("start", () => {
+      this.status === "ready" && this.setStatus("busy");
+    }), t.on("stop", () => {
+      this.status === "busy" && this.setStatus("ready");
+    }), t;
   }
   generateText(t, r) {
     return this.generator().generate(t, r);
   }
 }
 export {
-  s as default
+  n as default
 };

package/dist/Trainer.d.ts

@@ -2,7 +2,6 @@ import { default as NanoGPT } from './NanoGPTModel';
 import { ITokeniser } from './tokeniser/type';
 import { default as EE } from 'eventemitter3';
 export interface ITrainerOptions {
-    epochs?: number;
     batchSize?: number;
     learningRate?: number;
     maxSteps?: number;

package/dist/Trainer.js

@@ -1,9 +1,9 @@
-import { E as s } from "./index-SOhdqzHq.js";
-import n from "./training/FullTrainer.js";
-class o extends s {
+import { E as l } from "./index-SOhdqzHq.js";
+import o from "./training/FullTrainer.js";
+class d extends l {
   trainer;
   constructor(a, t) {
-    super(), this.trainer = new n(a.tf, a, t, 1e-3);
+    super(), this.trainer = new o(a.tf, a, t, 1e-3);
   }
   stop() {
   }
@@ -13,16 +13,17 @@ class o extends s {
       t?.batchSize || 32,
       t?.validationSplit || 0.1
     );
-    this.emit("start"), await this.trainer.trainOnDataset(
+    this.trainer.setLearningRate(t?.learningRate || 1e-3), this.emit("start"), await this.trainer.trainOnDataset(
       e,
       {
-        epochs: t?.epochs || 2,
         prompt: t?.prompt,
-        stepsPerEpoch: t?.maxSteps || 100,
         logInterval: t?.logInterval || 10,
         desiredLoss: t?.desiredLoss || 0.01,
+        maxSteps: t?.maxSteps || 1e3,
         onStep: async (i) => {
-          this.emit("log", i);
+          const s = this.listeners("log");
+          for (const n of s)
+            await n(i);
         }
       },
       r
@@ -30,5 +31,5 @@ class o extends s {
   }
 }
 export {
-  o as default
+  d as default
 };

package/dist/training/FullTrainer.js

@@ -1,75 +1,67 @@
-import { generateText as g } from "../utilities/generate.js";
-import T from "./Trainer.js";
-const b = {
-  epochs: 1,
-  stepsPerEpoch: 1e6,
+import { generateText as L } from "../utilities/generate.js";
+import w from "./Trainer.js";
+const g = {
   desiredLoss: 0.01,
-  logInterval: 1
+  logInterval: 1,
+  maxSteps: 1e3
 };
-class S extends T {
-  constructor(a, r, t, i = 3e-4) {
-    super(a, r, t, i);
+class S 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(a, r, t) {
-    const { epochs: i, stepsPerEpoch: n, desiredLoss: c, logInterval: L, onStep: h, onEpoch: o, prompt: l } = {
-      ...b,
-      ...r
+  async trainOnDataset(r, i, o) {
+    const { desiredLoss: n, logInterval: c, onStep: l, prompt: p, maxSteps: d } = {
+      ...g,
+      ...i
     }, s = {
-      epoch: 0,
       pass: 0,
       depth: 1,
       step: 0,
       stepSinceDepthChange: 0,
       lastLoss: 1e6,
-      epochLoss: 0,
       totalSteps: 0,
       losses: [],
       validationLosses: []
     };
     this.dummyPass(), this.model.trainable = !0;
     const m = Date.now();
-    for (s.epoch = 0; s.epoch < i; s.epoch++) {
-      s.step = 0, s.epochLoss = 0, s.pass = 0, s.depth = 1, s.stepSinceDepthChange = 0;
-      const u = await a.iterator();
-      try {
-        for (; !(n && s.step >= n || s.lastLoss < c); ) {
-          const e = await u.next();
-          if (e.done) break;
-          const f = e.value, w = this.trainBatch(s, f), p = {
-            epoch: s.epoch,
-            loss: s.lastLoss,
-            step: s.step,
-            time: Date.now() - m,
-            batchSize: f.xs.shape[0]
-          };
-          if (this.model.log.push(p), s.step % L === 0 && (await w, h)) {
-            if (l) {
-              const v = await g(this.tokenizer, this.model, l, 100, {
+    this.running = !0;
+    const u = await r.iterator();
+    try {
+      for (; this.running && !(s.lastLoss < n); ) {
+        const e = await u.next();
+        if (e.done) break;
+        const h = e.value, f = this.trainBatch(s, h), a = {
+          loss: s.lastLoss,
+          step: s.step,
+          time: Date.now() - m,
+          batchSize: h.xs.shape[0]
+        };
+        if (this.model.log.push(a), s.step % c === 0) {
+          if (await f, o)
+            try {
+              const t = await this.evaluateOnDataset(o, 5);
+              s.validationLosses.push(t), a.valLoss = t;
+            } catch (t) {
+              console.error("Validation error:", t);
+            }
+          if (l) {
+            if (p) {
+              const t = await L(this.tokenizer, this.model, p, 100, {
                 temperature: 0.8
               });
-              p.example = v;
+              a.example = t;
             }
-            await h(p);
+            await l(a);
           }
         }
-      } catch (e) {
-        throw console.error("Training error:", e), this.tf.dispose(), e;
+        s.step >= d && this.stop();
       }
-      const d = s.epochLoss / s.step;
-      if (t)
-        try {
-          const e = await this.evaluateOnDataset(t, 5);
-          s.validationLosses.push(e), o && await o(s.epoch, d, e);
-        } catch (e) {
-          console.error("Validation error:", e);
-        }
-      else
-        o && o(s.epoch, d);
-      if (this.tf.dispose(), s.lastLoss < c)
-        break;
+    } catch (e) {
+      throw console.error("Training error:", e), this.tf.dispose(), e;
     }
-    return { losses: s.losses, validationLosses: s.validationLosses };
+    return this.tf.dispose(), this.running = !1, { losses: s.losses, validationLosses: s.validationLosses };
   }
 }
 export {

package/dist/training/LayerTrainer.js

@@ -1,111 +1,91 @@
-import { generateText as v } from "../utilities/generate.js";
-import T from "./Trainer.js";
-import { schedule as k } from "./lwSchedule.js";
-const x = {
-  epochs: 1,
-  stepsPerEpoch: 1e6,
+import { generateText as d } from "../utilities/generate.js";
+import S from "./Trainer.js";
+import { schedule as u } from "./lwSchedule.js";
+const w = {
   desiredLoss: 0.01,
   logInterval: 1,
   stepsPerLayer: 400,
-  maxPasses: 3
+  maxPasses: 3,
+  maxSteps: 1e3
 };
-class D extends T {
+class b extends S {
   trainingPattern = [];
   startPass = 0;
   startLayer = 0;
-  constructor(o, e, t, h = 3e-4) {
-    if (super(o, e, t, h), this.trainingPattern = k[e.config.nLayer - 1] || [], e.log.length > 0) {
-      const r = e.log[e.log.length - 1];
-      r.pass !== void 0 && r.layer !== void 0 && (this.startPass = r.pass, this.startLayer = r.layer, console.log(`Resuming training from pass ${this.startPass}, layer ${this.startLayer}`));
+  constructor(r, a, e, p = 3e-4) {
+    if (super(r, a, e, p), this.trainingPattern = u[a.config.nLayer - 1] || [], a.log.length > 0) {
+      const i = a.log[a.log.length - 1];
+      i.pass !== void 0 && i.layer !== void 0 && (this.startPass = i.pass, this.startLayer = i.layer, console.log(`Resuming training from pass ${this.startPass}, layer ${this.startLayer}`));
     }
   }
-  applyTrainingPattern(o) {
-    const e = o < this.trainingPattern.length ? o : this.trainingPattern.length - 1, t = this.trainingPattern[e];
-    this.model.setSkipMask(t.skip), this.model.setTrainableMask(t.trainable), this.resetOptimizer(t.adam), console.log("Applied training pattern:", e, t);
+  applyTrainingPattern(r) {
+    const a = r < this.trainingPattern.length ? r : this.trainingPattern.length - 1, e = this.trainingPattern[a];
+    this.model.setSkipMask(e.skip), this.model.setTrainableMask(e.trainable), this.resetOptimizer(e.adam), console.log("Applied training pattern:", a, e);
   }
   // Train for multiple epochs using Dataset API - FIXED memory leaks
-  async trainOnDataset(o, e, t) {
-    const {
-      epochs: h,
-      stepsPerEpoch: r,
-      desiredLoss: c,
-      logInterval: P,
-      stepsPerLayer: d,
-      onLayerChange: n,
-      onPassComplete: g,
-      onStep: y,
-      onEpoch: p,
-      prompt: L
-    } = {
-      ...x,
-      ...e
-    }, s = {
-      epoch: 0,
+  async trainOnDataset(r, a, e) {
+    const { desiredLoss: p, logInterval: i, stepsPerLayer: L, onLayerChange: l, onPassComplete: h, onStep: c, prompt: g } = {
+      ...w,
+      ...a
+    }, t = {
       pass: 0,
       layerStep: 0,
       step: 0,
       stepSinceLayerChange: 0,
       lastLoss: 1e6,
-      epochLoss: 0,
       totalSteps: 0,
       losses: [],
       validationLosses: []
     };
     this.dummyPass();
-    const S = Date.now();
-    for (s.epoch = 0; s.epoch < h; s.epoch++) {
-      s.step = 0, s.epochLoss = 0, s.pass = this.startPass, s.layerStep = this.startLayer + this.startPass * this.model.config.nLayer, s.stepSinceLayerChange = 0, this.startPass = 0, this.startLayer = 0;
-      const u = await o.iterator();
-      this.applyTrainingPattern(s.layerStep % this.trainingPattern.length);
-      try {
-        for (; !(r && s.step >= r || s.lastLoss < c); ) {
-          const a = await u.next();
-          if (a.done) break;
-          const m = a.value, w = this.trainBatch(s, m);
-          s.stepSinceLayerChange++;
-          const l = {
-            epoch: s.epoch,
-            loss: s.lastLoss,
-            step: s.step,
-            time: Date.now() - S,
-            batchSize: m.xs.shape[0],
-            pass: s.pass,
-            layer: s.layerStep % this.model.config.nLayer
-          };
-          if (this.model.log.push(l), s.step % P === 0 && (await w, y)) {
-            if (L) {
-              const i = await v(this.tokenizer, this.model, L, 100, {
+    const m = Date.now();
+    this.startPass = 0, this.startLayer = 0;
+    const f = await r.iterator();
+    this.applyTrainingPattern(t.layerStep % this.trainingPattern.length);
+    try {
+      for (; !(t.lastLoss < p); ) {
+        const n = await f.next();
+        if (n.done) break;
+        const y = n.value, P = this.trainBatch(t, y);
+        t.stepSinceLayerChange++;
+        const o = {
+          loss: t.lastLoss,
+          step: t.step,
+          time: Date.now() - m,
+          batchSize: y.xs.shape[0],
+          pass: t.pass,
+          layer: t.layerStep % this.model.config.nLayer
+        };
+        if (this.model.log.push(o), t.step % i === 0) {
+          if (await P, e)
+            try {
+              const s = await this.evaluateOnDataset(e, 5);
+              t.validationLosses.push(s), o.valLoss = s;
+            } catch (s) {
+              console.error("Validation error:", s);
+            }
+          if (c) {
+            if (g) {
+              const s = await d(this.tokenizer, this.model, g, 100, {
                 temperature: 0.8,
                 topK: 10
               });
-              l.example = i;
+              o.example = s;
             }
-            await y(l);
-          }
-          if (s.stepSinceLayerChange >= d) {
-            let i;
-            t && (i = await this.evaluateOnDataset(t, 5), s.validationLosses.push(i), l.valLoss = i), s.layerStep++, s.layerStep % this.model.config.nLayer === 0 ? (n && await n(s.layerStep, s.pass, i), g && await g(s.pass), s.pass++) : n && await n(s.layerStep, s.pass, i), s.stepSinceLayerChange = 0, this.applyTrainingPattern(s.layerStep % this.trainingPattern.length);
+            await c(o);
           }
         }
-      } catch (a) {
-        throw console.error("Training error:", a), this.tf.dispose(), a;
-      }
-      const f = s.epochLoss / s.step;
-      if (t)
-        try {
-          const a = await this.evaluateOnDataset(t, 5);
-          s.validationLosses.push(a), p && await p(s.epoch, f, a);
-        } catch (a) {
-          console.error("Validation error:", a);
+        if (t.stepSinceLayerChange >= L) {
+          let s;
+          e && (s = await this.evaluateOnDataset(e, 5), t.validationLosses.push(s), o.valLoss = s), t.layerStep++, t.layerStep % this.model.config.nLayer === 0 ? (l && await l(t.layerStep, t.pass, s), h && await h(t.pass), t.pass++) : l && await l(t.layerStep, t.pass, s), t.stepSinceLayerChange = 0, this.applyTrainingPattern(t.layerStep % this.trainingPattern.length);
         }
-      else
-        p && p(s.epoch, f);
-      if (this.tf.dispose(), s.lastLoss < c)
-        break;
+      }
+    } catch (n) {
+      throw console.error("Training error:", n), this.tf.dispose(), n;
     }
-    return { losses: s.losses, validationLosses: s.validationLosses };
+    return this.tf.dispose(), { losses: t.losses, validationLosses: t.validationLosses };
   }
 }
 export {
-  D as default
+  b as default
 };

package/dist/training/Trainer.d.ts

@@ -4,10 +4,8 @@ import { default as NanoGPT, TrainingLogEntry } from '../NanoGPTModel';
 import { default as TF } from '@tensorflow/tfjs';
 import { default as AdamExt } from './AdamExt';
 export interface TrainingState {
-    epoch: number;
     step: number;
     lastLoss: number;
-    epochLoss: number;
     totalSteps: number;
     losses: number[];
     validationLosses: number[];
@@ -19,12 +17,10 @@ export interface AdamConfig {
     epsilon: number;
 }
 export interface TrainingOptions {
-    epochs: number;
-    stepsPerEpoch: number;
     desiredLoss: number;
     logInterval: number;
     prompt?: string;
-    onEpoch?: (e: number, loss: number, valLoss?: number) => Promise<void> | void;
+    maxSteps: number;
     onStep?: (log: TrainingLogEntry) => Promise<void> | void;
 }
 export default abstract class GPTTrainer {
@@ -34,7 +30,10 @@ export default abstract class GPTTrainer {
     protected datasetBuilder: DatasetBuilder;
     protected tf: typeof TF;
     protected learningRate: number;
+    protected running: boolean;
     constructor(tf: typeof TF, model: NanoGPT, tokenizer: ITokeniser, learningRate?: number);
+    setLearningRate(learningRate: number): void;
+    stop(): void;
     getOptimizer(): AdamExt;
     resetOptimizer(config?: AdamConfig): void;
     private printGradients;

package/dist/training/Trainer.js

@@ -9,6 +9,13 @@ class y {
   datasetBuilder;
   tf;
   learningRate;
+  running = !1;
+  setLearningRate(t) {
+    this.learningRate = t, this.resetOptimizer({ learningRateFactor: 1, beta1: 0.9, beta2: 0.99, epsilon: 1e-8 });
+  }
+  stop() {
+    this.running = !1;
+  }
   getOptimizer() {
     return this.optimizer;
   }
@@ -57,7 +64,7 @@ class y {
   async trainBatch(t, e) {
     try {
       const s = this.trainStep(e, !1, !1);
-      return e.xs.dispose(), e.ys.dispose(), t.step++, t.totalSteps++, s.array().then((a) => (t.lastLoss = a, t.losses.push(t.lastLoss), t.epochLoss += t.lastLoss, s.dispose(), t.lastLoss));
+      return e.xs.dispose(), e.ys.dispose(), t.step++, t.totalSteps++, s.array().then((a) => (t.lastLoss = a, 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;
     }

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@genai-fi/nanogpt",
-    "version": "0.1.1",
+    "version": "0.1.3",
     "type": "module",
     "main": "dist/main.js",
     "types": "dist/main.d.ts",
@@ -23,7 +23,7 @@
         "test": "vitest",
         "ci:test": "vitest --coverage --reporter=junit --outputFile=junit.xml",
         "coverage": "vitest run --coverage",
-        "train": "tsx scripts/train.ts --epochs 2 --batch 64",
+        "train": "tsx scripts/train.ts --batch 64",
         "generate": "tsx scripts/generate.ts",
         "evaluate": "tsx scripts/evaluate.ts",
         "debug": "tsx scripts/debug.ts"