@genai-fi/nanogpt 0.0.1

package/dist/training/LayerTrainer.js

@@ -0,0 +1,108 @@
+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,
+  desiredLoss: 0.01,
+  logInterval: 1,
+  stepsPerLayer: 400,
+  maxPasses: 3
+};
+class D extends T {
+  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}`));
+    }
+  }
+  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);
+  }
+  // Train for multiple epochs using Dataset API - FIXED memory leaks
+  async trainOnDataset(o, e, t) {
+    const {
+      epochs: h,
+      stepsPerEpoch: r,
+      desiredLoss: c,
+      logInterval: m,
+      stepsPerLayer: d,
+      onLayerChange: n,
+      onPassComplete: g,
+      onStep: y,
+      onEpoch: p,
+      prompt: L
+    } = {
+      ...x,
+      ...e
+    }, s = {
+      epoch: 0,
+      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 P = a.value, w = this.trainBatch(s, P);
+          s.stepSinceLayerChange++;
+          const l = {
+            epoch: s.epoch,
+            loss: s.lastLoss,
+            step: s.step,
+            time: Date.now() - S,
+            batchSize: P.xs.shape[0],
+            pass: s.pass,
+            layer: s.layerStep % this.model.config.nLayer
+          };
+          if (this.model.log.push(l), s.step % m === 0 && (await w, y)) {
+            if (L) {
+              const i = await v(this.tokenizer, this.model, L, 100, 0.8, 10);
+              l.example = i;
+            }
+            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);
+          }
+        }
+      } 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);
+        }
+      else
+        p && p(s.epoch, f);
+      if (this.tf.dispose(), s.lastLoss < c)
+        break;
+    }
+    return { losses: s.losses, validationLosses: s.validationLosses };
+  }
+}
+export {
+  D as default
+};

package/dist/training/Trainer.d.ts

@@ -0,0 +1,73 @@
+import { ITokeniser } from '../tokeniser/type';
+import { DatasetBuilder } from './DatasetBuilder';
+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[];
+}
+export interface AdamConfig {
+    learningRateFactor: number;
+    beta1: number;
+    beta2: number;
+    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;
+    onStep?: (log: TrainingLogEntry) => Promise<void> | void;
+}
+export default abstract class GPTTrainer {
+    protected tokenizer: ITokeniser;
+    protected model: NanoGPT;
+    protected optimizer: AdamExt;
+    protected datasetBuilder: DatasetBuilder;
+    protected tf: typeof TF;
+    protected learningRate: number;
+    constructor(tf: typeof TF, model: NanoGPT, tokenizer: ITokeniser, learningRate?: number);
+    getOptimizer(): AdamExt;
+    resetOptimizer(config?: AdamConfig): void;
+    private printGradients;
+    protected trainStep(batch: {
+        xs: TF.Tensor;
+        ys: TF.Tensor;
+    }, dummy?: boolean, print?: boolean): TF.Scalar;
+    protected dummyPass(): void;
+    protected trainBatch(state: TrainingState, batch: {
+        xs: TF.Tensor;
+        ys: TF.Tensor;
+    }): Promise<number>;
+    abstract trainOnDataset(dataset: TF.data.Dataset<{
+        xs: TF.Tensor;
+        ys: TF.Tensor;
+    }>, options: Partial<TrainingOptions>, validationDataset?: TF.data.Dataset<{
+        xs: TF.Tensor;
+        ys: TF.Tensor;
+    }>): Promise<{
+        losses: number[];
+        validationLosses: number[];
+    }>;
+    evaluateOnDataset(dataset: TF.data.Dataset<TF.TensorContainer>, maxBatches?: number): Promise<number>;
+    createTrainValidationSplit(textData: string[], batchSize?: number, validationSplit?: number): Promise<{
+        trainDataset: TF.data.Dataset<{
+            xs: TF.Tensor;
+            ys: TF.Tensor;
+        }>;
+        validationDataset: TF.data.Dataset<{
+            xs: TF.Tensor;
+            ys: TF.Tensor;
+        }>;
+    }>;
+    createDataset(textData: string[], batchSize?: number): Promise<TF.data.Dataset<TF.TensorContainer>>;
+    dispose(): void;
+}

package/dist/training/Trainer.js

@@ -0,0 +1,87 @@
+import { DatasetBuilder as h } from "./DatasetBuilder.js";
+import p from "./AdamExt.js";
+class y {
+  constructor(t, e, s, a = 1e-3) {
+    this.tokenizer = s, this.tf = t, this.model = e, this.learningRate = a, this.resetOptimizer(), this.datasetBuilder = new h(this.tf, s, e.config.blockSize);
+  }
+  model;
+  optimizer;
+  datasetBuilder;
+  tf;
+  learningRate;
+  getOptimizer() {
+    return this.optimizer;
+  }
+  resetOptimizer(t = { learningRateFactor: 1, beta1: 0.9, beta2: 0.99, epsilon: 1e-8 }) {
+    this.optimizer && this.optimizer.dispose();
+    const e = new p(
+      t.learningRateFactor * this.learningRate,
+      t.beta1,
+      t.beta2,
+      t.epsilon,
+      {
+        warmupSteps: 100,
+        decaySteps: 2e4,
+        minLearningRate: 1e-4,
+        weightDecay: 0
+      }
+    );
+    this.optimizer = e;
+  }
+  printGradients(t) {
+    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, e = !1, s = !1) {
+    return this.tf.tidy(() => {
+      const { xs: a, ys: o } = t, r = () => {
+        const { loss: l, logits: c } = this.model.forward(a, o, !0);
+        return c.dispose(), l;
+      }, { value: n, grads: i } = this.tf.variableGrads(r);
+      return e || (s && (console.log("-------"), this.printGradients(i), console.log("-------")), this.optimizer.applyGradients(i), this.tf.dispose(i)), n;
+    });
+  }
+  dummyPass() {
+    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 s = this.trainStep({ xs: t, ys: e }, !0);
+      s.dataSync(), s.dispose();
+    } catch (s) {
+      console.error("Error during dummy pass:", s);
+    } finally {
+      t.dispose(), e.dispose();
+    }
+  }
+  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));
+    } catch (s) {
+      throw console.error(`Error processing batch at step ${t.step}:`, s), this.tf.dispose(), s;
+    }
+  }
+  // Evaluate model on validation dataset - FIXED memory leaks
+  async evaluateOnDataset(t, e = 100) {
+    let s = 0, a = 0;
+    return await t.take(e).forEachAsync(async (o) => {
+      const { xs: r, ys: n } = o, { loss: i, logits: l } = this.model.forward(r, n, !1), d = i.arraySync();
+      i.dispose(), l.dispose(), s += d, a++;
+    }), s / a;
+  }
+  // Create training and validation datasets - FIXED memory leaks
+  async createTrainValidationSplit(t, e = 32, s = 0.1) {
+    const a = Math.floor(t.length * (1 - s)), o = t.slice(0, a), r = t.slice(a), n = await this.datasetBuilder.createTextDataset(o, e), i = await this.datasetBuilder.createTextDataset(r, e);
+    return { trainDataset: n, validationDataset: i };
+  }
+  async createDataset(t, e = 32) {
+    return await this.datasetBuilder.createTextDataset(t, e);
+  }
+  dispose() {
+    this.optimizer && this.optimizer.dispose(), this.tf.dispose();
+  }
+}
+export {
+  y as default
+};

package/dist/training/lwSchedule.d.ts

@@ -0,0 +1,7 @@
+import { AdamConfig } from './Trainer';
+export interface LWSchedule {
+    adam: AdamConfig;
+    skip: boolean[];
+    trainable: boolean[];
+}
+export declare const schedule: LWSchedule[][];

package/dist/training/lwSchedule.js

@@ -0,0 +1,162 @@
+const e = [
+  [
+    {
+      adam: {
+        learningRateFactor: 1,
+        beta1: 0.9,
+        beta2: 0.999,
+        epsilon: 1e-8
+      },
+      skip: [!1],
+      trainable: [!0]
+    }
+  ],
+  [
+    {
+      adam: {
+        learningRateFactor: 1,
+        beta1: 0.9,
+        beta2: 0.999,
+        epsilon: 1e-8
+      },
+      skip: [!0, !1],
+      trainable: [!1, !0]
+    },
+    {
+      adam: {
+        learningRateFactor: 1,
+        beta1: 0.9,
+        beta2: 0.999,
+        epsilon: 1e-8
+      },
+      skip: [!1, !1],
+      trainable: [!0, !1]
+    },
+    {
+      adam: {
+        learningRateFactor: 0.3333333333333333,
+        beta1: 0.95,
+        beta2: 0.999,
+        epsilon: 1e-8
+      },
+      skip: [!1, !1],
+      trainable: [!0, !0]
+    }
+  ],
+  [],
+  [
+    {
+      adam: {
+        learningRateFactor: 1,
+        beta1: 0.9,
+        beta2: 0.999,
+        epsilon: 1e-8
+      },
+      skip: [!0, !0, !0, !1],
+      trainable: [!1, !1, !1, !0]
+    },
+    {
+      adam: {
+        learningRateFactor: 1,
+        beta1: 0.9,
+        beta2: 0.999,
+        epsilon: 1e-8
+      },
+      skip: [!0, !0, !1, !1],
+      trainable: [!1, !1, !0, !1]
+    },
+    {
+      adam: {
+        learningRateFactor: 0.3333333333333333,
+        beta1: 0.95,
+        beta2: 0.999,
+        epsilon: 1e-8
+      },
+      skip: [!0, !0, !1, !1],
+      trainable: [!1, !1, !1, !0]
+    },
+    {
+      adam: {
+        learningRateFactor: 1,
+        beta1: 0.9,
+        beta2: 0.999,
+        epsilon: 1e-8
+      },
+      skip: [!0, !1, !1, !1],
+      trainable: [!1, !0, !1, !1]
+    },
+    {
+      adam: {
+        learningRateFactor: 0.3333333333333333,
+        beta1: 0.95,
+        beta2: 0.999,
+        epsilon: 1e-8
+      },
+      skip: [!0, !1, !1, !1],
+      trainable: [!1, !1, !0, !1]
+    },
+    {
+      adam: {
+        learningRateFactor: 0.16666666666666666,
+        beta1: 0.98,
+        beta2: 0.9999,
+        epsilon: 1e-8
+      },
+      skip: [!0, !1, !1, !1],
+      trainable: [!1, !1, !1, !0]
+    },
+    {
+      adam: {
+        learningRateFactor: 1,
+        beta1: 0.9,
+        beta2: 0.999,
+        epsilon: 1e-8
+      },
+      skip: [!1, !1, !1, !1],
+      trainable: [!0, !1, !1, !1]
+    },
+    {
+      adam: {
+        learningRateFactor: 0.3333333333333333,
+        beta1: 0.95,
+        beta2: 0.999,
+        epsilon: 1e-8
+      },
+      skip: [!1, !1, !1, !1],
+      trainable: [!1, !0, !1, !1]
+    },
+    {
+      adam: {
+        learningRateFactor: 0.16666666666666666,
+        beta1: 0.98,
+        beta2: 0.9999,
+        epsilon: 1e-8
+      },
+      skip: [!1, !1, !1, !1],
+      trainable: [!1, !1, !0, !1]
+    },
+    {
+      adam: {
+        learningRateFactor: 0.16666666666666666,
+        beta1: 0.98,
+        beta2: 0.9999,
+        epsilon: 1e-8
+      },
+      skip: [!1, !1, !1, !1],
+      trainable: [!1, !1, !1, !0]
+    },
+    {
+      adam: {
+        learningRateFactor: 0.16666666666666666,
+        beta1: 0.98,
+        beta2: 0.9999,
+        epsilon: 1e-8
+      },
+      skip: [!1, !1, !1, !1],
+      trainable: [!0, !0, !0, !0]
+    }
+  ]
+];
+export {
+  e as schedule
+};

package/dist/utilities/generate.d.ts

@@ -0,0 +1,3 @@
+import { ITokeniser } from '../tokeniser/type';
+import { default as NanoGPT } from '../NanoGPTModel';
+export declare function generateText(tokeniser: ITokeniser, model: NanoGPT, prompt: string, length: number, temperature?: number, topK?: number): Promise<string>;

package/dist/utilities/generate.js

@@ -0,0 +1,22 @@
+async function h(n, t, i, s, a = 1, r) {
+  if (s <= 0)
+    throw new Error("Length must be a positive integer");
+  if (a <= 0)
+    throw new Error("Temperature must be a positive number");
+  if (r !== void 0 && r <= 0)
+    throw new Error("topK must be a positive integer or undefined");
+  if (i.length === 0)
+    throw new Error("Prompt cannot be an empty string");
+  const c = await n.tokenise([i], !0), o = (await t.tf.tidy(() => {
+    let e = t.tf.tensor2d(c, [1, c[0].length], "int32");
+    for (let u = 0; u < s; u++) {
+      const f = t.generate(e, a, r), g = e;
+      e = t.tf.concat([e, f], 1), g.dispose(), f.dispose();
+    }
+    return e;
+  }).array())[0], d = o.indexOf(n.eosToken);
+  return d !== -1 && o.splice(d), await n.decode(o);
+}
+export {
+  h as generateText
+};

package/dist/utilities/load.d.ts

@@ -0,0 +1,7 @@
+import { default as TF } from '@tensorflow/tfjs';
+import { default as NanoGPT } from '../NanoGPTModel';
+import { ITokeniser } from '../tokeniser/type';
+export declare function loadModel(tf: typeof TF, data: Blob | Buffer | string): Promise<{
+    model: NanoGPT;
+    tokeniser: ITokeniser;
+}>;

package/dist/utilities/load.js

@@ -0,0 +1,47 @@
+import { z as k } from "../jszip.min-BLbRbbKt.js";
+import { importWeights as F } from "./weights.js";
+import z from "../tokeniser/CharTokeniser.js";
+import j from "../NanoGPTModel.js";
+function m(o) {
+  const a = o.tf.zeros([1, o.config.blockSize], "int32"), { logits: n, loss: s } = o.forward(a, void 0, !1);
+  n.dispose(), s && s.dispose(), a.dispose();
+}
+async function E(o) {
+  const e = await fetch(o);
+  if (!e.ok)
+    throw new Error(`Failed to fetch ${o}: ${e.statusText}`);
+  return e.arrayBuffer();
+}
+async function A(o, e) {
+  const a = typeof e == "string" ? await E(e) : e, n = await k.loadAsync(a), s = /* @__PURE__ */ new Map(), f = await n.file("manifest.json")?.async("string");
+  if (!f)
+    throw new Error("Manifest file not found in the zip archive");
+  const l = JSON.parse(f);
+  for (const [t, r] of Object.entries(l.weightSpec))
+    s.set(t, { spec: r, data: new Float32Array() });
+  const p = await n.file("tokeniser.json")?.async("string");
+  if (!p)
+    throw new Error("Tokeniser file not found in the zip archive");
+  const d = JSON.parse(p), y = new z(d.vocab), w = /* @__PURE__ */ new Map();
+  for (const t of Object.keys(n.files))
+    if (t.endsWith(".bin")) {
+      const r = t.replace(".bin", ""), h = await n.file(t).async("arraybuffer"), u = new Float32Array(h), c = s.get(r) || { spec: [], data: new Float32Array() };
+      c.data = u, s.set(r, c);
+      const b = await F(c, o);
+      w.set(r, b);
+    }
+  const i = new j(o, l.config);
+  m(i), i.loadWeights(w), m(i);
+  const g = await n.file("log.json")?.async("string");
+  if (g)
+    try {
+      const t = JSON.parse(g);
+      i.log = t;
+    } catch (t) {
+      throw console.error("Error parsing training log:", t), new Error(`Failed to parse training log: ${t}`);
+    }
+  return { model: i, tokeniser: y };
+}
+export {
+  A as loadModel
+};

package/dist/utilities/save.d.ts

@@ -0,0 +1,3 @@
+import { default as NanoGPT } from '../NanoGPTModel';
+import { ITokeniser } from '../tokeniser/type';
+export declare function saveModel(model: NanoGPT, tokeniser: ITokeniser): Promise<Blob>;

package/dist/utilities/save.js

@@ -0,0 +1,21 @@
+import { z as f } from "../jszip.min-BLbRbbKt.js";
+import { exportWeights as g } from "./weights.js";
+async function l(i, t) {
+  const o = i.saveWeights(), e = new f(), s = {};
+  for (const [n, r] of o) {
+    const a = await g(r);
+    s[n] = a.spec, e.file(`${n}.bin`, a.data.buffer, { binary: !0 });
+  }
+  return e.file("manifest.json", JSON.stringify({ weightSpec: s, config: i.config }), {
+    binary: !1
+  }), e.file(
+    "tokeniser.json",
+    JSON.stringify({ vocab: t.getVocab(), merges: await t.getMerges() }),
+    {
+      binary: !1
+    }
+  ), e.file("log.json", JSON.stringify(i.log), { binary: !1 }), e.generateAsync({ type: "blob" });
+}
+export {
+  l as saveModel
+};

package/dist/utilities/textLoader.d.ts

@@ -0,0 +1 @@
+export default function loadTextData(file: File | string): Promise<string[]>;