@genai-fi/nanogpt 0.0.1

package/dist/layers/CausalSelfAttention.d.ts

@@ -0,0 +1,22 @@
+import { default as TF } from '@tensorflow/tfjs';
+import { GPTConfig } from '../config';
+export default class CausalSelfAttention {
+    private config;
+    private cAttn;
+    private cProj;
+    private attnDropout;
+    private residDropout;
+    private bias;
+    private maskInf;
+    private tf;
+    private divisor;
+    private index;
+    private _trainable;
+    constructor(tf: typeof TF, index: number, config: GPTConfig);
+    get variables(): TF.Variable[];
+    get trainable(): boolean;
+    set trainable(value: boolean);
+    saveWeights(map: Map<string, TF.Tensor[]>): void;
+    loadWeights(weights: Map<string, TF.Tensor[]>): void;
+    call(x: TF.Tensor, training?: boolean): TF.Tensor;
+}

package/dist/layers/CausalSelfAttention.js

@@ -0,0 +1,75 @@
+class g {
+  config;
+  cAttn;
+  cProj;
+  attnDropout;
+  residDropout;
+  bias;
+  maskInf;
+  tf;
+  divisor;
+  index;
+  _trainable = !0;
+  constructor(s, e, t) {
+    this.config = t, this.tf = s, this.index = e, this.cAttn = this.tf.layers.dense({
+      units: 3 * t.nEmbed,
+      useBias: t.biasInLinear,
+      name: `block_${e}_attn_cAttn`,
+      kernelInitializer: this.tf.initializers.randomNormal({
+        mean: 0,
+        stddev: 0.02
+      }),
+      biasInitializer: "zeros"
+    }), this.cProj = this.tf.layers.dense({
+      units: t.nEmbed,
+      useBias: t.biasInLinear,
+      name: `block_${e}_attn_cProj`,
+      kernelInitializer: this.tf.initializers.randomNormal({
+        mean: 0,
+        stddev: 0.02 / Math.sqrt(2 * t.nLayer)
+      }),
+      biasInitializer: "zeros"
+    }), this.attnDropout = this.tf.layers.dropout({ rate: t.dropout }), this.residDropout = this.tf.layers.dropout({ rate: t.dropout }), this.bias = this.tf.linalg.bandPart(this.tf.ones([t.blockSize, t.blockSize]), -1, 0).cast("bool"), this.divisor = this.tf.scalar(1 / Math.sqrt(t.nEmbed / t.nHead)), this.maskInf = this.tf.zeros([t.blockSize, t.blockSize]).where(this.bias, -1 / 0);
+  }
+  get variables() {
+    return [
+      ...this.cAttn.trainableWeights.map((s) => s.read()),
+      ...this.cProj.trainableWeights.map((s) => s.read())
+    ];
+  }
+  get trainable() {
+    return this._trainable;
+  }
+  set trainable(s) {
+    this._trainable = s, this.cAttn.trainable = s, this.cProj.trainable = s;
+  }
+  saveWeights(s) {
+    s.set(`block_${this.index}_cAttn`, this.cAttn.getWeights()), s.set(`block_${this.index}_cProj`, this.cProj.getWeights());
+  }
+  loadWeights(s) {
+    this.cAttn.setWeights(s.get(`block_${this.index}_cAttn`) || []), this.cProj.setWeights(s.get(`block_${this.index}_cProj`) || []);
+  }
+  call(s, e = !1) {
+    return this.tf.tidy(() => {
+      const [t, i, n] = s.shape, r = this.cAttn.apply(s), [o, h, l] = this.tf.split(r, 3, -1);
+      r.dispose();
+      const a = n / this.config.nHead, d = this.tf.reshape(o, [t, i, this.config.nHead, a]);
+      o.dispose();
+      const b = d.transpose([0, 2, 1, 3]);
+      d.dispose();
+      const c = this.tf.reshape(h, [t, i, this.config.nHead, a]);
+      h.dispose();
+      const u = c.transpose([0, 2, 1, 3]);
+      c.dispose();
+      const p = this.tf.reshape(l, [t, i, this.config.nHead, a]);
+      l.dispose();
+      const f = p.transpose([0, 2, 1, 3]);
+      p.dispose();
+      const m = this.tf.matMul(b, u, !1, !0).mul(this.divisor), k = this.maskInf.slice([0, 0], [i, i]), _ = m.add(k), y = this.tf.softmax(_, -1), z = this.attnDropout.apply(y, { training: e }), A = this.tf.matMul(z, f).transpose([0, 2, 1, 3]), P = this.tf.reshape(A, [t, i, n]), j = this.cProj.apply(P);
+      return this.residDropout.apply(j, { training: e });
+    });
+  }
+}
+export {
+  g as default
+};

package/dist/layers/LayerNorm.d.ts

