@genai-fi/nanogpt 0.4.5 → 0.5.1

package/dist/BaseLayer-BhrMN8JO.js

@@ -0,0 +1,135 @@
+import { E as p, V as v, c as _, e as o, W as V } from "./index-iNhkcAEQ.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 E(l, t = !0, r, e) {
+  return p.makeVariable(l, t, r, e);
+}
+class m {
+  parent;
+  config;
+  _variables = /* @__PURE__ */ new Map();
+  _trainable = !0;
+  children = [];
+  constructor(t, r) {
+    this.config = t, this.parent = r, this.parent && this.parent.children.push(this);
+  }
+  getProfiler() {
+    return this.config.layerConfig.profiler;
+  }
+  startMemory() {
+    this.config.layerConfig.profiler?.startMemory();
+  }
+  endMemory(t) {
+    this.config.layerConfig.profiler?.endMemory(t);
+  }
+  addVariable(t, r) {
+    this._variables.set(t, r || null);
+  }
+  get variables() {
+    const t = Array.from(this._variables.values()).filter((e) => e !== null), r = this.children.flatMap((e) => e.variables);
+    return [...t, ...r];
+  }
+  get trainableVariables() {
+    const t = Array.from(this._variables.values()).filter(
+      (e) => e !== null && e.trainable
+    ), r = this.children.flatMap((e) => e.trainableVariables);
+    return [...t, ...r];
+  }
+  get trainable() {
+    return this._trainable;
+  }
+  set trainable(t) {
+    this._trainable = t, this._variables.forEach((r) => {
+      r && (r.trainable = t);
+    }), this.children.forEach((r) => {
+      r.trainable = t;
+    });
+  }
+  getVariable(t) {
+    const r = this._variables.get(t);
+    if (!r)
+      throw new Error(`Variable ${t} not found`);
+    return r;
+  }
+  hasVariable(t) {
+    return this._variables.get(t) !== null;
+  }
+  setVariable(t, r) {
+    if (!this._variables.has(t))
+      throw new Error(`Variable ${t} not found`);
+    this._variables.set(t, r);
+  }
+  saveWeights(t) {
+    this._variables.forEach((r, e) => {
+      r && t.set(e, [r.clone()]);
+    }), this.children.forEach((r) => {
+      r.saveWeights(t);
+    });
+  }
+  loadWeights(t) {
+    this._variables.forEach((r, e) => {
+      const i = t.get(e)?.[0];
+      if (!i)
+        throw new Error(`Weights for ${e} not found`);
+      r ? r.assign(i) : this._variables.set(e, E(i, this._trainable));
+    }), this.children.forEach((r) => {
+      r.loadWeights(t);
+    });
+  }
+  dispose() {
+    this._variables.forEach((t) => {
+      t?.dispose();
+    }), this._variables.clear();
+  }
+  build() {
+  }
+  dropout(t) {
+    return t;
+  }
+  call(t, ...r) {
+    this.build();
+    const e = this.forward(t, ...r);
+    if (t.training && e instanceof v) {
+      const i = this.dropout(e);
+      return i !== e && e.dispose(), i;
+    } else
+      return e;
+  }
+  callCheckpoint(t, ...r) {
+    return this.build(), this.checkpointingFn(t, ...r);
+  }
+  checkpointingFn(t, ...r) {
+    const e = this.trainableVariables, s = _((...a) => {
+      const h = a[a.length - 1], n = a.slice(0, r.length), c = this.forward(t, ...n);
+      return h(n), { value: c, gradFunc: (f, u) => {
+        const b = o().state.activeTape;
+        o().state.activeTape = [];
+        const d = V((...g) => this.forward(t, ...g.slice(0, n.length)))([...u, ...e], f);
+        return o().state.activeTape = b, d;
+      } };
+    })(...r, ...e);
+    if (t.training) {
+      const a = this.dropout(s);
+      return a !== s && s.dispose(), a;
+    } else
+      return s;
+  }
+}
+export {
+  m as B,
+  E as v
+};

package/dist/Generator.js

