@genai-fi/nanogpt 0.8.2 → 0.8.3

package/dist/Generator.d.ts

@@ -47,6 +47,9 @@ export default class Generator extends EE<'start' | 'stop' | 'tokens'> {
     getText(): string;
     getAttentionData(): number[][][][][];
     getProbabilitiesData(): number[][][];
-    getEmbeddingsData(): number[][][][];
+    getEmbeddingsData(): {
+        name: string;
+        tensor: number[][];
+    }[][];
     getTokens(): number[];
 }

package/dist/Generator.js

@@ -63,18 +63,18 @@ import { c as G } from "./concat-pHiVqR3L.js";
  * limitations under the License.
  * =============================================================================
  */
-function N(m, t, e, i = !1) {
-  const o = L(m, "logits", "multinomial"), s = o.size, n = o.rank;
+function N(h, t, e, i = !1) {
+  const o = L(h, "logits", "multinomial"), s = o.size, n = o.rank;
   if (s < 2)
     throw new Error(`Error in multinomial: you need at least 2 outcomes, but got ${s}.`);
   if (n > 2)
     throw new Error(`Rank of probabilities must be 1 or 2, but is ${n}`);
   e = e || Math.random();
-  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 a = { logits: n === 1 ? x(o, [1, -1]) : o }, l = { numSamples: t, seed: e, normalized: i }, m = C.runKernel(I, a, l);
+  return n === 1 ? x(m, [m.size]) : m;
 }
 const D = /* @__PURE__ */ A({ multinomial_: N }), H = [
-  ...Array.from({ length: 95 }, (m, t) => String.fromCharCode(t + 32)),
+  ...Array.from({ length: 95 }, (h, t) => String.fromCharCode(t + 32)),
   // ASCII
   // Spanish accented letters and punctuation
   ..."áéíóúüñ¿¡",
@@ -85,8 +85,8 @@ const D = /* @__PURE__ */ A({ multinomial_: N }), H = [
   // Cyrillic letters
   ..."абвгдеёжзийклмнопрстуфхцчшщъыьэюяАБВГДЕЁЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЫЬЭЮЯ"
 ];
-function U(m, t) {
-  return m.length === t ? m : m.length > t ? m.slice(0, t) : m.concat(Array(t - m.length).fill(""));
+function U(h, t) {
+  return h.length === t ? h : h.length > t ? h.slice(0, t) : h.concat(Array(t - h.length).fill(""));
 }
 class qt extends z {
   constructor(t, e) {
@@ -113,7 +113,7 @@ class qt extends z {
     const n = await t.decode([s]);
     if (i) {
       const d = await Promise.all(
-        i.map((a) => a.array().then((p) => p))
+        i.map((a) => a.array().then((l) => l))
       );
       i.forEach((a) => a.dispose()), this.attentionData.push(d);
     }
@@ -132,48 +132,47 @@ class qt extends z {
       } : void 0,
       cache: e,
       outputEmbeddings: i?.embeddings ?? !1
-    }, p = O(() => {
-      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]
-      ), g = d ? this.model.config.blockSize - u.shape[1] : 0, b = g > 0 ? _(u, [
+    }, l = O(() => {
+      const p = t, u = p.shape[1], r = u <= this.model.config.blockSize ? p : p.slice(
+        [0, u - this.model.config.blockSize],
+        [p.shape[0], this.model.config.blockSize]
+      ), f = d ? this.model.config.blockSize - r.shape[1] : 0, b = f > 0 ? _(r, [
         [0, 0],
-        [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]]);
+        [0, f]
+      ]) : r, [g] = this.model.forward(a, b), y = g.shape[1] - 1 - f, c = g.slice([0, y, 0], [g.shape[0], 1, g.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]])
         ), T.dispose());
-      }), f.dispose(), c.div(o).squeeze([1]);
+      }), g.dispose(), c.div(o).squeeze([1]);
     });
