import { E as m } from "./index-DvYrXKkX.js";
import g from "./training/PreTrainer.js";
import { createTrainValidationSplit as u } from "./training/validation.js";
import c from "./training/SFTTrainer.js";
import p from "./training/tasks/splitter.js";
const r = [];
for (let n = 0; n < 256; ++n)
  r.push((n + 256).toString(16).slice(1));
function w(n, t = 0) {
  return (r[n[t + 0]] + r[n[t + 1]] + r[n[t + 2]] + r[n[t + 3]] + "-" + r[n[t + 4]] + r[n[t + 5]] + "-" + r[n[t + 6]] + r[n[t + 7]] + "-" + r[n[t + 8]] + r[n[t + 9]] + "-" + r[n[t + 10]] + r[n[t + 11]] + r[n[t + 12]] + r[n[t + 13]] + r[n[t + 14]] + r[n[t + 15]]).toLowerCase();
}
const T = new Uint8Array(16);
function D() {
  return crypto.getRandomValues(T);
}
function d(n, t, a) {
  return crypto.randomUUID ? crypto.randomUUID() : k(n);
}
function k(n, t, a) {
  n = n || {};
  const i = n.random ?? n.rng?.() ?? D();
  if (i.length < 16)
    throw new Error("Random bytes length must be >= 16");
  return i[6] = i[6] & 15 | 64, i[8] = i[8] & 63 | 128, w(i);
}
class f extends m {
  trainer;
  trainingType = "pretraining";
  hasTrained = !1;
  trainDataset;
  validationDataset;
  totalTokens = 0;
  tokensProcessed = 0;
  log = [];
  progress = null;
  options = {
    batchSize: 32,
    sftMode: "full",
    logInterval: 10
  };
  tokenizer;
  constructor(t, a, i = "pretraining", e, s) {
    if (super(), t instanceof f) {
      const o = a || t.options, h = t.options;
      let l = !1;
      t.trainingType === "sft" && o.sftMode !== h.sftMode && (l = !0), i !== t.trainingType && (l = !0), l ? (t.trainingType === "sft" ? this.trainer = new c(t.model, t.tokenizer, o) : this.trainer = new g(t.model, t.tokenizer, o), this.trainingType = i, this.options = o, this.tokenizer = t.tokenizer) : (this.trainer = t.trainer, this.trainingType = i, this.options = o, this.trainer.updateOptimizer(this.options), this.log = t.log, this.progress = t.progress, this.totalTokens = t.totalTokens, this.tokenizer = t.tokenizer, o.batchSize === h.batchSize && (this.trainDataset = t.trainDataset, this.validationDataset = t.validationDataset));
      return;
    }
    if (!a)
      throw new Error("Tokeniser must be provided when initializing Trainer with a model");
    if (!t)
      throw new Error("Model must be provided when initializing Trainer");
    this.options = e || {
      batchSize: 32,
      sftMode: "full"
    }, i === "sft" ? this.trainer = new c(t, a, e, s) : this.trainer = new g(t, a, e, s), this.trainingType = i, this.tokenizer = a;
  }
  get model() {
    return this.trainer.model;
  }
  get optimizer() {
    return this.trainer.optimizer;
  }
  stop() {
    this.trainer.stop();
  }
  reset() {
    this.hasTrained = !1, this.log = [], this.trainer.reset();
  }
  dispose() {
    this.trainer.dispose(), this.removeAllListeners();
  }
  getTotalTokens() {
    return this.totalTokens;
  }
  setOptions(t) {
    const a = new Set(
      Object.keys(t).filter(
        (i) => t[i] !== this.options[i]
      )
    );
    if (this.trainer.isRunning) {
      if (a.has("batchSize"))
        throw new Error("Cannot change batch size during training");
      if (a.has("sftMode"))
        throw new Error("Cannot change SFT mode during training");
      if (a.has("loraConfig"))
        throw new Error("Cannot change LoRA configuration during training");
      if (a.has("validationSplit"))
        throw new Error("Cannot change validation split during training");
      if (a.has("trainableWeights"))
        throw new Error("Cannot change trainable weights during training");
      if (a.has("mixedPrecision"))
        throw new Error("Cannot change mixed precision setting during training");
      if (a.has("gradientCheckpointing"))
        throw new Error("Cannot change gradient checkpointing setting during training");
    }
    this.options = {
      ...this.options,
      ...t
    }, this.trainer.updateOptimizer(this.options), a.has("metrics") && this.trainer.setMetrics(t.metrics || []);
  }
  async prepare(t = [], a) {
    const i = this.options;
    if (a && (this.model.metaData.pretrainingData = a.map((e) => ({
      id: e.id,
      name: e.name
    }))), this.trainingType === "pretraining" && this.trainer instanceof g) {
      const { trainDataset: e, validationDataset: s, size: o } = await u(
        t,
        this.trainer.tokenizer,
        this.trainer.datasetBuilder,
        i?.batchSize || 32,
        i?.validationSplit || 0.1
      ), h = o * (1 - (i?.validationSplit || 0));
      this.trainDataset = e, this.validationDataset = s, this.totalTokens = h, this.options.epochSteps = Math.ceil(
        this.totalTokens / ((i?.batchSize || 32) * this.model.config.blockSize)
      ), this.trainer.updateOptimizer(this.options);
    } else if (this.trainingType === "sft" && this.trainer instanceof c) {
      if (t instanceof Uint16Array)
        throw new Error("SFT training requires Task[] input");
      if (i?.validationSplit && i.validationSplit > 0) {
        const e = p(t, i?.validationSplit), s = await this.trainer.datasetBuilder.createSFTDataset(
          [e.training],
          i?.batchSize || 32,
          -100
        ), o = await this.trainer.datasetBuilder.createSFTDataset(
          [e.validation],
          i?.batchSize || 32,
          -100
        );
        this.validationDataset = o, this.trainDataset = s;
      } else {
        const e = await this.trainer.datasetBuilder.createSFTDataset(
          t,
          i?.batchSize || 32,
          -100
        );
        this.trainDataset = e;
      }
      this.totalTokens = t.reduce((e, s) => e + s.length, 0), this.options.epochSteps = Math.ceil(
        this.totalTokens / ((i?.batchSize || 32) * this.model.config.blockSize)
      ), this.trainer.updateOptimizer(this.options);
    }
  }
  configureModel(t) {
    const a = t?.sftMode || "full";
    if (this.trainingType === "pretraining" && (this.trainer.model.hasLoRA() && this.trainer.model.detachLoRA(), this.trainer.model.weightStore.setTrainable(["*"])), this.trainingType === "sft") {
      if (a === "lora") {
        const i = this.trainer.model;
        if (t?.loraName)
          if (i.hasLoRA(t.loraName)) {
            if (i.attachLoRA(t.loraName), t.loraConfig) {
              const e = i.lora;
              (e.alpha !== t.loraConfig.alpha || e.rank !== t.loraConfig.rank) && (i.detachLoRA(), i.deleteLoRA(t.loraName), i.createLoRA(t.loraName, t.loraConfig), i.attachLoRA(t.loraName), console.warn("Resetting LoRA with new configuration."));
            }
          } else if (t.loraConfig)
            i.createLoRA(t.loraName, t.loraConfig), i.attachLoRA(t.loraName);
          else
            throw new Error(
              `LoRA configuration must be provided to create LoRA with name ${t.loraName}`
            );
        else if (t?.loraConfig)
          if (i.hasLoRA()) {
            const e = i.lora;
            if (e.alpha !== t.loraConfig.alpha || e.rank !== t.loraConfig.rank) {
              i.detachLoRA();
              const s = t.loraName || d();
              i.createLoRA(s, t.loraConfig), i.attachLoRA(s);
            }
          } else {
            const e = t.loraName || d();
            i.createLoRA(e, t.loraConfig), i.attachLoRA(e);
          }
        else if (!i.hasLoRA()) throw new Error("LoRA configuration must be provided for lora SFT mode");
      } else
        this.trainer.model.hasLoRA() && this.trainer.model.detachLoRA();
      a === "last-layer" ? this.trainer.model.weightStore.setTrainable([
        `block_${this.trainer.model.config.nLayer - 1}_*`,
        "token_embedding"
      ]) : a === "full" && this.trainer.model.weightStore.setTrainable(["*"]);
    }
    t?.trainableWeights && this.trainer.model.weightStore.setTrainable(t.trainableWeights);
  }
  async train() {
    const t = this.options;
    if (!this.trainDataset)
      throw new Error("Dataset not prepared");
    this.hasTrained || this.trainer.setLearningRate(t?.learningRate || 1e-3), this.hasTrained = !0, this.emit("start"), this.model.metaData.pretrainingSettings = t;
    const a = Date.now();
    this.log.length > 0 && this.trainer.resumeFromLog(this.log[this.log.length - 1]), this.trainer.setGradientCheckpointing(t?.gradientCheckpointing || !1), this.trainer.setMixedPrecision(t?.mixedPrecision || !1), this.trainer.setLabelSmoothing(t?.labelSmoothing || 0), this.trainer.setDropout(t?.dropout || 0), this.trainer.setLayerDrop(t?.layerDrop || 0), this.configureModel(t), await this.trainer.trainOnDataset(
      this.trainDataset,
      {
        ...t,
        onStep: async (e) => {
          this.log.push(e), this.progress = {
            lastLog: e,
            progress: e.totalTokens / this.totalTokens,
            remaining: Math.max(0, (this.totalTokens - e.totalTokens) / e.totalTokens * e.duration)
          }, this.tokensProcessed = e.totalTokens;
          const s = this.listeners("log");
          for (const o of s)
            await o(e, this.progress);
        }
      },
      this.validationDataset
    ), this.model.metaData.actionLog = this.model.metaData.actionLog || [];
    const i = Date.now();
    this.model.metaData.actionLog.push({
      action: "pretrain",
      timestamp: i,
      duration: i - a,
      tokensProcessed: this.tokensProcessed,
      options: t
    }), this.emit("stop");
  }
  async step(t) {
    if (!this.trainDataset)
      throw new Error("Dataset not prepared");
    this.hasTrained || this.trainer.setLearningRate(t?.learningRate || 1e-3), this.hasTrained = !0, this.emit("start");
    const { log: a } = await this.trainer.stepDataset(this.trainDataset, t || {}, this.validationDataset), i = this.listeners("log");
    for (const e of i)
      await e(a, {
        lastLog: a,
        progress: a.totalTokens / this.totalTokens,
        remaining: Math.max(0, (this.totalTokens - a.totalTokens) / a.totalTokens * a.duration)
      });
    this.emit("stop");
  }
  getLog() {
    return this.log;
  }
  getProgress() {
    return this.progress;
  }
  isPrepared() {
    return this.trainDataset !== void 0 && this.validationDataset !== void 0;
  }
}
export {
  f as default
};
