@genai-fi/nanogpt 0.19.0 → 0.20.0

package/dist/training/AdamW.js

@@ -1,138 +0,0 @@
-import { adamAdjust as B } from "../ops/adamAdjust.js";
-import { adamMoments as N } from "../ops/adamMoments.js";
-import { O as S, h as b, t as c, a as M, d as w } from "../index-CUXkjxiT.js";
-import R from "./LRScheduler.js";
-import { clipScale as f } from "../ops/globalNorm.js";
-import { save_safetensors as v, load_safetensors as A } from "../utilities/safetensors.js";
-import { z as O } from "../zeros-DvZpK8s6.js";
-class _ extends S {
-  constructor(t) {
-    super(), this.config = t, this.accBeta1 = t.accBeta1 ?? t.beta1, this.accBeta2 = t.accBeta2 ?? t.beta2, this.learningRate = t.learningRate, this.beta1 = t.beta1, this.beta2 = t.beta2, this.weightDecay = t.weightDecay, this.lossScaling = t.lossScaling, this.clipNorm = t.clipNorm, this.orthGrad = t.orthoGrad ?? !1, t.epsilon === null || t.epsilon === void 0 ? this.epsilon = b().backend.epsilon() : this.epsilon = t.epsilon, this.lrScheduler = new R(t.learningRate, t);
-  }
-  className = "AdamW";
-  accBeta1 = 0;
-  accBeta2 = 0;
-  accumulatedMoments = [];
-  learningRate;
-  beta1;
-  beta2;
-  lossScaling;
-  weightDecay;
-  epsilon = null;
-  lrScheduler;
-  clipNorm;
-  orthGradEpsilon = 1e-30;
-  orthGrad;
-  get lr() {
-    return this.learningRate;
-  }
-  saveMoments() {
-    const t = {};
-    return this.accumulatedMoments.forEach((e) => {
-      t[e.originalName] = e.variable;
-    }), v(t);
-  }
-  async loadMoments(t) {
-    const e = await A(t);
-    Object.entries(e).forEach(([a, s]) => {
-      const n = s.variable(!1);
-      this.accumulatedMoments.push({ originalName: a, variable: n });
-    });
-  }
-  serializeConfig() {
-    return {
-      learningRate: this.learningRate,
-      beta1: this.beta1,
-      beta2: this.beta2,
-      accBeta1: this.accBeta1,
-      accBeta2: this.accBeta2,
-      epsilon: this.epsilon ?? void 0,
-      weightDecay: this.weightDecay,
-      lossScaling: this.lossScaling,
-      clipNorm: this.clipNorm,
-      orthoGrad: this.orthGrad,
-      ...this.lrScheduler.serializeConfig()
-    };
-  }
-  orthogonalizeGradient(t, e) {
-    return c(() => {
-      const a = t.reshape([-1]), s = e.reshape([-1]), n = a.mul(a).sum().add(this.orthGradEpsilon), h = a.mul(s).sum().div(n), o = s.sub(a.mul(h)), l = s.norm(), i = o.norm().add(this.orthGradEpsilon);
-      return o.mul(l.div(i)).reshape(e.shape);
-    });
-  }
-  updateConfig(t) {
-    const e = { ...this.config, ...t };
-    this.learningRate = e.learningRate, this.beta1 = e.beta1, this.beta2 = e.beta2, this.weightDecay = e.weightDecay, this.lossScaling = e.lossScaling, this.epsilon = e.epsilon ?? this.epsilon, this.clipNorm = e.clipNorm, this.lrScheduler.updateConfig(e, e.learningRate);
-  }
-  applyGradients(t) {
-    const e = this.lrScheduler.getNextLR();
-    this.learningRate = e;
-    const a = Array.isArray(t) ? t.map((n) => n.name) : Object.keys(t), s = c(() => {
-      const n = 1 - this.accBeta1, h = 1 - this.accBeta2;
-      let o;
-      if (this.clipNorm !== void 0) {
-        const l = a.map((i, r) => Array.isArray(t) ? t[r].tensor : t[i]);
-        o = f(l, 1 / this.lossScaling, this.clipNorm);
-      } else
-        o = M(1 / this.lossScaling);
-      return a.forEach((l, i) => {
-        const r = b().registeredVariables[l], p = !1;
-        this.accumulatedMoments[i] == null && (this.accumulatedMoments[i] = {
-          originalName: `${l}/m`,
-          variable: c(() => O([...r.shape, 2]).variable(p))
-        });
-        const m = Array.isArray(t) ? t[i].tensor : t[l];
-        if (m == null)
-          return;
-        const u = this.orthGrad ? this.orthogonalizeGradient(r, m) : m, d = this.accumulatedMoments[i].variable, g = N(d, u, this.beta1, this.beta2, o);
-        d.assign(g), this.orthGrad && u.dispose();
-        const y = B(
-          g,
-          r,
-          n,
-          h,
-          this.epsilon ?? 1e-8,
-          this.learningRate,
-          // Only apply weight decay if the variable is multi-dimensional (e.g. weights, not biases)
-          r.shape.length > 1 ? this.weightDecay : 0
-        );
-        r.assign(y);
-      }), this.accBeta1 = this.accBeta1 * this.beta1, this.accBeta2 = this.accBeta2 * this.beta2, o;
-    });
-    return this.incrementIterations(), s;
-  }
-  dispose() {
-    this.accumulatedMoments != null && w(this.accumulatedMoments.map((t) => t.variable));
-  }
-  async getWeights() {
-    const t = [...this.accumulatedMoments];
-    return [await this.saveIterations()].concat(
-      t.map((e) => ({ name: e.originalName, tensor: e.variable }))
-    );
-  }
-  async setWeights(t) {
-    t = await this.extractIterations(t), c(() => {
-      this.accBeta1 = Math.pow(this.beta1, this.iterations_ + 1), this.accBeta2 = Math.pow(this.beta2, this.iterations_ + 1);
-    });
-    const e = t.length / 2, a = !1;
-    this.accumulatedMoments = t.slice(0, e).map((s) => ({
-      originalName: s.name,
-      variable: s.tensor.variable(a)
-    }));
-  }
-  getConfig() {
-    return {
-      learningRate: this.learningRate,
-      beta1: this.beta1,
-      beta2: this.beta2,
-      epsilon: this.epsilon
-    };
-  }
-  /** @nocollapse */
-  static fromConfig(t, e) {
-    return new t(e.learningRate, e.beta1, e.beta2, e.epsilon);
-  }
-}
-export {
-  _ as AdamWOptimizer
-};

