import y from "./Evaluator.js";
import { t as L, Z as k, k as u, l as p, b as S } from "../index-DSGwv2Yx.js";
import w from "../utilities/profile.js";
import { createTensorStatistics as b } from "../checks/weights.js";
import { calculateLoss as x, calculateAccuracy as P } from "./loss.js";
import { AdamWOptimizer as T } from "./AdamW.js";
import { z as v } from "../zeros-Bw0puq_w.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 G {
  constructor(s, e, n, l) {
    this.tokenizer = e, this.model = s, this.optimizerConfig = {
      ...D,
      ...n,
      lossScaling: s.lossScaling
    };
    const m = l || new T(this.optimizerConfig);
    l && l.updateConfig(this.optimizerConfig), this.optimizer = m;
  }
  model;
  optimizer;
  running = !1;
  lastState;
  _gradientCheckpointing = !1;
  _mixedPrecision = !1;
  maskedLoss = !1;
  optimizerConfig;
  metrics = /* @__PURE__ */ new Set();
  _labelSmoothing = 0;
  _layerDrop = 0;
  _dropout = 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, e, n = !1, l = !1) {
    return L(() => {
      this.model.getProfiler()?.startMemory();
      const { xs: m, ys: i } = e, d = () => {
        const r = this.model.forward(
          {
            training: !0,
            checkpointing: this._gradientCheckpointing,
            mixedPrecision: this._mixedPrecision,
            dropout: this._dropout,
            layerDrop: this._layerDrop
          },
          m
        ), o = x(r, i, this.maskedLoss, !1, this._labelSmoothing);
        this.metrics.has("accuracy") && (s.accuracy = P(r, i), u(s.accuracy)), r.dispose();
        const a = o.mul(S(this.optimizerConfig.lossScaling));
        return o.dispose(), a;
      }, { value: t, grads: c } = k(d);
      if (n)
        this.model.getProfiler()?.endMemory("Training");
      else {
        const r = this.optimizer.applyGradients(c);
        this.metrics.has("gradientNorm") ? (s.gradientNorm = r, u(r)) : (s.gradientNorm = void 0, r.dispose());
        const o = Object.keys(c);
        this.model.weightStore.touchVariables(o), this.model.getProfiler()?.endMemory("Training"), l ? (s.gradients = c, Object.values(c).forEach((a) => u(a))) : p(c);
      }
      return t.mul(S(1 / this.optimizerConfig.lossScaling));
    });
  }
  async dummyPass() {
    const s = v([1, this.model.config.blockSize], "int32"), e = v([1, this.model.config.blockSize], "int32");
    try {
      const n = this.trainStep({}, { xs: s, ys: e }, !0);
      await n.data(), n.dispose();
    } catch (n) {
      console.error("Error during dummy pass:", n);
    } finally {
      s.dispose(), e.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, e, n) {
    const { logInterval: l = 10 } = {
      ...z,
      ...e
    };
    e.metrics && this.setMetrics(e.metrics);
    const m = Date.now(), i = this.createEmptyState();
    this.lastState = i, await this.dummyPass(), this.metrics.has("memoryUsage") && (this.model.getProfiler() || this.model.setProfiler(new w())), this.running = !0, i.logStartTime = m;
    const d = n ? new y(this.model, n, void 0, this.maskedLoss) : void 0, t = await s.iterator();
    try {
      for (; this.running; ) {
        const c = await t.next();
        if (c.done) break;
        const r = c.value, o = this.trainStep(i, r, !1);
        r.xs.dispose(), r.ys.dispose(), i.step++, i.totalSteps++, i.step % l === 0 ? await this.performLogging(o, r.xs.shape[0], e, d) : (i.gradientNorm && (i.gradientNorm.dispose(), i.gradientNorm = void 0), i.accuracy && (i.accuracy.dispose(), i.accuracy = void 0)), o.dispose();
      }
    } catch (c) {
      throw console.error("Training error:", c), c;
    }
    throw this.model.trainingState = {
      steps: i.totalSteps,
      learningRate: this.optimizer.lr,
      batchSize: e.batchSize || 32,
      loss: i.lastLoss,
      tokensProcessed: i.totalSteps * (e.batchSize || 32) * this.model.config.blockSize,
      duration: i.trainingDuration
    }, p(), this.running = !1, new Error("No log returned before training stopped.");
  }
  async performLogging(s, e, n, l) {
    const m = n?.onStep, i = this.metrics.has("gradientStatistics"), d = (await s.data())[0], t = this.lastState;
    t.lastLoss = d;
    const c = Date.now();
    t.trainingDuration += c - t.logStartTime;
    const r = t.totalSteps * e * 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: e,
      learningRate: this.metrics.has("learningRate") ? this.optimizer.lr : void 0,
      duration: t.trainingDuration,
      totalTokens: r,
      tokensPerSecond: r / (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: e,
      loss: t.lastLoss,
      tokensProcessed: r,
      duration: t.trainingDuration
    }, i && t.gradients) {
      const a = /* @__PURE__ */ new Map();
      for (const [h, g] of Object.entries(t.gradients))
        a.set(h, await b(g)), g.dispose();
      o.gradientMetrics = a;
    }
    if (l)
      try {
        const a = await l.evaluate(5);
        Array.isArray(a) ? o.validationMetrics = { loss: a[0].loss, accuracy: a[0].accuracy } : (t.validationLosses.push(a.loss), o.validationMetrics = {
          accuracy: a.accuracy,
          loss: a.loss,
          perplexity: this.metrics.has("perplexity") ? Math.exp(a.loss) : void 0
        });
      } catch (a) {
        console.error("Validation error:", a);
      }
    m && await m(o), t.logStartTime = Date.now();
  }
  async trainOnDataset(s, e, n) {
    const { logInterval: l = 10, maxEpochs: m = 1 / 0 } = {
      ...z,
      ...e
    }, i = m * (e?.epochSteps || 1e3);
    e.metrics && this.setMetrics(e.metrics);
    const d = Date.now(), t = this.createEmptyState();
    this.lastState = t, await this.dummyPass(), e?.metrics?.includes("memoryUsage") && (this.model.getProfiler() || this.model.setProfiler(new w())), this.running = !0, t.logStartTime = d;
    const c = n ? new y(this.model, n, void 0, this.maskedLoss) : void 0, r = await s.iterator();
    try {
      for (; this.running; ) {
        const o = await r.next();
        if (o.done) break;
        const a = o.value, h = t.step % l === 0, g = (e?.metrics?.includes("gradientStatistics") || !1) && h, f = this.trainStep(t, a, !1, g);
        a.xs.dispose(), a.ys.dispose(), t.step++, t.totalSteps++, h ? await this.performLogging(f, a.xs.shape[0], e, c) : (t.gradientNorm && (t.gradientNorm.dispose(), t.gradientNorm = void 0), t.accuracy && (t.accuracy.dispose(), t.accuracy = void 0)), f.dispose(), t.step >= i && this.stop();
      }
    } catch (o) {
      throw console.error("Training error:", o), p(), o;
    }
    return p(), this.running = !1, { losses: t.losses, validationLosses: t.validationLosses };
  }
}
export {
  G as default
};
