import { z as y } from "../jszip.min-BZhlzntC.js";
import b from "../tokeniser/CharTokeniser.js";
import { save_safetensors as _ } from "../utilities/safetensors.js";
import { VERSION as m } from "./load.js";
async function d(i, a, o, t) {
  const g = /* @__PURE__ */ new Map();
  i.weightStore.saveWeights(g);
  const n = new y();
  if (t?.optimizer) {
    const e = await t.optimizer.saveMoments();
    n.file("optimizer.safetensors", e, { binary: !0 }), n.file("optimizer_config.json", JSON.stringify(t.optimizer.serializeConfig()), {
      binary: !1
    });
  }
  t?.trainingLog && n.file("training_log.json", JSON.stringify(t.trainingLog, void 0, 4), {
    binary: !1
  });
  const s = {};
  g.forEach((e, f) => {
    e.length === 1 && (s[f] = e[0]);
  });
  const c = await _(s);
  n.file("model.safetensors", c, { binary: !0 });
  const p = i.config.modelType;
  let r;
  p === "GenAI_NanoGPT_v1" ? r = {
    model_type: "GenAI_NanoGPT_v1",
    vocab_size: a.getVocab().length,
    hidden_size: i.config.nEmbed,
    num_hidden_layers: i.config.nLayer,
    num_attention_heads: i.config.nHead,
    block_size: i.config.blockSize,
    mlpFactor: i.config.mlpFactor,
    useRope: i.config.useRope
  } : r = {
    model_type: "GenAI_NanoGPT_v2",
    vocab_size: a.getVocab().length,
    hidden_size: i.config.nEmbed,
    num_hidden_layers: i.config.nLayer,
    num_attention_heads: i.config.nHead,
    block_size: i.config.blockSize,
    mlpFactor: i.config.mlpFactor,
    loraConfig: i.config.loraConfig ? Object.fromEntries(i.config.loraConfig) : void 0,
    loraName: i.config.loraName,
    windowSize: i.config.windowSize
  }, n.file("config.json", JSON.stringify(r, void 0, 4), {
    binary: !1
  });
  const l = {
    version: m,
    application: "@genai-fi/nanogpt",
    meta: o?.metadata,
    name: o?.name,
    training: i.metaData?.training || void 0,
    reference: i.metaData?.url || void 0,
    phase: i.metaData?.phase || void 0,
    pretrainingData: i.metaData?.pretrainingData || void 0,
    pretrainingSettings: i.metaData?.pretrainingSettings || void 0,
    generationSettings: i.metaData?.generationSettings || void 0,
    actionLog: i.metaData?.actionLog || void 0
  };
  if (n.file("meta.json", JSON.stringify(l, void 0, 4), {
    binary: !1
  }), n.file(
    "tokeniser.json",
    JSON.stringify({
      type: a instanceof b ? "char" : "bpe",
      vocab: a.getVocab(),
      merges: a.getMerges(),
      datasetID: a.datasetID,
      id: a.id
    }),
    {
      binary: !1
    }
  ), o?.files)
    for (const [e, f] of Object.entries(o.files))
      n.file(e, JSON.stringify(f), { binary: !1 });
  return n.generateAsync({ type: "blob" });
}
export {
  d as saveModel
};