package/dist/training/BasicTrainer.d.ts

@@ -1,63 +0,0 @@
-import { ITokeniser } from '../tokeniser/type';
-import { Scalar, Tensor } from '@tensorflow/tfjs-core';
-import { Dataset } from '@tensorflow/tfjs-data';
-import { default as Model, ModelForwardAttributes } from '../models/model';
-import { AdamWOptimizerConfig, TrainingLogEntry, TrainingMetrics, TrainingOptions, TrainingState } from './types';
-import { AdamWOptimizer } from './AdamW';
-export default class BasicTrainer {
-    tokenizer: ITokeniser;
-    model: Model<ModelForwardAttributes>;
-    optimizer: AdamWOptimizer;
-    protected running: boolean;
-    protected lastState?: TrainingState;
-    protected _gradientCheckpointing: boolean;
-    protected _mixedPrecision: boolean;
-    protected maskedLoss: boolean;
-    protected optimizerConfig: AdamWOptimizerConfig;
-    protected metrics: Set<TrainingMetrics>;
-    protected _labelSmoothing: number;
-    protected _layerDrop: number;
-    protected _dropout: number;
-    constructor(model: Model<ModelForwardAttributes>, tokenizer: ITokeniser, optConfig?: Partial<AdamWOptimizerConfig>, optimizer?: AdamWOptimizer);
-    setLossMasking(): void;
-    setGradientCheckpointing(enabled: boolean): void;
-    setMixedPrecision(enabled: boolean): void;
-    setLabelSmoothing(smoothing: number): void;
-    setDropout(dropout: number): void;
-    setLayerDrop(layerDrop: number): void;
-    setLearningRate(learningRate: number): void;
-    setMetrics(metrics: TrainingMetrics[]): void;
-    reset(): void;
-    stop(): void;
-    get isRunning(): boolean;
-    getOptimizer(): AdamWOptimizer;
-    updateOptimizer(config?: Partial<AdamWOptimizerConfig>): void;
-    resumeFromLog(log: TrainingLogEntry): void;
-    protected trainStep(state: Partial<TrainingState>, batch: {
-        xs: Tensor;
-        ys: Tensor;
-    }, dummy?: boolean, keepGrads?: boolean): Scalar;
-    private dummyPass;
-    dispose(): void;
-    private createEmptyState;
-    stepDataset(dataset: Dataset<{
-        xs: Tensor;
-        ys: Tensor;
-    }>, options: Partial<TrainingOptions>, validationDataset?: Dataset<{
-        xs: Tensor;
-        ys: Tensor;
-    }>): Promise<{
-        log: TrainingLogEntry;
-    }>;
-    private performLogging;
-    trainOnDataset(dataset: Dataset<{
-        xs: Tensor;
-        ys: Tensor;
-    }>, options: Partial<TrainingOptions>, validationDataset?: Dataset<{
-        xs: Tensor;
-        ys: Tensor;
-    }>): Promise<{
-        losses: number[];
-        validationLosses: number[];
-    }>;
-}