@@ -0,0 +1,12 @@
+import { default as TF } from '@tensorflow/tfjs';
+export default class LayerNorm {
+    private gamma;
+    private epsilon;
+    private tf;
+    constructor(tf: typeof TF, shape: number[], epsilon?: number, name?: string);
+    get trainableWeights(): TF.Variable[];
+    set trainable(value: boolean);
+    getWeights(): TF.Tensor[];
+    setWeights(weights: TF.Tensor[]): void;
+    apply(x: TF.Tensor): TF.Tensor;
+}

package/dist/layers/LayerNorm.js

@@ -0,0 +1,30 @@
+class u {
+  gamma;
+  //private beta: TF.Variable;
+  epsilon;
+  tf;
+  constructor(a, s, t = 1e-5, e = "") {
+    this.tf = a, this.epsilon = t, this.gamma = a.variable(a.ones(s), !0, `${e}_gamma`, "float32");
+  }
+  get trainableWeights() {
+    return [this.gamma];
+  }
+  set trainable(a) {
+    this.gamma.trainable = a;
+  }
+  getWeights() {
+    return [this.gamma];
+  }
+  setWeights(a) {
+    this.gamma.assign(a[0]);
+  }
+  apply(a) {
+    return this.tf.tidy(() => {
+      const s = a.mean(-1, !0), t = a.sub(s), n = t.square().mean(-1, !0).add(this.epsilon).rsqrt();
+      return t.mul(n).mul(this.gamma);
+    });
+  }
+}
+export {
+  u as default
+};

package/dist/layers/MLP.d.ts

@@ -0,0 +1,17 @@
+import { default as TF } from '@tensorflow/tfjs';
+import { GPTConfig } from '../config';
+export default class MLP {
+    private cFc;
+    private cProj;
+    private dropout;
+    private tf;
+    private index;
+    private _trainable;
+    constructor(tf: typeof TF, index: number, config: GPTConfig);
+    get variables(): TF.Variable[];
+    get trainable(): boolean;
+    set trainable(value: boolean);
+    saveWeights(map: Map<string, TF.Tensor[]>): void;
+    loadWeights(weights: Map<string, TF.Tensor[]>): void;
+    call(x: TF.Tensor, training?: boolean): TF.Tensor;
+}

package/dist/layers/MLP.js

@@ -0,0 +1,57 @@
+class n {
+  cFc;
+  cProj;
+  dropout;
+  tf;
+  index;
+  _trainable = !0;
+  constructor(t, i, e) {
+    this.tf = t, this.index = i, this.cFc = this.tf.layers.dense({
+      units: 4 * e.nEmbed,
+      activation: "gelu",
+      useBias: e.biasInLinear,
+      kernelInitializer: this.tf.initializers.randomNormal({
+        mean: 0,
+        stddev: 0.02
+      }),
+      biasInitializer: "zeros",
+      name: `block_${i}_mlp_cFc`
+    }), this.cProj = this.tf.layers.dense({
+      units: e.nEmbed,
+      useBias: e.biasInLinear,
+      kernelInitializer: this.tf.initializers.randomNormal({
+        mean: 0,
+        stddev: 0.02 / Math.sqrt(2 * e.nLayer)
+      }),
+      biasInitializer: "zeros",
+      name: `block_${i}_mlp_cProj`
+    }), this.dropout = this.tf.layers.dropout({ rate: e.dropout });
+  }
+  get variables() {
+    return [
+      ...this.cFc.trainableWeights.map((t) => t.read()),
+      ...this.cProj.trainableWeights.map((t) => t.read())
+    ];
+  }
+  get trainable() {
+    return this._trainable;
+  }
+  set trainable(t) {
+    this._trainable = t, this.cFc.trainable = t, this.cProj.trainable = t;
+  }
+  saveWeights(t) {
+    t.set(`block_${this.index}_mlpHidden`, this.cFc.getWeights()), t.set(`block_${this.index}_mlpOut`, this.cProj.getWeights());
+  }
+  loadWeights(t) {
+    this.cFc.setWeights(t.get(`block_${this.index}_mlpHidden`) || []), this.cProj.setWeights(t.get(`block_${this.index}_mlpOut`) || []);
+  }
+  call(t, i = !1) {
+    return this.tf.tidy(() => {
+      const e = this.cFc.apply(t), s = this.cProj.apply(e);
+      return this.dropout.apply(s, { training: i });
+    });
+  }
+}
+export {
+  n as default
+};

package/dist/layers/TiedEmbedding.d.ts

@@ -0,0 +1,22 @@
+import { default as TF } from '@tensorflow/tfjs';
+export default class TiedEmbeddingOutputLayer {
+    private vocabSize;
+    private embedDim;
+    private tf;
+    private tiedWeights;
+    private initializer;
+    constructor(tf: typeof TF, config: {
+        vocabSize: number;
+        embedDim: number;
+        name?: string;
+    }, name?: string);
+    get variables(): TF.Variable[];
+    embed(inputs: TF.Tensor): TF.Tensor;
+    project(inputs: TF.Tensor): TF.Tensor;
+    getWeights(): TF.Tensor[];
+    setWeights(weights: TF.Tensor[]): void;
+    getConfig(): {
+        vocabSize: number;
+        embedDim: number;
+    };
+}