@genai-fi/nanogpt 0.7.3 → 0.8.0

package/dist/NanoGPTModel.d.ts

@@ -1,52 +0,0 @@
-import { GPTConfig } from './config';
-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;
-    step: number;
-    time: number;
-    example?: string;
-    batchSize: number;
-    gradientNorm?: number;
-    learningRate?: number;
-}
-export interface GenerateOptions {
-    temperature?: number;
-    topK?: number;
-    topP?: number;
-    usePadding?: boolean;
-    attentionScores?: boolean;
-    includeProbabilities?: boolean;
-}
-export interface ModelForwardAttributes extends ForwardAttributes {
-    cache?: KVCache[];
-    attentionScores?: AttentionScores;
-    seed?: number;
-}
-export default class NanoGPT extends BaseLayer<ModelForwardAttributes> {
-    private wte;
-    private wpe?;
-    private drop;
-    private blocks;
-    private lnF;
-    private ropeCache?;
-    log: TrainingLogEntry[];
-    constructor(config?: Partial<GPTConfig>);
-    get checkpointing(): boolean;
-    set checkpointing(value: boolean);
-    private inputPhase;
-    setSkipMask(mask: boolean[]): void;
-    setTrainableMask(mask: boolean[]): void;
-    private validateInput;
-    private calculateLoss;
-    forward(attrs: ModelForwardAttributes, idx: Tensor, targets?: Tensor): Tensor[];
-    generate(idx: Tensor, cache?: KVCache[], options?: GenerateOptions): Promise<{
-        output: Tensor;
-        probabilities?: Tensor;
-        attention?: Tensor[];
-    }>;
-    getNumParams(): number;
-    dispose(): void;
-}

package/dist/NanoGPTModel.js