package/dist/training/BasicTrainer.js

@@ -1,265 +0,0 @@
-import S from "./Evaluator.js";
-import { t as k, v as x, k as y, d as u, a as w } from "../index-CUXkjxiT.js";
-import v from "../utilities/profile.js";
-import { createTensorStatistics as N } from "../checks/weights.js";
-import { calculateLoss as b, calculateAccuracy as P } from "./loss.js";
-import { AdamWOptimizer as T } from "./AdamW.js";
-import { z as L } from "../zeros-DvZpK8s6.js";
-const z = {
-  logInterval: 1,
-  maxEpochs: 100,
-  sftMode: "full",
-  batchSize: 32
-}, D = {
-  learningRate: 3e-4,
-  beta1: 0.9,
-  beta2: 0.99,
-  epsilon: 1e-8,
-  weightDecay: 0.01,
-  warmupSteps: 100,
-  decayEpochs: 100,
-  epochSteps: 1e4,
-  minLearningRate: 1e-5,
-  lossScaling: 1
-};
-class B {
-  constructor(s, i, n, l) {
-    this.tokenizer = i, this.model = s, this.optimizerConfig = {
-      ...D,
-      ...n,
-      lossScaling: s.lossScaling
-    };
-    const d = l || new T(this.optimizerConfig);
-    l && l.updateConfig(this.optimizerConfig), this.optimizer = d;
-  }
-  model;
-  optimizer;
-  running = !1;
-  lastState;
-  _gradientCheckpointing = !1;
-  _mixedPrecision = !1;
-  maskedLoss = !1;
-  optimizerConfig;
-  metrics = /* @__PURE__ */ new Set();
-  _labelSmoothing = 0;
-  _layerDrop = 0;
-  _dropout = 0;
-  setLossMasking() {
-    this.maskedLoss = !0;
-  }
-  setGradientCheckpointing(s) {
-    this._gradientCheckpointing = s;
-  }
-  setMixedPrecision(s) {
-    this._mixedPrecision = s;
-  }
-  setLabelSmoothing(s) {
-    this._labelSmoothing = s;
-  }
-  setDropout(s) {
-    this._dropout = s;
-  }
-  setLayerDrop(s) {
-    this._layerDrop = s;
-  }
-  setLearningRate(s) {
-    this.optimizerConfig.learningRate = s, this.updateOptimizer();
-  }
-  setMetrics(s) {
-    this.metrics = new Set(s);
-  }
-  reset() {
-    this.lastState = void 0, this.running = !1;
-  }
-  stop() {
-    this.running = !1;
-  }
-  get isRunning() {
-    return this.running;
-  }
-  getOptimizer() {
-    return this.optimizer;
-  }
-  updateOptimizer(s) {
-    s && (this.optimizerConfig = { ...this.optimizerConfig, ...s }), this.optimizer.updateConfig(this.optimizerConfig);
-  }
-  resumeFromLog(s) {
-    (!this.lastState || this.lastState.step === 0) && (this.lastState = {
-      losses: [],
-      validationLosses: [],
-      logStartTime: 0,
-      step: s.step,
-      lastLoss: s.trainingMetrics.loss,
-      totalSteps: s.step,
-      trainingDuration: s.duration
-    });
-  }
-  // A single forward pass, backward pass, and optimizer step
-  trainStep(s, i, n = !1, l = !1) {
-    return k(() => {
-      this.model.getProfiler()?.startMemory();
-      const { xs: d, ys: r } = i, m = () => {
-        const a = this.model.forward(
-          {
-            training: !0,
-            checkpointing: this._gradientCheckpointing,
-            mixedPrecision: this._mixedPrecision,
-            dropout: this._dropout,
-            layerDrop: this._layerDrop,
-            ropePositionOffset: 0
-          },
-          d
-        ), o = b(a, r, this.maskedLoss, !1, this._labelSmoothing);
-        this.metrics.has("accuracy") && (s.accuracy = P(a, r), y(s.accuracy)), a.dispose();
-        const e = o.mul(w(this.optimizerConfig.lossScaling));
-        return o.dispose(), e;
-      }, { value: t, grads: c } = x(m);
-      if (n)
-        this.model.getProfiler()?.endMemory("Training");
-      else {
-        const a = this.optimizer.applyGradients(c);
-        this.metrics.has("gradientNorm") ? (s.gradientNorm = a, y(a)) : (s.gradientNorm = void 0, a.dispose());
-        const o = Object.keys(c);
-        this.model.weightStore.touchVariables(o), this.model.getProfiler()?.endMemory("Training"), l ? (s.gradients = c, Object.values(c).forEach((e) => y(e))) : u(c);
-      }
-      return t.mul(w(1 / this.optimizerConfig.lossScaling));
-    });
-  }
-  async dummyPass() {
-    const s = L([1, this.model.config.blockSize], "int32"), i = L([1, this.model.config.blockSize], "int32");
-    try {
-      const n = this.trainStep({}, { xs: s, ys: i }, !0);
-      await n.data(), n.dispose();
-    } catch (n) {
-      console.error("Error during dummy pass:", n);
-    } finally {
-      s.dispose(), i.dispose();
-    }
-  }
-  dispose() {
-    this.optimizer && this.optimizer.dispose();
-  }
-  createEmptyState() {
-    return {
-      step: 0,
-      lastLoss: 1e6,
-      totalSteps: 0,
-      losses: [],
-      validationLosses: [],
-      logStartTime: 0,
-      trainingDuration: 0,
-      ...this.lastState || {}
-    };
-  }
-  async stepDataset(s, i, n) {
-    const { logInterval: l = 10 } = {
-      ...z,
-      ...i
-    };
-    i.metrics && this.setMetrics(i.metrics);
-    const d = Date.now(), r = this.createEmptyState();
-    this.lastState = r, await this.dummyPass(), this.metrics.has("memoryUsage") && (this.model.getProfiler() || this.model.setProfiler(new v())), this.running = !0, r.logStartTime = d;
-    const m = n ? new S(this.model, n, this.maskedLoss) : void 0, t = await s.iterator();
-    try {
-      for (; this.running; ) {
-        const c = await t.next();
-        if (c.done) break;
-        const a = c.value, o = this.trainStep(r, a, !1);
-        if (i.debug) {
-          const e = (await o.data())[0];
-          if (isNaN(e) || !isFinite(e))
-            throw console.error("Invalid loss value:", e), console.error("Batch xs:", a.xs.toString()), console.error("Batch ys:", a.ys.toString()), console.error("State:", r), new Error("Loss is NaN or Infinity");
-          console.log(`Step ${r.step}: Loss = ${e}`);
-        }
-        a.xs.dispose(), a.ys.dispose(), r.step++, r.totalSteps++, r.step % l === 0 ? await this.performLogging(o, a.xs.shape[0], i, m) : (r.gradientNorm && (r.gradientNorm.dispose(), r.gradientNorm = void 0), r.accuracy && (r.accuracy.dispose(), r.accuracy = void 0)), o.dispose();
-      }
-    } catch (c) {
-      throw console.error("Training error:", c), c;
-    }
-    throw this.model.trainingState = {
-      steps: r.totalSteps,
-      learningRate: this.optimizer.lr,
-      batchSize: i.batchSize || 32,
-      loss: r.lastLoss,
-      tokensProcessed: r.totalSteps * (i.batchSize || 32) * this.model.config.blockSize,
-      duration: r.trainingDuration
-    }, u(), this.running = !1, new Error("No log returned before training stopped.");
-  }
-  async performLogging(s, i, n, l) {
-    const d = n?.onStep, r = this.metrics.has("gradientStatistics"), m = (await s.data())[0], t = this.lastState;
-    t.lastLoss = m;
-    const c = Date.now();
-    t.trainingDuration += c - t.logStartTime;
-    const a = t.totalSteps * i * this.model.config.blockSize, o = {
-      trainingMetrics: {
-        loss: t.lastLoss,
-        perplexity: this.metrics.has("perplexity") ? Math.exp(t.lastLoss) : void 0,
-        accuracy: t.accuracy ? (await t.accuracy.data())[0] : void 0
-      },
-      step: t.step,
-      time: Date.now() - t.logStartTime,
-      gradientNorm: t.gradientNorm ? (await t.gradientNorm.data())[1] : void 0,
-      batchSize: i,
-      learningRate: this.metrics.has("learningRate") ? this.optimizer.lr : void 0,
-      duration: t.trainingDuration,
-      totalTokens: a,
-      tokensPerSecond: a / (t.trainingDuration / 1e3),
-      memoryUsage: this.metrics.has("memoryUsage") ? this.model.getProfiler()?.getPeakMemory() || 0 : void 0
-    };
-    if (t.gradientNorm && (t.gradientNorm.dispose(), t.gradientNorm = void 0), t.accuracy && (t.accuracy.dispose(), t.accuracy = void 0), this.model.trainingState = {
-      steps: t.totalSteps,
-      learningRate: this.optimizer.lr,
-      batchSize: i,
-      loss: t.lastLoss,
-      tokensProcessed: a,
-      duration: t.trainingDuration
-    }, r && t.gradients) {
-      const e = /* @__PURE__ */ new Map();
-      for (const [h, g] of Object.entries(t.gradients))
-        e.set(h, await N(g)), g.dispose();
-      o.gradientMetrics = e;
-    }
-    if (l)
-      try {
-        const e = await l.evaluate(5);
-        Array.isArray(e) ? o.validationMetrics = { loss: e[0].loss, accuracy: e[0].accuracy } : (t.validationLosses.push(e.loss), o.validationMetrics = {
-          accuracy: e.accuracy,
-          loss: e.loss,
-          perplexity: this.metrics.has("perplexity") ? Math.exp(e.loss) : void 0
-        });
-      } catch (e) {
-        console.error("Validation error:", e);
-      }
-    d && await d(o), t.logStartTime = Date.now();
-  }
-  async trainOnDataset(s, i, n) {
-    const { logInterval: l = 10, maxEpochs: d = 1 / 0 } = {
-      ...z,
-      ...i
-    }, r = d * (i?.epochSteps || 1e3);
-    i.metrics && this.setMetrics(i.metrics);
-    const m = Date.now(), t = this.createEmptyState();
-    this.lastState = t, await this.dummyPass(), i?.metrics?.includes("memoryUsage") && (this.model.getProfiler() || this.model.setProfiler(new v())), this.running = !0, t.logStartTime = m;
-    const c = n ? new S(this.model, n, this.maskedLoss) : void 0, a = await s.iterator();
-    try {
-      for (; this.running; ) {
-        const o = await a.next();
-        if (o.done) break;
-        const e = o.value, h = t.step % l === 0, g = (i?.metrics?.includes("gradientStatistics") || !1) && h, f = this.trainStep(t, e, !1, g);
-        if (i.debug) {
-          const p = (await f.data())[0];
-          if (isNaN(p) || !isFinite(p))
-            throw console.error("Invalid loss value:", p), console.error("Batch xs:", await e.xs.array()), console.error("Batch ys:", await e.ys.array()), console.error("State:", t), new Error("Loss is NaN or Infinity");
-          console.log(`Step ${t.step}: Loss = ${p}`);
-        }
-        e.xs.dispose(), e.ys.dispose(), t.step++, t.totalSteps++, h ? await this.performLogging(f, e.xs.shape[0], i, c) : (t.gradientNorm && (t.gradientNorm.dispose(), t.gradientNorm = void 0), t.accuracy && (t.accuracy.dispose(), t.accuracy = void 0)), f.dispose(), t.step >= r && this.stop();
-      }
-    } catch (o) {
-      throw console.error("Training error:", o), u(), o;
-    }
-    return u(), this.running = !1, { losses: t.losses, validationLosses: t.validationLosses };
-  }
-}
-export {
-  B as default
-};