@@ -1,70 +1,73 @@
 import { E as u } from "./index-Dwqa6Zy2.js";
-import "./index--6vO-cOz.js";
-import { t as d } from "./tensor2d-DUr_htjt.js";
-import { c as p } from "./concat-DvWM7HGZ.js";
+import "./index-iNhkcAEQ.js";
+import { t as d } from "./tensor2d-tSxWdFMH.js";
+import { c as p } from "./concat-Cxbo2sOz.js";
 class w extends u {
-  constructor(s, e) {
-    super(), this.model = s, this.tokeniser = e;
+  constructor(i, t) {
+    super(), this.model = i, this.tokeniser = t;
   }
   active = !1;
-  async tokenisePrompt(s) {
-    const e = s ? await this.tokeniser.tokenise([s], !0) : [[this.tokeniser.eosToken]];
-    return d(e, [1, e[0].length], "int32");
+  async tokenisePrompt(i) {
+    const t = i ? await this.tokeniser.tokenise([i], !0) : [[this.tokeniser.eosToken]];
+    return d(t, [1, t[0].length], "int32");
   }
-  async generateNoCache(s, e) {
-    let t = await this.tokenisePrompt(s), n = s || "";
-    const o = e?.maxLength ?? 1e3;
-    for (let a = 0; a < o && this.active; a++) {
+  async generateNoCache(i, t) {
+    let s = await this.tokenisePrompt(i), o = i || "";
+    const n = t?.maxLength ?? 1e3;
+    for (let r = 0; r < n && this.active; r++) {
       const {
-        output: i,
-        attention: c,
-        probabilities: l
-      } = this.model.generate(t, void 0, e), h = t;
-      t = p([t, i], 1), h.dispose();
-      const r = await this.processResponse(i, c, l);
-      if (i.dispose(), r === null)
+        output: e,
+        attention: a,
+        probabilities: c
+      } = this.model.generate(s, void 0, t), h = s;
+      s = p([s, e], 1), h.dispose();
+      const l = await this.processResponse(e, a, c);
+      if (e.dispose(), l === null)
         break;
-      n += r;
+      o += l;
     }
-    return t.dispose(), n;
+    return s.dispose(), o;
   }
-  async processResponse(s, e, t) {
-    const n = (await s.array())[0][0];
-    if (n === this.tokeniser.eosToken)
+  async processResponse(i, t, s) {
+    const o = (await i.array())[0][0];
+    if (o === this.tokeniser.eosToken)
       return null;
-    const o = await this.tokeniser.decode([n]);
-    let a;
-    e && (a = await e.array(), e.dispose());
-    let i;
-    return t && (i = await t.array(), t.dispose()), this.emit("tokens", [n], o, a, i), o;
+    const n = await this.tokeniser.decode([o]);
+    let r;
+    t && (r = await Promise.all(t.map((a) => a.array().then((c) => c))), t.forEach((a) => a.dispose()));
+    let e;
+    return s && (e = await s.array(), s.dispose()), this.emit("tokens", [o], n, r, e), n;
   }
-  async generateCache(s, e) {
-    let t = await this.tokenisePrompt(s), n = s || "";
-    const o = new Array(this.model.config.gpt.nLayer).fill(void 0), a = e?.maxLength ?? 1e3;
-    for (let i = 0; i < a && this.active; i++) {
+  async generateCache(i, t) {
+    let s = await this.tokenisePrompt(i), o = i || "";
+    const n = new Array(this.model.config.gpt.nLayer);
+    for (let e = 0; e < this.model.config.gpt.nLayer; e++)
+      n[e] = { k: void 0, v: void 0, length: 0, cumulativeLength: 0 };
+    const r = t?.maxLength ?? 1e3;
+    for (let e = 0; e < r && this.active; e++) {
       const {
-        output: c,
-        attention: l,
-        probabilities: h
-      } = this.model.generate(t, o, {
-        ...e,
+        output: a,
+        probabilities: c,
+        attention: h
+      } = this.model.generate(s, n, {
+        ...t,
         usePadding: !1
       });
-      t.dispose(), t = c;
-      const r = await this.processResponse(c, l, h);
-      if (r === null)
+      s.dispose(), s = a;
+      const l = await this.processResponse(a, h, c);
+      if (l === null)
         break;
-      n += r;
+      o += l;
     }
-    return o.forEach((i) => {
-      i && (i.k.dispose(), i.v.dispose());
-    }), t.dispose(), n;
+    return n.forEach((e) => {
+      e && (e.k && e.k.dispose(), e.v && e.v.dispose());
+    }), s.dispose(), o;
   }
-  async generate(s, e) {
-    const t = s && s.length > this.model.config.gpt.blockSize ? s.slice(-this.model.config.gpt.blockSize) : s;
+  async generate(i, t) {
+    const s = i && i.length > this.model.config.gpt.blockSize ? i.slice(-this.model.config.gpt.blockSize) : i;
     this.active = !0, this.emit("start");
-    const o = await (this.model.config.gpt.useRope && !e?.noCache && !e?.includeAttention ? this.generateCache(t, e) : this.generateNoCache(t, e));
-    return this.active = !1, this.emit("stop"), o;
+    const n = await (this.model.config.gpt.useRope && !t?.noCache ? this.generateCache(s, t) : this.generateNoCache(s, t));
+    return this.active = !1, this.emit("stop"), n;
   }
   stop() {
     this.active = !1;

package/dist/NanoGPTModel.d.ts

@@ -1,7 +1,7 @@
 import { GPTConfig } from './config';
-import { KVCache } from './layers/CausalSelfAttention';
-import { default as BaseLayer } from './layers/BaseLayer';
-import { Tensor, Variable } from '@tensorflow/tfjs-core';
+import { AttentionScores, KVCache } from './layers/CausalSelfAttention';
+import { default as BaseLayer, ForwardAttributes } from './layers/BaseLayer';
+import { Tensor } from '@tensorflow/tfjs-core';
 export interface TrainingLogEntry {
     loss: number;
     valLoss?: number;
@@ -14,10 +14,15 @@ export interface GenerateOptions {
     temperature?: number;
     topK?: number;
     usePadding?: boolean;
-    includeAttention?: boolean;
+    attentionScores?: boolean;
     includeProbabilities?: boolean;
 }
-export default class NanoGPT extends BaseLayer {
+export interface ModelForwardAttributes extends ForwardAttributes {
+    cache?: KVCache[];
+    attentionScores?: AttentionScores;
+    seed?: number;
+}
+export default class NanoGPT extends BaseLayer<ModelForwardAttributes> {
     private wte;
     private wpe?;
     private drop;
@@ -28,25 +33,16 @@ export default class NanoGPT extends BaseLayer {
     constructor(config?: Partial<GPTConfig>);
     get checkpointing(): boolean;
     set checkpointing(value: boolean);
-    get variables(): Variable[];
-    saveWeights(): Map<string, Tensor[]>;
-    loadWeights(weights: Map<string, Tensor[]>): void;
     private inputPhase;
     setSkipMask(mask: boolean[]): void;
     setTrainableMask(mask: boolean[]): void;
-    set trainable(value: boolean);
     private validateInput;
     private calculateLoss;
-    private computeAttentionRollout;
-    forward(idx: Tensor, targets?: Tensor, training?: boolean, includeAttention?: boolean, cache?: (KVCache | undefined)[]): {
-        logits: Tensor;
-        loss?: Tensor;
-        attention?: Tensor;
-    };
-    generate(idx: Tensor, cache?: (KVCache | undefined)[], options?: GenerateOptions): {
+    forward(attrs: ModelForwardAttributes, idx: Tensor, targets?: Tensor): Tensor[];
+    generate(idx: Tensor, cache?: KVCache[], options?: GenerateOptions): {
         output: Tensor;
-        attention?: Tensor;
         probabilities?: Tensor;
+        attention?: Tensor[];
     };
     getNumParams(): number;
     dispose(): void;