@@ -1,203 +0,0 @@
-import { defaultConfig as M } from "./config.js";
-import v from "./layers/TransformerBlock.js";
-import { T as x, r as T } from "./TiedEmbedding-BxOerUmB.js";
-import F from "./layers/RoPECache.js";
-import O from "./layers/RMSNorm.js";
-import { estimateParameterCount as _ } from "./utilities/parameters.js";
-import { createSoftmaxCrossEntropyWithGrad as D } from "./training/sparseCrossEntropy.js";
-import K from "./layers/BaseLayer.js";
-import { E as N, D as R, p as q } from "./random_width-sZORGo5k.js";
-import { B as A, C as B, E as G, ad as V, t as C, o as j, b as z, w as U } from "./index-BoWRt-10.js";
-import W from "./utilities/multinomialCPU.js";
-import { m as H, t as J } from "./register_all_kernels-BwDSRN-f.js";
-import { r as P } from "./reshape-CdBq1WJ6.js";
-import { r as Q } from "./range-CRuAh-gd.js";
-import { s as $ } from "./softmax-Bv_6lyMX.js";
-import { g as X } from "./gather-CMMy2KEG.js";
-/**
- * @license
- * Copyright 2020 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 Y(u, t, o, e = !1) {
-  const l = B(u, "logits", "multinomial"), s = l.size, r = l.rank;
-  if (s < 2)
-    throw new Error(`Error in multinomial: you need at least 2 outcomes, but got ${s}.`);
-  if (r > 2)
-    throw new Error(`Rank of probabilities must be 1 or 2, but is ${r}`);
-  o = o || Math.random();
-  const n = { logits: r === 1 ? P(l, [1, -1]) : l }, a = { numSamples: t, seed: o, normalized: e }, i = G.runKernel(V, n, a);
-  return r === 1 ? P(i, [i.size]) : i;
-}
-const I = /* @__PURE__ */ A({ multinomial_: Y });
-/**
- * @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 Z(u) {
-  return new R(u);
-}
-function tt(u) {
-  return new N(u);
-}
-class bt extends K {
-  wte;
-  // Token embeddings
-  wpe;
-  // Position embeddings
-  drop;
-  // Dropout
-  blocks;
-  lnF;
-  // Final layer norm
-  ropeCache;
-  log = [];
-  // Training log
-  constructor(t = {}) {
-    super({ gpt: { ...M, ...t }, layerConfig: {} }), this.wte = new x(this.config, "token_embedding", this), this.config.gpt.useRope === !1 ? this.wpe = tt({
-      inputDim: this.config.gpt.blockSize,
-      outputDim: this.config.gpt.nEmbed,
-      name: "positional_embedding",
-      embeddingsInitializer: T({ mean: 0, stddev: 0.02 })
-    }) : (this.ropeCache = new F(this.config.gpt), this.config.layerConfig.ropeCache = this.ropeCache), this.drop = Z({ rate: this.config.gpt.dropout }), this.blocks = [];
-    for (let o = 0; o < this.config.gpt.nLayer; o++)
-      this.blocks.push(new v(o, this.config, this));
-    this.lnF = new O(this.config, "final_rms_norm", this);
-  }
-  get checkpointing() {
-    return this.config.layerConfig.checkpointing === !0;
-  }
-  set checkpointing(t) {
-    this.config.layerConfig.checkpointing = t;
-  }
-  inputPhase(t, o, e = !1) {
-    return C(() => {
-      const l = this.wte.embed(t);
-      if (this.config.gpt.useRope === !1) {
-        const [, s] = t.shape, r = this.config.gpt.blockSize, g = Q(0, s, 1, "int32"), n = H(j(g, z(o, "int32")), z(r, "int32")), a = this.wpe.apply(n), i = l.add(a);
-        return this.drop.apply(i, { training: e });
-      } else
-        return this.drop.apply(l, { training: e });
-    });
-  }
-  setSkipMask(t) {
-    if (t.length !== this.blocks.length)
-      throw new Error(`Mask length ${t.length} does not match number of blocks ${this.blocks.length}`);
-    for (let o = 0; o < this.blocks.length; o++)
-      this.blocks[o].skipped = t[o];
-  }
-  setTrainableMask(t) {
-    if (t.length !== this.blocks.length)
-      throw new Error(`Mask length ${t.length} does not match number of blocks ${this.blocks.length}`);
-    for (let o = 0; o < this.blocks.length; o++)
-      this.blocks[o].trainable = t[o];
-  }
-  validateInput(t) {
-    if (t.shape.length !== 2)
-      throw new Error(`Invalid input shape: expected [batch_size, sequence_length], got ${t.shape}`);
-    if (t.shape[1] > this.config.gpt.blockSize)
-      throw new Error(`Input sequence length ${t.shape[1]} isn't block size ${this.config.gpt.blockSize}`);
-    if (t.dtype !== "int32")
-      throw new Error(`Input tensor must be of type int32, got ${t.dtype}`);
-  }
-  calculateLoss(t, o) {
-    try {
-      return D()(t, o).mean();
-    } catch (e) {
-      throw console.error("Error computing loss:", e), new Error(`Loss computation failed: ${e}`);
-    }
-  }
-  forward(t, o, e) {
-    return this.validateInput(o), C(() => {
-      this.startMemory();
-      const l = t.cache?.[0]?.length ?? 0;
-      let s = this.inputPhase(o, l, t.training);
-      if (t.cache && t.cache.length !== this.blocks.length)
-        throw console.error("Cache", t.cache), new Error(
-          `Cache length ${t.cache.length} does not match number of blocks ${this.blocks.length}`
-        );
-      for (let n = 0; n < this.blocks.length; n++) {
-        const a = this.blocks[n], i = Math.random() * 1e9, d = {
-          training: t.training,
-          seed: i,
-          attentionScores: t.attentionScores,
-          pastKV: t.cache ? t.cache[n] : void 0
-        }, S = this.config.layerConfig.checkpointing && t.training ? a.callCheckpoint(d, s) : a.call(d, s);
-        s.dispose(), s = S;
-      }
-      s = this.lnF.call(t, s);
-      const r = this.wte.project(s);
-      s.dispose();
-      let g;
-      return e && (g = this.calculateLoss(r, e)), this.endMemory("Forward"), g ? [r, g] : [r];
-    });
-  }
-  async generate(t, o, e) {
-    const l = e?.temperature ?? 1, s = e?.topK, r = e?.topP, g = e?.usePadding ?? !1, n = {
-      training: !1,
-      attentionScores: e?.attentionScores ? {
-        attentionOut: []
-      } : void 0,
-      cache: o
-    }, a = C(() => {
-      const p = t, m = p.shape[1], h = m <= this.config.gpt.blockSize ? p : p.slice(
-        [0, m - this.config.gpt.blockSize],
-        [p.shape[0], this.config.gpt.blockSize]
-      ), b = g ? this.config.gpt.blockSize - h.shape[1] : 0, w = b > 0 ? q(h, [
-        [0, 0],
-        [0, b]
-      ]) : h, [f] = this.forward(n, w), E = f.shape[1] - 1 - b, c = f.slice([0, E, 0], [f.shape[0], 1, f.shape[2]]);
-      return n.attentionScores?.attentionOut && n.attentionScores.attentionOut.forEach((y, L) => {
-        y.shape[1] !== 1 && (n.attentionScores.attentionOut[L] = U(
-          y.slice([0, E, 0], [y.shape[0], 1, y.shape[2]])
-        ), y.dispose());
-      }), f.dispose(), c.div(l).squeeze([1]);
-    });
-    let i;
-    if (r) {
-      const p = $(a), m = await p.array();
-      p.dispose();
-      const h = m[0].map((c, k) => ({ prob: c, index: k })).sort((c, k) => k.prob - c.prob);
-      let b = 0;
-      const w = new Array(h.length).fill(0);
-      for (const c of h)
-        if (b += c.prob, w[c.index] = c.prob, b >= r)
-          break;
-      const f = w.reduce((c, k) => c + k, 0), E = w.map((c) => c / f);
-      i = W(E);
-    } else if (s) {
-      const { values: p, indices: m } = J(a, s), h = I(p, 1);
-      i = X(m, h, 1), p.dispose(), m.dispose(), h.dispose();
-    } else
-      i = I(a, 1);
-    let d;
-    e?.includeProbabilities && (d = $(a));
-    const S = i.reshape([1, 1]);
-    return i.dispose(), i = S, a.dispose(), { output: i, probabilities: d, attention: n.attentionScores?.attentionOut };
-  }
-  getNumParams() {
-    return _(this.config.gpt);
-  }
-  dispose() {
-    this.wte.dispose(), this.wpe && this.wpe.dispose(), this.drop.dispose(), this.blocks.forEach((t) => t.dispose()), this.lnF.dispose();
-  }
-}
-export {
-  bt as default
-};

package/dist/TiedEmbedding-BxOerUmB.js

@@ -1,43 +0,0 @@
-import { R as a, d as s } from "./random_width-sZORGo5k.js";
-import "./index-BoWRt-10.js";
-import o from "./layers/BaseLayer.js";
-import { v as m } from "./variable-BuddVFLa.js";
-import { g as d } from "./gather-CMMy2KEG.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 n(i) {
-  return new a(i);
-}
-class S extends o {
-  vocabSize;
-  embedDim;
-  initializer;
-  WEIGHTS;
-  constructor(t, e, r) {
-    super(t, r), this.WEIGHTS = e, this.vocabSize = t.gpt.vocabSize, this.embedDim = t.gpt.nEmbed, this.initializer = n({
-      mean: 0,
-      stddev: 0.02
-    }), this.addVariable(this.WEIGHTS, m(this.initializer.apply([this.vocabSize, this.embedDim]), !0));
-  }
-  embed(t) {
-    return d(this.getVariable(this.WEIGHTS), t, 0);
-  }
-  project(t) {
-    return s(t, this.getVariable(this.WEIGHTS).transpose());
-  }
-  // Dummy, should not be used.
-  forward(t, e) {
-    return this.project(e);
-  }
-}
-export {
-  S as T,
-  n as r
-};

package/dist/utilities/generate.d.ts

@@ -1,3 +0,0 @@
-import { ITokeniser } from '../tokeniser/type';
-import { default as NanoGPT, GenerateOptions } from '../NanoGPTModel';
-export declare function generateText(tokeniser: ITokeniser, model: NanoGPT, prompt: string, length: number, options: GenerateOptions): Promise<string>;

package/dist/utilities/generate.js

@@ -1,22 +0,0 @@
-import "../index-BoWRt-10.js";
-import { t as m } from "../tensor2d-wxPAnDQy.js";
-import { c as u } from "../concat-CsxrgovM.js";
-async function v(o, r, a, c, f) {
-  if (c <= 0)
-    throw new Error("Length must be a positive integer");
-  if (a.length === 0)
-    throw new Error("Prompt cannot be an empty string");
-  const p = await o.tokenise([a], !0), s = r.config.gpt.useRope ? new Array(r.config.gpt.nLayer).fill(void 0) : void 0;
-  let e = m(p, [1, p[0].length], "int32"), t = e;
-  for (let g = 0; g < c; g++) {
-    const { output: n } = await r.generate(e, s, f), T = e, l = t;
-    t = u([t, n], 1), e = s ? n : u([e, n], 1), T.dispose(), l.dispose(), s || n.dispose();
-  }
-  const w = await t.array();
-  t.dispose();
-  const i = w[0], d = i.indexOf(o.eosToken);
-  return d !== -1 && i.splice(d), await o.decode(i);
-}
-export {
-  v as generateText
-};

package/dist/utilities/save.d.ts

@@ -1,9 +0,0 @@
-import { default as NanoGPT } from '../NanoGPTModel';
-import { ITokeniser } from '../tokeniser/type';
-export interface SaveOptions {
-    includeLog?: boolean;
-    name?: string;
-    metadata?: Record<string, unknown>;
-    files?: Record<string, unknown>;
-}
-export declare function saveModel(model: NanoGPT, tokeniser: ITokeniser, options?: SaveOptions): Promise<Blob>;

package/dist/utilities/save.js

@@ -1,61 +0,0 @@
-import { j as p } from "../jszip.min-CjP2V1VV.js";
-import b from "../tokeniser/CharTokeniser.js";
-import { save_safetensors as l } from "./safetensors.js";
-import { VERSION as y } from "../loader/load.js";
-async function N(e, a, n) {
-  const f = n?.includeLog ?? !0, s = /* @__PURE__ */ new Map();
-  e.saveWeights(s);
-  const i = new p(), r = {};
-  s.forEach((t, o) => {
-    t.length === 1 && (r[o] = t[0]);
-  });
-  const g = await l(r);
-  i.file("model.safetensors", g, { binary: !0 });
-  const c = {
-    model_type: "GenAI_NanoGPT_1",
-    vocab_size: a.getVocab().length,
-    hidden_size: e.config.gpt.nEmbed,
-    num_hidden_layers: e.config.gpt.nLayer,
-    num_attention_heads: e.config.gpt.nHead,
-    block_size: e.config.gpt.blockSize,
-    dropout: e.config.gpt.dropout,
-    biasInLinear: e.config.gpt.biasInLinear,
-    biasInLayerNorm: e.config.gpt.biasInLayerNorm,
-    mlpFactor: e.config.gpt.mlpFactor,
-    useRope: e.config.gpt.useRope
-  };
-  if (i.file("config.json", JSON.stringify(c, void 0, 4), {
-    binary: !1
-  }), i.file(
-    "meta.json",
-    JSON.stringify(
-      {
-        version: y,
-        application: "@genai-fi/nanogpt",
-        meta: n?.metadata,
-        name: n?.name
-      },
-      void 0,
-      4
-    ),
-    {
-      binary: !1
-    }
-  ), i.file(
-    "tokeniser.json",
-    JSON.stringify({
-      type: a instanceof b ? "char" : "bpe",
-      vocab: a.getVocab(),
-      merges: await a.getMerges()
-    }),
-    {
-      binary: !1
-    }
-  ), f && i.file("log.json", JSON.stringify(e.log), { binary: !1 }), n?.files)
-    for (const [t, o] of Object.entries(n.files))
-      i.file(t, JSON.stringify(o), { binary: !1 });
-  return i.generateAsync({ type: "blob" });
-}
-export {
-  N as saveModel
-};