@genai-fi/nanogpt 0.8.1 → 0.8.2

package/dist/Generator.d.ts

@@ -45,7 +45,8 @@ export default class Generator extends EE<'start' | 'stop' | 'tokens'> {
     generate(prompt?: string, options?: IGenerateOptions): Promise<string>;
     stop(): void;
     getText(): string;
-    getAttentionData(): number[][][][];
+    getAttentionData(): number[][][][][];
     getProbabilitiesData(): number[][][];
+    getEmbeddingsData(): number[][][][];
     getTokens(): number[];
 }

package/dist/Generator.js

@@ -73,7 +73,7 @@ function N(m, t, e, i = !1) {
   const a = { logits: n === 1 ? x(o, [1, -1]) : o }, p = { numSamples: t, seed: e, normalized: i }, l = C.runKernel(I, a, p);
   return n === 1 ? x(l, [l.size]) : l;
 }
-const S = /* @__PURE__ */ A({ multinomial_: N }), H = [
+const D = /* @__PURE__ */ A({ multinomial_: N }), H = [
   ...Array.from({ length: 95 }, (m, t) => String.fromCharCode(t + 32)),
   // ASCII
   // Spanish accented letters and punctuation
@@ -112,7 +112,9 @@ class qt extends z {
       return null;
     const n = await t.decode([s]);
     if (i) {
-      const d = await Promise.all(i.map((a) => a.array().then((p) => p)));
+      const d = await Promise.all(
+        i.map((a) => a.array().then((p) => p))
+      );
       i.forEach((a) => a.dispose()), this.attentionData.push(d);
     }
     if (o) {
@@ -134,10 +136,10 @@ class qt extends z {
       const r = t, h = r.shape[1], u = h <= this.model.config.blockSize ? r : r.slice(
         [0, h - this.model.config.blockSize],
         [r.shape[0], this.model.config.blockSize]
-      ), k = d ? this.model.config.blockSize - u.shape[1] : 0, b = k > 0 ? _(u, [
+      ), g = d ? this.model.config.blockSize - u.shape[1] : 0, b = g > 0 ? _(u, [
         [0, 0],
-        [0, k]
-      ]) : u, [f] = this.model.forward(a, b), y = f.shape[1] - 1 - k, c = f.slice([0, y, 0], [f.shape[0], 1, f.shape[2]]);
+        [0, g]
+      ]) : u, [f] = this.model.forward(a, b), y = f.shape[1] - 1 - g, c = f.slice([0, y, 0], [f.shape[0], 1, f.shape[2]]);
       return a.attentionScores?.attentionOut && a.attentionScores.attentionOut.forEach((T, E) => {
         T.shape[1] !== 1 && (a.attentionScores.attentionOut[E] = R(
           T.slice([0, y, 0], [T.shape[0], 1, T.shape[2]])
@@ -148,19 +150,19 @@ class qt extends z {
     if (n) {
       const r = v(p), h = await r.array();
       r.dispose();
-      const u = h[0].map((c, g) => ({ prob: c, index: g })).sort((c, g) => g.prob - c.prob);
-      let k = 0;
+      const u = h[0].map((c, k) => ({ prob: c, index: k })).sort((c, k) => k.prob - c.prob);
+      let g = 0;
       const b = new Array(u.length).fill(0);
       for (const c of u)
-        if (k += c.prob, b[c.index] = c.prob, k >= n)
+        if (g += c.prob, b[c.index] = c.prob, g >= n)
           break;
-      const f = b.reduce((c, g) => c + g, 0), y = b.map((c) => c / f);
+      const f = b.reduce((c, k) => c + k, 0), y = b.map((c) => c / f);
       l = $(y);
     } else if (s) {
-      const { values: r, indices: h } = K(p, s), u = S(r, 1);
+      const { values: r, indices: h } = K(p, s), u = D(r, 1);
       l = q(h, u, 1), r.dispose(), h.dispose(), u.dispose();
     } else
-      l = S(p, 1);
+      l = D(p, 1);
     let w;
     i?.includeProbabilities && (w = v(p)), a.embeddings && this.embeddingsData.push(
       await Promise.all(
@@ -170,8 +172,8 @@ class qt extends z {
         })
       )
     );
-    const D = l.reshape([1, 1]);
-    return l.dispose(), l = D, p.dispose(), { output: l, probabilities: w, attention: a.attentionScores?.attentionOut };
+    const S = l.reshape([1, 1]);
+    return l.dispose(), l = S, p.dispose(), { output: l, probabilities: w, attention: a.attentionScores?.attentionOut };
   }
   /** Generate multiple tokens in a loop and produce text */
   async _generate(t) {
@@ -239,6 +241,9 @@ class qt extends z {
   getProbabilitiesData() {
     return this.probabilitiesData;
   }
+  getEmbeddingsData() {
+    return this.embeddingsData;
+  }
   getTokens() {
     return this.tokens;
   }

package/dist/checks/index.d.ts

@@ -6,6 +6,7 @@ import { execute as normRMSGrad } from './normRMSGrad';
 import { execute as appendCache } from './appendCache';
 import { execute as attentionMask } from './attentionMask';
 import { default as runCheck } from './check';
+import { createWeightStatistics, createTensorStatistics } from './weights';
 declare const checks: {
     rope: typeof rope;
     qkv: typeof qkv;
@@ -15,5 +16,7 @@ declare const checks: {
     appendCache: typeof appendCache;
     attentionMask: typeof attentionMask;
     runCheck: typeof runCheck;
+    createLayerWeightStatistics: typeof createWeightStatistics;
+    createWeightStatistics: typeof createTensorStatistics;
 };
 export default checks;

package/dist/checks/index.js

@@ -1,20 +1,23 @@
 import { execute as e } from "./rope.js";
 import { execute as t } from "./normRMS.js";
-import { execute as o } from "./qkv.js";
-import { execute as r } from "./gelu.js";
-import { execute as c } from "./normRMSGrad.js";
-import { execute as m } from "./appendCache.js";
-import { execute as u } from "./attentionMask.js";
-import x from "./check.js";
+import { execute as r } from "./qkv.js";
+import { execute as c } from "./gelu.js";
+import { execute as o } from "./normRMSGrad.js";
+import { execute as a } from "./appendCache.js";
+import { execute as i } from "./attentionMask.js";
+import m from "./check.js";
+import { createTensorStatistics as s, createWeightStatistics as u } from "./weights.js";
 const d = {
   rope: e,
-  qkv: o,
-  gelu: r,
+  qkv: r,
+  gelu: c,
   normRMS: t,
-  normRMSGrad: c,
-  appendCache: m,
-  attentionMask: u,
-  runCheck: x
+  normRMSGrad: o,
+  appendCache: a,
+  attentionMask: i,
+  runCheck: m,
+  createLayerWeightStatistics: u,
+  createWeightStatistics: s
 };
 export {
   d as default

package/dist/checks/qkv.js

@@ -1,9 +1,34 @@
-import { s as c, e as m } from "../index-DdmHGZjq.js";
-import { t as i } from "../tensor3d-BOukqWwr.js";
-import { t as u } from "../tensor2d-CObBWBkW.js";
-async function w(a) {
-  await c(a);
-  const o = i(
+import { z as i, A as u, B as c, s as l, e as h } from "../index-DdmHGZjq.js";
+import { t as f } from "../tensor2d-CObBWBkW.js";
+/**
+ * @license
+ * Copyright 2018 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+function m(t, e, n) {
+  if (i(t), e != null && e.length !== 3)
+    throw new Error("tensor3d() requires shape to have three numbers");
+  const r = u(t, n);
+  if (r.length !== 3 && r.length !== 1)
+    throw new Error("tensor3d() requires values to be number[][][] or flat/TypedArray");
+  if (r.length === 1 && e == null)
+    throw new Error("tensor3d() requires shape to be provided when `values` are a flat array");
+  return c(t, e, r, n);
+}
+async function y(t) {
+  await l(t);
+  const e = m(
     [
       [
         [0.1, 0.2],
@@ -11,15 +36,15 @@ async function w(a) {
       ]
     ],
     [1, 2, 2]
-  ), r = u(
+  ), n = f(
     [
       [0.5, 0.6, 0.9, 1, 1.3, 1.4],
       [0.7, 0.8, 1.1, 1.2, 1.5, 1.6]
     ],
     [2, 6]
-  ), t = m().runKernel("QKV", { x: o, kernel: r }, { heads: 1 }), s = await t[0].array(), n = await t[1].array(), e = await t[2].array();
-  return [s, n, e];
+  ), r = h().runKernel("QKV", { x: e, kernel: n }, { heads: 1 }), o = await r[0].array(), a = await r[1].array(), s = await r[2].array();
+  return [o, a, s];
 }
 export {
-  w as execute
+  y as execute
 };

package/dist/checks/rope.js

@@ -1,9 +1,9 @@
+import t from "../layers/RoPECache.js";
 import { s as c, e as i } from "../index-DdmHGZjq.js";
-import { t as m } from "../tensor4d-DLtk7Nxh.js";
-import { t } from "../tensor3d-BOukqWwr.js";
-async function y(n) {
-  await c(n);
-  const s = m(
+import { t as p } from "../tensor4d-DLtk7Nxh.js";
+async function y(a) {
+  await c(a);
+  const o = p(
     [
       [
         [
@@ -13,8 +13,25 @@ async function y(n) {
       ]
     ],
     [1, 1, 2, 2]
-  ), e = t([0.5, 0.6], [2, 1, 1]), o = t([0.9, 1], [2, 1, 1]), r = i().runKernel("Rope", { x: s, sin: e, cos: o }, { pastLen: 0 });
-  return Array.isArray(r) ? r.map((a) => a.array()) : r.array();
+  ), n = {
+    biasInLayerNorm: !1,
+    vocabSize: 20,
+    nEmbed: 16,
+    nHead: 2,
+    nLayer: 1,
+    biasInLinear: !1,
+    dropout: 0,
+    blockSize: 128,
+    mlpFactor: 4,
+    useRope: !0
+  }, e = new t(n);
+  e.ensureRopeCache(120);
+  const r = i().runKernel(
+    "Rope",
+    { x: o, sin: e.getSin(), cos: e.getCos() },
+    { pastLen: 20 }
+  );
+  return Array.isArray(r) ? r.map((s) => s.array()) : r.array();
 }
 export {
   y as execute

package/dist/checks/weights.d.ts

@@ -0,0 +1,16 @@
+import { default as BaseLayer } from '../layers/BaseLayer';
+import { Tensor } from '@tensorflow/tfjs-core';
+export interface TensorStatistics {
+    mean: number;
+    std: number;
+    min: number;
+    max: number;
+    sparsity: number;
+    isFinite: boolean;
+    hasNaN: boolean;
+    closeToZeroCount: number;
+}
+export declare function createTensorStatistics(weight: Tensor | number[]): Promise<TensorStatistics>;
+export declare function createWeightStatistics(layer: BaseLayer): Promise<{
+    [key: string]: TensorStatistics;
+}>;

package/dist/checks/weights.js

@@ -0,0 +1,29 @@
+async function d(s) {
+  const e = Array.isArray(s) ? s : await s.data(), a = e.length;
+  let n = 0, l = 0, i = e[0], r = e[0], u = 0, h = !0, f = !1, m = 0;
+  for (let c = 0; c < a; c++) {
+    const t = e[c];
+    n += t, l += t * t, t < i && (i = t), t > r && (r = t), t === 0 && u++, Math.abs(t) < 1e-8 && m++, Number.isNaN(t) && (f = !0);
+  }
+  const o = n / a, y = l / a - o * o, N = Math.sqrt(y), b = u / a;
+  return {
+    mean: o,
+    std: N,
+    min: i,
+    max: r,
+    sparsity: b,
+    isFinite: h,
+    hasNaN: f,
+    closeToZeroCount: m
+  };
+}
+async function S(s) {
+  const e = s.trainableVariables, a = {};
+  for (const n of e)
+    a[n.name] = await d(n);
+  return a;
+}
+export {
+  d as createTensorStatistics,
+  S as createWeightStatistics
+};

package/dist/layers/CausalSelfAttention.js

@@ -1,16 +1,16 @@
 import { attentionMask as g } from "../ops/attentionMask.js";
 import O from "./BaseLayer.js";
-import { qkv as P } from "../ops/qkv.js";
+import { qkv as _ } from "../ops/qkv.js";
 import { rope as v } from "../ops/rope.js";
 import { appendCache as V } from "../ops/appendCache.js";
-import { k as c, t as C } from "../index-DdmHGZjq.js";
-import { fusedSoftmax as T } from "../ops/fusedSoftmax.js";
-import { d as L } from "../random_width-DKGeiFuR.js";
-import { v as b } from "../variable-DPFOJyRG.js";
-import { r as k, d as y } from "../dropout-CcKSfOYE.js";
-import { r as N } from "../reshape-WeJkT3ja.js";
-import { m as R } from "../mat_mul-Dpy2mMRu.js";
-class $ extends O {
+import { k as c, t as P } from "../index-DdmHGZjq.js";
+import { fusedSoftmax as b } from "../ops/fusedSoftmax.js";
+import { d as C } from "../random_width-DKGeiFuR.js";
+import { v as k } from "../variable-DPFOJyRG.js";
+import { r as T, d as L } from "../dropout-CcKSfOYE.js";
+import { r as j } from "../reshape-WeJkT3ja.js";
+import { m as x } from "../mat_mul-Dpy2mMRu.js";
+class W extends O {
   divisor;
   index;
   units;
@@ -23,34 +23,34 @@ class $ extends O {
   build() {
     this.hasVariable(this.ATTN) === !1 && this.setVariable(
       this.ATTN,
-      b(
-        k([this.config.nEmbed, this.units], 0, 0.02),
-        !0
-        //`block_${this.index}_attn_cAttn_kernel`
+      k(
+        T([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.nEmbed], 0, 0.02),
-        !0
-        //`block_${this.index}_attn_cProj_kernel`
+      k(
+        T([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.dropout : 0, o);
+    const e = g(t, i, this.divisor), n = b(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.
   getAttentionScoresWithPast(t, i, s) {
-    const o = g(t, i, this.divisor, s), e = T(o, 0, 0);
+    const o = g(t, i, this.divisor, s), e = b(o, 0, 0);
     return o.dispose(), e;
   }
   getQKV(t) {
-    return P(t, this.getVariable(this.ATTN), this.config.nHead);
+    return _(t, this.getVariable(this.ATTN), this.config.nHead);
   }
   getOutputProjection(t) {
-    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));
+    const i = t.shape[0], s = t.shape[2], o = this.config.nEmbed, e = t.transpose([0, 2, 1, 3]), n = j(e, [i, s, o]), r = C(n, this.getVariable(this.PROJ));
     return n.dispose(), e.dispose(), r;
   }
   updateCache(t, i, s) {
@@ -62,19 +62,19 @@ class $ extends O {
     s.length = d, s.cumulativeLength = h, s.k = c(r), s.v = c(p);
   }
   forward(t, i) {
-    return C(() => {
+    return P(() => {
       this.startMemory();
       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;
+      const u = t.pastKV?.k ? t.pastKV.k : d, l = t.pastKV?.v ? t.pastKV.v : e;
       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 m = x(a, l), f = t.attentionScores !== void 0 && t.attentionScores.attentionOut !== void 0;
+      f || a.dispose(), t.pastKV || l.dispose();
+      const A = this.getOutputProjection(m);
+      if (m.dispose(), f && t.attentionScores && t.attentionScores.attentionOut !== void 0) {
         const K = a.shape[1], S = a.shape[2];
         t.attentionScores.attentionOut?.push(
           c(a.slice([0, 0, 0, 0], [1, -1, -1, -1]).reshape([K, S, -1]))
@@ -85,12 +85,12 @@ class $ extends O {
   }
   dropout(t) {
     if (this.config.dropout > 0) {
-      const i = y(t, this.config.dropout);
+      const i = L(t, this.config.dropout);
       return t.dispose(), i;
     } else
       return t;
   }
 }
 export {
-  $ as default
+  W as default
 };

package/dist/layers/MLP.js

@@ -1,10 +1,10 @@
 import { t as p } from "../index-DdmHGZjq.js";
 import u from "./BaseLayer.js";
 import { matMulGelu as M } from "../ops/matMulGelu.js";
-import { v as o } from "../variable-DPFOJyRG.js";
-import { r as h, d as f } from "../dropout-CcKSfOYE.js";
-import { r as d } from "../reshape-WeJkT3ja.js";
-import { m as c } from "../mat_mul-Dpy2mMRu.js";
+import { v as a } from "../variable-DPFOJyRG.js";
+import { r as d, d as c } from "../dropout-CcKSfOYE.js";
+import { r as h } from "../reshape-WeJkT3ja.js";
+import { m as b } from "../mat_mul-Dpy2mMRu.js";
 class H extends u {
   index;
   hiddenUnits;
@@ -16,32 +16,32 @@ class H extends u {
   build() {
     this.hasVariable(this.MLPHIDDEN) === !1 && this.setVariable(
       this.MLPHIDDEN,
-      o(
-        h([this.config.nEmbed, this.hiddenUnits], 0, 0.02),
-        !0
-        //`block_${this.index}_attn_cAttn_kernel`
+      a(
+        d([this.config.nEmbed, this.hiddenUnits], 0, 0.02),
+        !0,
+        `block_${this.index}_mlpHidden_kernel`
       )
     ), this.hasVariable(this.MLPOUT) === !1 && this.setVariable(
       this.MLPOUT,
-      o(
-        h([this.hiddenUnits, this.config.nEmbed], 0, 0.02 / Math.sqrt(2 * this.config.nLayer)),
-        !0
-        //`block_${this.index}_attn_cProj_kernel`
+      a(
+        d([this.hiddenUnits, this.config.nEmbed], 0, 0.02 / Math.sqrt(2 * this.config.nLayer)),
+        !0,
+        `block_${this.index}_mlpOut_kernel`
       )
     );
   }
   forward(i, t) {
     return p(() => {
       this.startMemory();
-      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 l = d(m, [s, r, e]);
-      return this.endMemory("MLP"), l;
+      const [s, e, r] = t.shape, n = h(t, [s * e, r]), o = M(n, this.getVariable(this.MLPHIDDEN)), l = b(o, this.getVariable(this.MLPOUT));
+      o.dispose();
+      const m = h(l, [s, e, r]);
+      return this.endMemory("MLP"), m;
     });
   }
   dropout(i) {
     if (this.config.dropout > 0) {
-      const t = f(i, this.config.dropout);
+      const t = c(i, this.config.dropout);
       return i.dispose(), t;
     }
     return i;

package/dist/layers/TiedEmbedding.js

@@ -2,28 +2,28 @@ import { d as r } from "../random_width-DKGeiFuR.js";
 import "../index-DdmHGZjq.js";
 import { r as a } from "../exports_initializers-DKk7-bsx.js";
 import s from "./BaseLayer.js";
-import { v as m } from "../variable-DPFOJyRG.js";
-import { g as o } from "../gather-CPg6ZlQA.js";
+import { v as o } from "../variable-DPFOJyRG.js";
+import { g as m } from "../gather-CPg6ZlQA.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({
+  constructor(i, t, e) {
+    super(i, e), this.WEIGHTS = t, 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));
+    }), this.addVariable(this.WEIGHTS, o(this.initializer.apply([this.vocabSize, this.embedDim]), !0, t));
   }
   embed(i) {
-    return o(this.getVariable(this.WEIGHTS), i, 0);
+    return m(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);
+  forward(i, t) {
+    return this.project(t);
   }
 }
 export {

package/dist/main.d.ts

@@ -26,3 +26,4 @@ export declare const layers: {
 };
 export { default as AdamExt } from './training/AdamExt';
 export { default as checks } from './checks';
+export type { TensorStatistics } from './checks/weights';

package/package.json

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

package/dist/tensor3d-BOukqWwr.js

@@ -1,30 +0,0 @@
-import { z as o, A as a, B as s } from "./index-DdmHGZjq.js";
-/**
- * @license
- * Copyright 2018 Google LLC. All Rights Reserved.
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- * =============================================================================
- */
-function h(n, r, t) {
-  if (o(n), r != null && r.length !== 3)
-    throw new Error("tensor3d() requires shape to have three numbers");
-  const e = a(n, t);
-  if (e.length !== 3 && e.length !== 1)
-    throw new Error("tensor3d() requires values to be number[][][] or flat/TypedArray");
-  if (e.length === 1 && r == null)
-    throw new Error("tensor3d() requires shape to be provided when `values` are a flat array");
-  return s(n, r, e, t);
-}
-export {
-  h as t
-};