@genai-fi/nanogpt 0.7.2 → 0.8.0

package/dist/layers/BaseLayer.js

@@ -1,22 +1,28 @@
-import { T as g, y as p, e as o, A as v } from "../index-BoWRt-10.js";
-import { v as _ } from "../variable-BuddVFLa.js";
-class M {
+import { T as p, I as g, e as o, J as v } from "../index-CUQrfsw_.js";
+import { v as _ } from "../variable-C9hihzDB.js";
+class T {
   parent;
   config;
   _variables = /* @__PURE__ */ new Map();
   _trainable = !0;
   children = [];
+  profiler;
   constructor(t, r) {
     this.config = t, this.parent = r, this.parent && this.parent.children.push(this);
   }
   getProfiler() {
-    return this.config.layerConfig.profiler;
+    return this.profiler;
+  }
+  setProfiler(t) {
+    this.profiler = t || void 0, this.children.forEach((r) => {
+      r.setProfiler(t);
+    });
   }
   startMemory() {
-    this.config.layerConfig.profiler?.startMemory();
+    this.profiler?.startMemory();
   }
   endMemory(t) {
-    this.config.layerConfig.profiler?.endMemory(t);
+    this.profiler?.endMemory(t);
   }
   addVariable(t, r) {
     this._variables.set(t, r || null);
@@ -41,11 +47,17 @@ class M {
       r.trainable = t;
     });
   }
-  getVariable(t) {
-    const r = this._variables.get(t);
-    if (!r)
+  getVariable(t, r = !1) {
+    const e = this._variables.get(t);
+    if (!e && r)
+      for (const i of this.children) {
+        const s = i.getVariable(t, !0);
+        if (s)
+          return s;
+      }
+    if (!e)
       throw new Error(`Variable ${t} not found`);
-    return r;
+    return e;
   }
   hasVariable(t) {
     return this._variables.get(t) !== null;
@@ -85,7 +97,7 @@ class M {
   call(t, ...r) {
     this.build();
     const e = this.forward(t, ...r);
-    if (t.training && e instanceof g) {
+    if (t.training && e instanceof p) {
       const i = this.dropout(e);
       return i !== e && e.dispose(), i;
     } else
@@ -95,7 +107,7 @@ class M {
     return this.build(), this.checkpointingFn(t, ...r);
   }
   checkpointingFn(t, ...r) {
-    const e = this.trainableVariables, s = p((...a) => {
+    const e = this.trainableVariables, s = g((...a) => {
       const l = a[a.length - 1], n = a.slice(0, r.length), h = this.forward(t, ...n);
       return l(n), { value: h, gradFunc: (c, f) => {
         const u = o().state.activeTape;
@@ -112,5 +124,5 @@ class M {
   }
 }
 export {
-  M as default
+  T as default
 };

package/dist/layers/CausalSelfAttention.d.ts

@@ -1,5 +1,6 @@
-import { default as BaseLayer, ForwardAttributes, GPTLayerConfig } from './BaseLayer';
+import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
 import { Tensor } from '@tensorflow/tfjs-core';
+import { GPTConfig } from '../models/config';
 export type KVCache = {
     k?: Tensor;
     v?: Tensor;
@@ -22,7 +23,7 @@ export default class CausalSelfAttention extends BaseLayer<AttentionForwardAttri
     private projUnits;
     private ATTN;
     private PROJ;
-    constructor(index: number, config: GPTLayerConfig, parent?: BaseLayer);
+    constructor(index: number, config: GPTConfig, parent?: BaseLayer);
     protected build(): void;
     private getAttentionScores;
     private getAttentionScoresWithPast;

package/dist/layers/CausalSelfAttention.js

@@ -3,14 +3,14 @@ import O from "./BaseLayer.js";
 import { qkv as P } from "../ops/qkv.js";
 import { rope as v } from "../ops/rope.js";
 import { appendCache as V } from "../ops/appendCache.js";
-import { w as c, t as C } from "../index-BoWRt-10.js";
+import { k as c, t as C } from "../index-CUQrfsw_.js";
 import { fusedSoftmax as T } from "../ops/fusedSoftmax.js";
-import { d as y } from "../random_width-sZORGo5k.js";
-import { v as b } from "../variable-BuddVFLa.js";
-import { r as k, d as L } from "../dropout-DYs5QFGQ.js";
-import { r as N } from "../reshape-CdBq1WJ6.js";
-import { m as R } from "../mat_mul-8m8pfdcx.js";
-class W extends O {
+import { d as L } from "../random_width-D8Pwy_na.js";
+import { v as b } from "../variable-C9hihzDB.js";
+import { r as k, d as y } from "../dropout-sx0sjVAT.js";
+import { r as N } from "../reshape-DEfQGSin.js";
+import { m as R } from "../mat_mul-DMkduNJu.js";
+class $ extends O {
   divisor;
   index;
   units;
@@ -18,27 +18,27 @@ class W extends O {
   ATTN;
   PROJ;
   constructor(t, i, s) {
-    super(i, s), this.index = t, this.units = i.gpt.nEmbed * 3, this.projUnits = i.gpt.nEmbed, this.ATTN = `block_${this.index}_cAttn`, this.PROJ = `block_${this.index}_cProj`, this.addVariable(this.ATTN), this.addVariable(this.PROJ), this.divisor = 1 / Math.sqrt(i.gpt.nEmbed / i.gpt.nHead);
+    super(i, s), this.index = t, this.units = i.nEmbed * 3, this.projUnits = i.nEmbed, this.ATTN = `block_${this.index}_cAttn`, this.PROJ = `block_${this.index}_cProj`, this.addVariable(this.ATTN), this.addVariable(this.PROJ), this.divisor = 1 / Math.sqrt(i.nEmbed / i.nHead);
   }
   build() {
     this.hasVariable(this.ATTN) === !1 && this.setVariable(
       this.ATTN,
       b(
-        k([this.config.gpt.nEmbed, this.units], 0, 0.02),
+        k([this.config.nEmbed, this.units], 0, 0.02),
         !0
         //`block_${this.index}_attn_cAttn_kernel`
       )
     ), this.hasVariable(this.PROJ) === !1 && this.setVariable(
       this.PROJ,
       b(
-        k([this.projUnits, this.config.gpt.nEmbed], 0, 0.02),
+        k([this.projUnits, this.config.nEmbed], 0, 0.02),
         !0
         //`block_${this.index}_attn_cProj_kernel`
       )
     );
   }
   getAttentionScores(t, i, s, o) {
-    const e = g(t, i, this.divisor), n = T(e, s ? this.config.gpt.dropout : 0, o);
+    const e = g(t, i, this.divisor), n = T(e, s ? this.config.dropout : 0, o);
     return e.dispose(), n;
   }
   // Attention with optional past. If pastLen > 0 and T_cur == 1, no mask needed.
@@ -47,50 +47,50 @@ class W extends O {
     return o.dispose(), e;
   }
   getQKV(t) {
-    return P(t, this.getVariable(this.ATTN), this.config.gpt.nHead);
+    return P(t, this.getVariable(this.ATTN), this.config.nHead);
   }
   getOutputProjection(t) {
-    const i = t.shape[0], s = t.shape[2], o = this.config.gpt.nEmbed, e = t.transpose([0, 2, 1, 3]), n = N(e, [i, s, o]), p = y(n, this.getVariable(this.PROJ));
-    return n.dispose(), e.dispose(), p;
+    const i = t.shape[0], s = t.shape[2], o = this.config.nEmbed, e = t.transpose([0, 2, 1, 3]), n = N(e, [i, s, o]), r = L(n, this.getVariable(this.PROJ));
+    return n.dispose(), e.dispose(), r;
   }
   updateCache(t, i, s) {
-    const o = this.config.gpt.blockSize, e = t.shape[2], n = s.length || 0, p = V(t, o, n, s.k);
+    const o = this.config.blockSize, e = t.shape[2], n = s.length || 0, r = V(t, o, n, s.k);
     t.dispose(), s.k && s.k.dispose();
-    const a = V(i, o, n, s.v);
+    const p = V(i, o, n, s.v);
     i.dispose(), s.v && s.v.dispose();
     const d = Math.min(n + e, o), h = s.cumulativeLength + e;
-    s.length = d, s.cumulativeLength = h, s.k = c(p), s.v = c(a);
+    s.length = d, s.cumulativeLength = h, s.k = c(r), s.v = c(p);
   }
   forward(t, i) {
     return C(() => {
       this.startMemory();
-      const [s, o, e] = this.getQKV(i), n = t.pastKV ? t.pastKV.cumulativeLength : 0, p = this.config.layerConfig.ropeCache, a = p ? v(s, p, n) : s, d = p ? v(o, p, n) : o;
-      p && (s.dispose(), o.dispose());
+      const [s, o, e] = this.getQKV(i), n = t.pastKV ? t.pastKV.cumulativeLength : 0, r = t.ropeCache, p = r ? v(s, r, n) : s, d = r ? v(o, r, n) : o;
+      r && (s.dispose(), o.dispose());
       const h = t.pastKV ? t.pastKV.length : 0;
       t.pastKV && !t.training && this.updateCache(d, e, t.pastKV);
       const u = t.pastKV?.k ? t.pastKV.k : d, m = t.pastKV?.v ? t.pastKV.v : e;
-      let r;
-      h > 0 ? r = this.getAttentionScoresWithPast(a, u, h) : r = this.getAttentionScores(a, u, t.training, t.seed || 0), a.dispose(), t.pastKV || u.dispose();
-      const l = R(r, m), f = t.attentionScores !== void 0 && t.attentionScores.attentionOut !== void 0;
-      f || r.dispose(), t.pastKV || m.dispose();
+      let a;
+      h > 0 ? a = this.getAttentionScoresWithPast(p, u, h) : a = this.getAttentionScores(p, u, t.training, t.seed || 0), p.dispose(), t.pastKV || u.dispose();
+      const l = R(a, m), f = t.attentionScores !== void 0 && t.attentionScores.attentionOut !== void 0;
+      f || a.dispose(), t.pastKV || m.dispose();
       const A = this.getOutputProjection(l);
       if (l.dispose(), f && t.attentionScores && t.attentionScores.attentionOut !== void 0) {
-        const K = r.shape[1], S = r.shape[2];
+        const K = a.shape[1], S = a.shape[2];
         t.attentionScores.attentionOut?.push(
-          c(r.slice([0, 0, 0, 0], [1, -1, -1, -1]).reshape([K, S, -1]))
+          c(a.slice([0, 0, 0, 0], [1, -1, -1, -1]).reshape([K, S, -1]))
         );
       }
       return this.endMemory("CausalSelfAttention"), A;
     });
   }
   dropout(t) {
-    if (this.config.gpt.dropout > 0) {
-      const i = L(t, this.config.gpt.dropout);
+    if (this.config.dropout > 0) {
+      const i = y(t, this.config.dropout);
       return t.dispose(), i;
     } else
       return t;
   }
 }
 export {
-  W as default
+  $ as default
 };

package/dist/layers/MLP.d.ts

@@ -1,11 +1,12 @@
 import { Tensor } from '@tensorflow/tfjs-core';
-import { default as BaseLayer, ForwardAttributes, GPTLayerConfig } from './BaseLayer';
+import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
+import { GPTConfig } from '../main';
 export default class MLP extends BaseLayer {
     private index;
     private hiddenUnits;
     private MLPHIDDEN;
     private MLPOUT;
-    constructor(index: number, config: GPTLayerConfig, parent?: BaseLayer);
+    constructor(index: number, config: GPTConfig, parent?: BaseLayer);
     protected build(): void;
     forward(_: ForwardAttributes, x: Tensor): Tensor;
     protected dropout(x: Tensor): Tensor;

package/dist/layers/MLP.js

@@ -1,56 +1,52 @@
-import { t as l } from "../index-BoWRt-10.js";
+import { t as p } from "../index-CUQrfsw_.js";
 import u from "./BaseLayer.js";
 import { matMulGelu as M } from "../ops/matMulGelu.js";
-import { v as o } from "../variable-BuddVFLa.js";
-import { r as h, d as f } from "../dropout-DYs5QFGQ.js";
-import { r as d } from "../reshape-CdBq1WJ6.js";
-import { m as c } from "../mat_mul-8m8pfdcx.js";
-class V extends u {
+import { v as o } from "../variable-C9hihzDB.js";
+import { r as h, d as f } from "../dropout-sx0sjVAT.js";
+import { r as d } from "../reshape-DEfQGSin.js";
+import { m as c } from "../mat_mul-DMkduNJu.js";
+class H extends u {
   index;
   hiddenUnits;
   MLPHIDDEN;
   MLPOUT;
   constructor(i, t, s) {
-    super(t, s), this.index = i, this.hiddenUnits = t.gpt.mlpFactor * t.gpt.nEmbed, this.MLPHIDDEN = `block_${this.index}_mlpHidden`, this.MLPOUT = `block_${this.index}_mlpOut`, this.addVariable(this.MLPHIDDEN), this.addVariable(this.MLPOUT);
+    super(t, s), this.index = i, this.hiddenUnits = t.mlpFactor * t.nEmbed, this.MLPHIDDEN = `block_${this.index}_mlpHidden`, this.MLPOUT = `block_${this.index}_mlpOut`, this.addVariable(this.MLPHIDDEN), this.addVariable(this.MLPOUT);
   }
   build() {
     this.hasVariable(this.MLPHIDDEN) === !1 && this.setVariable(
       this.MLPHIDDEN,
       o(
-        h([this.config.gpt.nEmbed, this.hiddenUnits], 0, 0.02),
+        h([this.config.nEmbed, this.hiddenUnits], 0, 0.02),
         !0
         //`block_${this.index}_attn_cAttn_kernel`
       )
     ), this.hasVariable(this.MLPOUT) === !1 && this.setVariable(
       this.MLPOUT,
       o(
-        h(
-          [this.hiddenUnits, this.config.gpt.nEmbed],
-          0,
-          0.02 / Math.sqrt(2 * this.config.gpt.nLayer)
-        ),
+        h([this.hiddenUnits, this.config.nEmbed], 0, 0.02 / Math.sqrt(2 * this.config.nLayer)),
         !0
         //`block_${this.index}_attn_cProj_kernel`
       )
     );
   }
   forward(i, t) {
-    return l(() => {
+    return p(() => {
       this.startMemory();
-      const [s, r, e] = t.shape, n = d(t, [s * r, e]), a = M(n, this.getVariable(this.MLPHIDDEN)), p = c(a, this.getVariable(this.MLPOUT));
+      const [s, r, e] = t.shape, n = d(t, [s * r, e]), a = M(n, this.getVariable(this.MLPHIDDEN)), m = c(a, this.getVariable(this.MLPOUT));
       a.dispose();
-      const m = d(p, [s, r, e]);
-      return this.endMemory("MLP"), m;
+      const l = d(m, [s, r, e]);
+      return this.endMemory("MLP"), l;
     });
   }
   dropout(i) {
-    if (this.config.gpt.dropout > 0) {
-      const t = f(i, this.config.gpt.dropout);
+    if (this.config.dropout > 0) {
+      const t = f(i, this.config.dropout);
       return i.dispose(), t;
     }
     return i;
   }
 }
 export {
-  V as default
+  H as default
 };

package/dist/layers/PositionEmbedding.d.ts

@@ -0,0 +1,9 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+import { default as BaseLayer } from './BaseLayer';
+import { GPTConfig, ModelForwardAttributes } from '../main';
+export default class PositionEmbedding extends BaseLayer {
+    private wpe?;
+    private drop;
+    constructor(config: GPTConfig, name?: string, parent?: BaseLayer);
+    forward(attrs: ModelForwardAttributes, x: Tensor): Tensor;
+}

package/dist/layers/PositionEmbedding.js

@@ -0,0 +1,45 @@
+import { t as c, a8 as u, b as i } from "../index-CUQrfsw_.js";
+import f from "./BaseLayer.js";
+import { E as g, D as h } from "../random_width-D8Pwy_na.js";
+import { r as b } from "../exports_initializers-DAKM8UO9.js";
+import { m as l } from "../mod-uUuj4gSb.js";
+import { r as w } from "../range-LVHrSLdi.js";
+/**
+ * @license
+ * Copyright 2018 Google LLC
+ *
+ * Use of this source code is governed by an MIT-style
+ * license that can be found in the LICENSE file or at
+ * https://opensource.org/licenses/MIT.
+ * =============================================================================
+ */
+function E(t) {
+  return new h(t);
+}
+function x(t) {
+  return new g(t);
+}
+class q extends f {
+  wpe;
+  // Position embeddings
+  drop;
+  // Dropout
+  constructor(o, n = "", r) {
+    super(o, r), this.wpe = x({
+      inputDim: this.config.blockSize,
+      outputDim: this.config.nEmbed,
+      name: n,
+      embeddingsInitializer: b({ mean: 0, stddev: 0.02 })
+    }), this.drop = E({ rate: this.config.dropout });
+  }
+  forward(o, n) {
+    const r = o.cache?.[0]?.length ?? 0;
+    return c(() => {
+      const [, s] = n.shape, e = this.config.blockSize, a = w(0, s, 1, "int32"), m = l(u(a, i(r, "int32")), i(e, "int32")), d = this.wpe.apply(m), p = n.add(d);
+      return this.drop.apply(p, { training: o.training });
+    });
+  }
+}
+export {
+  q as default
+};

package/dist/layers/RMSNorm.d.ts

@@ -1,7 +1,8 @@
 import { Tensor } from '@tensorflow/tfjs-core';
-import { default as BaseLayer, ForwardAttributes, GPTLayerConfig } from './BaseLayer';
+import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
+import { GPTConfig } from '../main';
 export default class RMSNorm extends BaseLayer {
     private GAMMA;
-    constructor(config: GPTLayerConfig, name?: string, parent?: BaseLayer);
+    constructor(config: GPTConfig, name?: string, parent?: BaseLayer);
     forward(_: ForwardAttributes, x: Tensor): Tensor;
 }

package/dist/layers/RMSNorm.js

@@ -1,12 +1,12 @@
-import { t as s } from "../index-BoWRt-10.js";
+import { t as s } from "../index-CUQrfsw_.js";
 import e from "./BaseLayer.js";
 import { normRMS as a } from "../ops/normRMS.js";
-import { v as i } from "../variable-BuddVFLa.js";
-import { o as m } from "../ones-Dj0SDhHf.js";
-class f extends e {
+import { v as i } from "../variable-C9hihzDB.js";
+import { o as m } from "../ones-ZdgQGBCP.js";
+class l extends e {
   GAMMA;
   constructor(r, t = "", o) {
-    super(r, o), this.GAMMA = t, this.addVariable(this.GAMMA, i(m([r.gpt.nEmbed]), !0, this.GAMMA, "float32"));
+    super(r, o), this.GAMMA = t, this.addVariable(this.GAMMA, i(m([r.nEmbed]), !0, this.GAMMA, "float32"));
   }
   forward(r, t) {
     return s(() => {
@@ -17,5 +17,5 @@ class f extends e {
   }
 }
 export {
-  f as default
+  l as default
 };

package/dist/layers/RoPECache.d.ts

@@ -1,5 +1,5 @@
 import { Tensor } from '@tensorflow/tfjs-core';
-import { GPTConfig } from '../config';
+import { GPTConfig } from '../models/config';
 export default class RoPECache {
     private readonly config;
     readonly rotaryDim: number;

package/dist/layers/RoPECache.js

@@ -1,7 +1,7 @@
-import { b as t, x as h, t as n, w as p } from "../index-BoWRt-10.js";
-import { r as c } from "../reciprocal-BvGAyKyu.js";
-import { c as f, s as m } from "../sin-BeA3tsEd.js";
-import { r as a } from "../range-CRuAh-gd.js";
+import { b as t, x as h, t as n, k as p } from "../index-CUQrfsw_.js";
+import { r as c } from "../reciprocal-CaR9e67G.js";
+import { c as f, s as m } from "../sin-D2CKKmyR.js";
+import { r as a } from "../range-LVHrSLdi.js";
 class D {
   constructor(o) {
     this.config = o;

package/dist/layers/TiedEmbedding.d.ts

@@ -1,11 +1,12 @@
 import { Tensor } from '@tensorflow/tfjs-core';
-import { default as BaseLayer, ForwardAttributes, GPTLayerConfig } from './BaseLayer';
+import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
+import { GPTConfig } from '../models/config';
 export default class TiedEmbeddingOutputLayer extends BaseLayer {
     private vocabSize;
     private embedDim;
     private initializer;
     private WEIGHTS;
-    constructor(config: GPTLayerConfig, name: string, parent?: BaseLayer);
+    constructor(config: GPTConfig, name: string, parent?: BaseLayer);
     embed(inputs: Tensor): Tensor;
     project(inputs: Tensor): Tensor;
     forward(_: ForwardAttributes, x: Tensor): Tensor;

package/dist/layers/TiedEmbedding.js

@@ -1,9 +1,31 @@
-import "../random_width-sZORGo5k.js";
-import "../index-BoWRt-10.js";
-import { T as e } from "../TiedEmbedding-BxOerUmB.js";
-import "./BaseLayer.js";
-import "../variable-BuddVFLa.js";
-import "../gather-CMMy2KEG.js";
+import { d as r } from "../random_width-D8Pwy_na.js";
+import "../index-CUQrfsw_.js";
+import { r as a } from "../exports_initializers-DAKM8UO9.js";
+import s from "./BaseLayer.js";
+import { v as m } from "../variable-C9hihzDB.js";
+import { g as o } from "../gather-C1siEkdp.js";
+class S extends s {
+  vocabSize;
+  embedDim;
+  initializer;
+  WEIGHTS;
+  constructor(i, e, t) {
+    super(i, t), this.WEIGHTS = e, this.vocabSize = i.vocabSize, this.embedDim = i.nEmbed, this.initializer = a({
+      mean: 0,
+      stddev: 0.02
+    }), this.addVariable(this.WEIGHTS, m(this.initializer.apply([this.vocabSize, this.embedDim]), !0));
+  }
+  embed(i) {
+    return o(this.getVariable(this.WEIGHTS), i, 0);
+  }
+  project(i) {
+    return r(i, this.getVariable(this.WEIGHTS).transpose());
+  }
+  // Dummy, should not be used.
+  forward(i, e) {
+    return this.project(e);
+  }
+}
 export {
-  e as default
+  S as default
 };

package/dist/layers/TransformerBlock.d.ts

@@ -1,6 +1,7 @@
 import { AttentionScores, KVCache } from './CausalSelfAttention';
-import { default as BaseLayer, ForwardAttributes, GPTLayerConfig } from './BaseLayer';
+import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
 import { Tensor } from '@tensorflow/tfjs-core';
+import { GPTConfig } from '../models/config';
 interface BlockAttributes extends ForwardAttributes {
     pastKV?: KVCache;
     seed?: number;
@@ -13,7 +14,7 @@ export default class Block extends BaseLayer<BlockAttributes> {
     private mlp;
     private index;
     skipped: boolean;
-    constructor(index: number, config: GPTLayerConfig, parent?: BaseLayer);
+    constructor(index: number, config: GPTConfig, parent?: BaseLayer);
     private getMLPOutput;
     forward(attrs: BlockAttributes, x: Tensor): Tensor;
     dispose(): void;

package/dist/layers/TransformerBlock.js

@@ -2,7 +2,7 @@ import l from "./CausalSelfAttention.js";
 import r from "./MLP.js";
 import o from "./RMSNorm.js";
 import d from "./BaseLayer.js";
-import { t as p } from "../index-BoWRt-10.js";
+import { t as p } from "../index-CUQrfsw_.js";
 class k extends d {
   ln1;
   attn;

package/dist/loader/load.d.ts

@@ -1,5 +1,5 @@
-import { default as NanoGPT } from '../NanoGPTModel';
 import { ITokeniser } from '../tokeniser/type';
+import { default as Model, ModelForwardAttributes } from '../models/model';
 export declare const VERSION = 2;
 export interface Metadata {
     version: string;
@@ -7,7 +7,7 @@ export interface Metadata {
     name?: string;
 }
 export declare function loadModel(data: Blob | Buffer | string): Promise<{
-    model: NanoGPT;
+    model: Model<ModelForwardAttributes>;
     tokeniser: ITokeniser;
     name?: string;
 }>;

package/dist/loader/loadHF.d.ts

@@ -1,7 +1,7 @@
-import { default as NanoGPT } from '../NanoGPTModel';
 import { ITokeniser } from '../main';
+import { default as Model, ModelForwardAttributes } from '../models/model';
 export default function loadHuggingFace(name: string): Promise<{
-    model: NanoGPT;
+    model: Model<ModelForwardAttributes>;
     tokeniser: ITokeniser;
     name?: string;
 }>;

package/dist/loader/loadTransformers.d.ts

@@ -1,5 +1,6 @@
 import { ITokeniser } from '../tokeniser/type';
-import { default as NanoGPT } from '../NanoGPTModel';
+import { default as Model, ModelForwardAttributes } from '../models/model';
+import { TrainingState } from '../training/Trainer';
 export interface TransformersConfig {
     model_type: string;
     vocab_size: number;
@@ -22,10 +23,11 @@ export interface TransformersMetadata {
     name?: string;
     version: number;
     application: string;
+    training?: TrainingState;
     [key: string]: unknown;
 }
 export default function loadTransformers(config: TransformersConfig, tokeniser: TransformersTokeniser, metadata: TransformersMetadata, weightData: ArrayBuffer): Promise<{
-    model: NanoGPT;
+    model: Model<ModelForwardAttributes>;
     tokeniser: ITokeniser;
     name?: string;
 }>;

package/dist/loader/loadTransformers.js

@@ -1,11 +1,12 @@
-import b from "../NanoGPTModel.js";
-import c from "../tokeniser/CharTokeniser.js";
-import l from "../tokeniser/bpe.js";
-import { load_safetensors as u } from "../utilities/safetensors.js";
-import { a0 as y } from "../index-BoWRt-10.js";
-import { dummyPassAsync as h } from "../utilities/dummy.js";
+import l from "../tokeniser/CharTokeniser.js";
+import c from "../tokeniser/bpe.js";
+import { load_safetensors as b } from "../utilities/safetensors.js";
+import { a1 as y } from "../index-CUQrfsw_.js";
+import { dummyPassAsync as u } from "../utilities/dummy.js";
+import _ from "../models/factory.js";
 async function L(e, a, r, t) {
   const n = {
+    modelType: e.model_type || "GenAI_NanoGPT_v1",
     vocabSize: e.vocab_size,
     blockSize: e.block_size,
     nLayer: e.num_hidden_layers,
@@ -16,12 +17,12 @@ async function L(e, a, r, t) {
     biasInLayerNorm: e.biasInLayerNorm,
     mlpFactor: e.mlpFactor,
     useRope: e.useRope
-  }, m = (a.type ?? "char") === "char" ? new c(a.vocab) : new l(a.vocab, a.merges), i = await u(t), s = /* @__PURE__ */ new Map();
+  }, m = (a.type ?? "char") === "char" ? new l(a.vocab) : new c(a.vocab, a.merges), i = await b(t), s = /* @__PURE__ */ new Map();
   for (const [p, d] of Object.entries(i))
     s.set(p, [d]);
   y();
-  const o = new b(n);
-  return await h(o), o.loadWeights(s), { model: o, tokeniser: m, name: r.name };
+  const o = _(n);
+  return await u(o), o.loadWeights(s), { model: o, tokeniser: m, name: r.name };
 }
 export {
   L as default

package/dist/loader/newZipLoad.d.ts

@@ -1,8 +1,8 @@
 import { ITokeniser } from '../main';
-import { default as NanoGPT } from '../NanoGPTModel';
 import { default as zip } from 'jszip';
+import { default as Model, ModelForwardAttributes } from '../models/model';
 export default function loadZipFile(zipFile: zip): Promise<{
-    model: NanoGPT;
+    model: Model<ModelForwardAttributes>;
     tokeniser: ITokeniser;
     name?: string;
 }>;

package/dist/loader/oldZipLoad.d.ts

@@ -1,7 +1,7 @@
 import { default as zip } from 'jszip';
 import { ITokeniser } from '../main';
-import { default as NanoGPT } from '../NanoGPTModel';
+import { default as Model, ModelForwardAttributes } from '../models/model';
 export default function loadOldModel(zipFile: zip): Promise<{
-    model: NanoGPT;
+    model: Model<ModelForwardAttributes>;
     tokeniser: ITokeniser;
 }>;