package/dist/training/DatasetBuilder.d.ts

@@ -1,26 +0,0 @@
-import { Tensor } from '@tensorflow/tfjs-core';
-import { Conversation, ITokeniser } from '../tokeniser/type';
-import { Dataset } from '@tensorflow/tfjs-data';
-export declare const PAGE_FACTOR = 8;
-export declare function flattenTokens(textData: Conversation[][], tokenizer: ITokeniser): Uint16Array;
-export declare function flattenTokensWithMask(textData: Conversation[][], tokenizer: ITokeniser): {
-    tokens: Uint16Array;
-    mask: Uint8Array;
-};
-export declare function shuffle(array: Uint32Array): Uint32Array;
-export interface DatasetState {
-    shuffledIndexes: Uint32Array;
-    step: number;
-}
-export declare class DatasetBuilder {
-    tokenizer: ITokeniser;
-    blockSize: number;
-    constructor(tokenizer: ITokeniser, blockSize?: number);
-    createTextDataset(flatTokens: Uint16Array, batchSize?: number, indexes?: Uint32Array, mask?: Uint8Array, ignoreIndex?: number): Promise<{
-        dataset: Dataset<{
-            xs: Tensor;
-            ys: Tensor;
-        }>;
-        state: DatasetState;
-    }>;
-}