-    let l;
+    let m;
     if (n) {
-      const r = v(p), h = await r.array();
-      r.dispose();
-      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 (g += c.prob, b[c.index] = c.prob, g >= n)
+      const p = v(l), u = await p.array();
+      p.dispose();
+      const r = u[0].map((c, k) => ({ prob: c, index: k })).sort((c, k) => k.prob - c.prob);
+      let f = 0;
+      const b = new Array(r.length).fill(0);
+      for (const c of r)
+        if (f += c.prob, b[c.index] = c.prob, f >= n)
           break;
-      const f = b.reduce((c, k) => c + k, 0), y = b.map((c) => c / f);
-      l = $(y);
+      const g = b.reduce((c, k) => c + k, 0), y = b.map((c) => c / g);
+      m = $(y);
     } else if (s) {
-      const { values: r, indices: h } = K(p, s), u = D(r, 1);
-      l = q(h, u, 1), r.dispose(), h.dispose(), u.dispose();
+      const { values: p, indices: u } = K(l, s), r = D(p, 1);
+      m = q(u, r, 1), p.dispose(), u.dispose(), r.dispose();
     } else
-      l = D(p, 1);
+      m = D(l, 1);
     let w;
-    i?.includeProbabilities && (w = v(p)), a.embeddings && this.embeddingsData.push(
-      await Promise.all(
-        a.embeddings.map(async (r) => {
-          const h = await r.array();
-          return r.dispose(), h;
-        })
-      )
-    );
-    const S = l.reshape([1, 1]);
-    return l.dispose(), l = S, p.dispose(), { output: l, probabilities: w, attention: a.attentionScores?.attentionOut };
+    if (i?.includeProbabilities && (w = v(l)), a.embeddings) {
+      const p = a.embeddings.map(async (r) => {
+        const f = await r.tensor.array();
+        return r.tensor.dispose(), { name: r.name, tensor: f };
+      }), u = await Promise.all(p);
+      this.embeddingsData.push(u);
+    }
+    const S = m.reshape([1, 1]);
+    return m.dispose(), m = S, l.dispose(), { output: m, probabilities: w, attention: a.attentionScores?.attentionOut };
   }
   /** Generate multiple tokens in a loop and produce text */
   async _generate(t) {
@@ -191,8 +190,8 @@ class qt extends z {
       if (this.cache)
         e.dispose(), e = s;
       else {
-        const p = e;
-        e = G([e, s], 1), p.dispose();
+        const l = e;
+        e = G([e, s], 1), l.dispose();
       }
       const a = await this.processResponse(this.actualTokeniser, s, d, n);
       if (this.cache || s.dispose(), a === null)

package/dist/layers/BaseLayer.d.ts

@@ -6,6 +6,11 @@ export interface ForwardAttributes {
     training: boolean;
     checkpointing?: boolean;
     ropeCache?: RoPECache;
+    outputEmbeddings?: boolean;
+    embeddings?: {
+        name: string;
+        tensor: Tensor;
+    }[];
 }
 export default abstract class BaseLayer<ATTR extends ForwardAttributes = ForwardAttributes> {
     readonly parent?: BaseLayer;

package/dist/layers/TransformerBlock.js

@@ -1,32 +1,32 @@
 import l from "./CausalSelfAttention.js";
-import r from "./MLP.js";
+import p from "./MLP.js";
 import o from "./RMSNorm.js";
-import d from "./BaseLayer.js";
-import { t as p } from "../index-DdmHGZjq.js";
-class k extends d {
+import m from "./BaseLayer.js";
+import { k as n, t as h } from "../index-DdmHGZjq.js";
+class k extends m {
   ln1;
   attn;
   ln2;
   mlp;
   index;
   skipped = !1;
-  constructor(t, s, i) {
-    super(s, i), this.index = t, this.ln1 = new o(s, `block_${this.index}_rms1`, this), this.attn = new l(this.index, s, this), this.ln2 = new o(s, `block_${this.index}_rms2`, this), this.mlp = new r(this.index, s, this);
+  constructor(i, s, e) {
+    super(s, e), this.index = i, this.ln1 = new o(s, `block_${this.index}_rms1`, this), this.attn = new l(this.index, s, this), this.ln2 = new o(s, `block_${this.index}_rms2`, this), this.mlp = new p(this.index, s, this);
   }
-  getMLPOutput(t, s) {
-    const i = this.ln2.call({ training: s }, t), e = this.mlp.call({ training: s }, i);
-    i.dispose();
-    const n = t.add(e);
-    return t.dispose(), e.dispose(), n;
+  getMLPOutput(i, s) {
+    const e = this.ln2.call({ training: s.training }, i), t = this.mlp.call({ training: s.training }, e);
+    s.outputEmbeddings ? (n(e), s.embeddings.push({ name: `block_ln2_${this.index}`, tensor: e })) : e.dispose();
+    const d = i.add(t);
+    return i.dispose(), s.outputEmbeddings ? (n(t), s.embeddings.push({ name: `block_mlp_out_${this.index}`, tensor: t })) : t.dispose(), d;
   }
-  forward(t, s) {
-    return p(() => {
+  forward(i, s) {
+    return h(() => {
       if (this.skipped)
         return s;
-      const i = this.ln1.call(t, s), e = this.attn.call(t, i);
-      i.dispose();
-      const n = s.add(e);
-      return e.dispose(), this.getMLPOutput(n, t.training);
+      const e = this.ln1.call(i, s), t = this.attn.call(i, e);
+      i.outputEmbeddings ? (n(e), i.embeddings.push({ name: `block_ln1_${this.index}`, tensor: e })) : e.dispose();
+      const d = s.add(t);
+      return i.outputEmbeddings ? (n(t), i.embeddings.push({ name: `block_attn_out_${this.index}`, tensor: t })) : t.dispose(), this.getMLPOutput(d, i);
     });
   }
   dispose() {

package/dist/models/NanoGPTV1.js

@@ -1,6 +1,6 @@
-import { defaultConfig as m } from "./config.js";
-import f from "../layers/TransformerBlock.js";
-import u from "../layers/TiedEmbedding.js";
+import { defaultConfig as a } from "./config.js";
+import u from "../layers/TransformerBlock.js";
+import f from "../layers/TiedEmbedding.js";
 import g from "../layers/RoPECache.js";
 import b from "../layers/RMSNorm.js";
 import { t as l, k as p } from "../index-DdmHGZjq.js";
@@ -17,9 +17,9 @@ class R extends w {
   // Final layer norm
   ropeCache;
   constructor(e = {}) {
-    super({ ...m, ...e }), this.wte = new u(this.config, "token_embedding", this), this.config.useRope === !1 ? this.wpe = new k(this.config, "positional_embedding", this) : this.ropeCache = new g(this.config), this.blocks = [];
+    super({ ...a, ...e }), this.wte = new f(this.config, "token_embedding", this), this.config.useRope === !1 ? this.wpe = new k(this.config, "positional_embedding", this) : this.ropeCache = new g(this.config), this.blocks = [];
     for (let i = 0; i < this.config.nLayer; i++)
-      this.blocks.push(new f(i, this.config, this));
+      this.blocks.push(new u(i, this.config, this));
     this.lnF = new b(this.config, "final_rms_norm", this);
   }
   getClassName() {
@@ -47,17 +47,15 @@ class R extends w {
         );
       for (let t = 0; t < this.blocks.length; t++) {
         const c = this.blocks[t], d = Math.random() * 1e9, r = {
-          ropeCache: e.ropeCache,
-          training: e.training,
+          ...e,
           seed: d,
-          attentionScores: e.attentionScores,
           pastKV: e.cache ? e.cache[t] : void 0
-        }, a = e.checkpointing && e.training ? c.callCheckpoint(r, o) : c.call(r, o);
-        e.outputEmbeddings ? (p(o), e.embeddings.push(o)) : o.dispose(), o = a;
+        }, m = e.checkpointing && e.training ? c.callCheckpoint(r, o) : c.call(r, o);
+        e.outputEmbeddings ? (p(o), e.embeddings.push({ name: `block_output_${t}`, tensor: o })) : o.dispose(), o = m;
       }
       o = this.lnF.call(e, o);
       const n = this.wte.project(o);
-      e.outputEmbeddings ? (p(o), e.embeddings.push(o)) : o.dispose();
+      e.outputEmbeddings ? (p(o), e.embeddings.push({ name: "final_norm_output", tensor: o })) : o.dispose();
       let h;
       return s && (h = this.calculateLoss(n, s)), this.endMemory("Forward"), h ? [n, h] : [n];
     });

package/dist/models/model.d.ts

@@ -5,8 +5,6 @@ export interface ModelForwardAttributes extends ForwardAttributes {
     cache?: KVCache[];
     attentionScores?: AttentionScores;
     seed?: number;
-    outputEmbeddings?: boolean;
-    embeddings?: Tensor[];
 }
 interface TrainingState {
     steps: number;

package/package.json

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