package/dist/training/DatasetBuilder.js

@@ -1,86 +0,0 @@
-import { t as x } from "../index-CUXkjxiT.js";
-import { d as g, i as m } from "../dataset-CGGp1z9P.js";
-import "../index-Cp39cXWe.js";
-function p(e) {
-  return g(async () => {
-    const t = await e();
-    return m(() => t.next());
-  });
-}
-const I = 8;
-function z(e, t) {
-  const r = e.map((c) => t.encodeConversation(c)).flat();
-  return new Uint16Array(r);
-}
-function A(e, t) {
-  const s = e.map((i) => t.encodeConversation(i, !1, !0));
-  console.log("Tokenised Texts with Mask:", s);
-  const r = s.map((i) => i.tokens).flat(), c = s.map((i) => i.mask).flat();
-  return { tokens: new Uint16Array(r), mask: new Uint8Array(c.map((i) => i ? 1 : 0)) };
-}
-function u(e) {
-  for (let t = e.length - 1; t > 0; t--) {
-    const s = Math.floor(Math.random() * (t + 1));
-    [e[t], e[s]] = [e[s], e[t]];
-  }
-  return e;
-}
-class S {
-  tokenizer;
-  blockSize;
-  constructor(t, s = 128) {
-    this.tokenizer = t, this.blockSize = s;
-  }
-  // Create dataset from text files
-  async createTextDataset(t, s = 32, r, c, i = 65535) {
-    if (t.length < this.blockSize + 1)
-      throw new Error(`Not enough tokens (${t.length}) for block size ${this.blockSize}`);
-    const o = {
-      shuffledIndexes: new Uint32Array(t.length),
-      step: 0
-    };
-    if (r)
-      o.shuffledIndexes = r;
-    else {
-      o.shuffledIndexes = new Uint32Array(t.length);
-      for (let n = 0; n < t.length; n++)
-        o.shuffledIndexes[n] = n;
-      u(o.shuffledIndexes);
-    }
-    const d = (function* () {
-      for (; ; ) {
-        const n = o.shuffledIndexes[o.step++];
-        if (o.step >= o.shuffledIndexes.length && (o.step = 0, u(o.shuffledIndexes)), n + this.blockSize + 1 > t.length)
-          continue;
-        const a = new Int32Array(t.subarray(n, n + this.blockSize)), k = t.subarray(n + 1, n + this.blockSize + 1), l = new Int32Array(k);
-        if (c) {
-          let h = 0;
-          for (let f = 0; f < l.length; f++)
-            c[n + 1 + f] === 0 && (l[f] = i, h++);
-          if (h === l.length)
-            continue;
-        }
-        yield { xs: a, ys: l };
-      }
-    }).bind(this);
-    return {
-      dataset: p(d).batch(s).map((n) => {
-        const a = n;
-        return x(() => ({
-          xs: a.xs.cast("int32"),
-          ys: a.ys.cast("int32")
-          // this.tf.oneHot(batchData.ys.cast('int32'), this.tokenizer.vocabSize),
-        }));
-      }).prefetch(2),
-      // Smaller prefetch to reduce memory pressure
-      state: o
-    };
-  }
-}
-export {
-  S as DatasetBuilder,
-  I as PAGE_FACTOR,
-  z as flattenTokens,
-  A as flattenTokensWithMask,
-  u as shuffle
-};

package/dist/training/Evaluator.d.ts

@@ -1,19 +0,0 @@
-import { Dataset } from '@tensorflow/tfjs-data';
-import { TensorContainer } from '@tensorflow/tfjs-core';
-import { default as Model, ModelForwardAttributes } from '../models/model';
-interface Result {
-    loss: number;
-    accuracy: number;
-}
-export default class Evaluator {
-    private model;
-    private iterator?;
-    private xs?;
-    private ys?;
-    private masked;
-    constructor(model: Model<ModelForwardAttributes>, dataset: Dataset<TensorContainer>, masked?: boolean);
-    dispose(): void;
-    private calculateBatchLoss;
-    evaluate(maxBatches?: number): Promise<Result | Result[]>;
-}
-export {};

package/dist/training/Evaluator.js

@@ -1,39 +0,0 @@
-import { t as d } from "../index-CUXkjxiT.js";
-import { calculateLoss as f, calculateAccuracy as p } from "./loss.js";
-class b {
-  constructor(o, t, a) {
-    this.model = o, this.masked = !!a, this.iterator = t.iterator();
-  }
-  iterator;
-  xs;
-  ys;
-  masked = !1;
-  dispose() {
-    this.xs && this.xs.dispose(), this.ys && this.ys.dispose();
-  }
-  async calculateBatchLoss(o, t, a, r) {
-    const [l, e] = d(() => {
-      const s = this.model.forward({ training: !1 }, o), h = f(s, t, r, a), y = p(s, t);
-      return s.dispose(), [h, y];
-    }), u = await l.array(), n = await e.array(), c = u, i = n;
-    return e.dispose(), l.dispose(), Array.isArray(c) ? c.map((s) => ({ loss: s, accuracy: i })) : { loss: c, accuracy: i };
-  }
-  async evaluate(o = 100) {
-    let t = 0, a = 0, r = 0;
-    if (this.iterator) {
-      const l = await this.iterator;
-      for (let e = 0; e < o; e++) {
-        const u = await l.next();
-        if (u.done) break;
-        const n = u.value, { xs: c, ys: i } = n, s = await this.calculateBatchLoss(c, i, !1, this.masked);
-        c.dispose(), i.dispose(), t += s.loss, a += s.accuracy, r++;
-      }
-      return { loss: t / r, accuracy: a / r };
-    } else if (this.xs && this.ys)
-      return this.calculateBatchLoss(this.xs, this.ys, !0, !0);
-    throw new Error("No data available for evaluation");
-  }
-}
-export {
-  b as default
-};

package/dist/training/LRScheduler.d.ts

@@ -1,12 +0,0 @@
-import { LRSchedulerConfig } from './types';
-export default class LRScheduler {
-    protected learningRate: number;
-    private config;
-    private step;
-    private startLearningRate;
-    constructor(learningRate: number, config: LRSchedulerConfig);
-    serializeConfig(): LRSchedulerConfig;
-    updateConfig(newConfig: Partial<LRSchedulerConfig>, learningRate?: number): void;
-    get lr(): number;
-    getNextLR(): number;
-}

package/dist/training/LRScheduler.js

@@ -1,34 +0,0 @@
-class o {
-  constructor(i, t) {
-    this.learningRate = i, this.config = t, this.startLearningRate = i, t.step !== void 0 && (this.step = t.step);
-  }
-  step = 0;
-  startLearningRate;
-  serializeConfig() {
-    return {
-      ...this.config,
-      step: this.step
-    };
-  }
-  updateConfig(i, t) {
-    this.config = { ...this.config, ...i }, t !== void 0 && (this.startLearningRate = t);
-  }
-  get lr() {
-    return this.learningRate;
-  }
-  getNextLR() {
-    const i = this.step;
-    if (this.config.warmupSteps > 0 && i < this.config.warmupSteps) {
-      const r = (i + 1) / this.config.warmupSteps, e = this.startLearningRate * r;
-      return this.learningRate = e, this.step++, e;
-    }
-    const t = this.config.epochSteps * this.config.decayEpochs;
-    if (i >= t || t <= this.config.warmupSteps)
-      return this.learningRate = this.config.minLearningRate, this.step++, this.config.minLearningRate;
-    const n = (i - this.config.warmupSteps) / (t - this.config.warmupSteps), a = 0.5 * (1 + Math.cos(Math.PI * n)), s = this.config.minLearningRate + a * (this.startLearningRate - this.config.minLearningRate);
-    return this.learningRate = s, this.step++, s;
-  }
-}
-export {
-  o as default
-};