import torch
import json
import os
from safetensors.torch import save_file

CKPT_PATH = "models/trained/Exocore.pt"
OUTPUT_DIR = "models/hf_exocorev1"

os.makedirs(OUTPUT_DIR, exist_ok=True)

ckpt = torch.load(CKPT_PATH, map_location="cpu", weights_only=False)

model_state = ckpt["model_state"]
config = ckpt["config"]
num_layers = config["num_hidden_layers"]

hf_state = {}
for k, v in model_state.items():
    if k.startswith("embed_tokens."):
        hf_k = "model." + k
    elif k.startswith("layers."):
        parts = k.split(".")
        layer_idx = parts[1]
        rest = ".".join(parts[2:])
        if rest.startswith("attn."):
            attn_part = rest.replace("attn.", "self_attn.")
            hf_k = f"model.layers.{layer_idx}.{attn_part}"
        elif rest.startswith("mlp."):
            hf_k = f"model.layers.{layer_idx}.{rest}"
        elif rest.startswith("input_layernorm."):
            hf_k = f"model.layers.{layer_idx}.input_layernorm.{rest.split('.', 1)[1]}"
        elif rest.startswith("post_attention_layernorm."):
            hf_k = f"model.layers.{layer_idx}.post_attention_layernorm.{rest.split('.', 1)[1]}"
        else:
            hf_k = f"model.layers.{layer_idx}.{rest}"
    elif k.startswith("norm."):
        hf_k = "model." + k
    elif k.startswith("lm_head."):
        hf_k = k
    else:
        hf_k = k
    hf_state[hf_k] = v.to(torch.bfloat16)

save_file(hf_state, os.path.join(OUTPUT_DIR, "model.safetensors"))

import shutil
shutil.copy("config.json", os.path.join(OUTPUT_DIR, "config.json"))

with open(os.path.join(OUTPUT_DIR, "model.safetensors.index.json"), "w") as f:
    json.dump({
        "metadata": {"total_size": sum(v.numel() * 2 for v in hf_state.values())},
        "weight_map": {k: "model.safetensors" for k in hf_state}
    }, f, indent=2)

tok_json = ckpt.get("tokenizer_json", "")
if tok_json:
    with open(os.path.join(OUTPUT_DIR, "tokenizer.json"), "w") as f:
        f.write(tok_json if isinstance(tok_json, str) else json.dumps(tok_json))

print(f"Converted! Output in {OUTPUT_DIR}")
print(f"Number of tensors: {len(hf_state)}")
print(f"Sample keys: {list(hf_state.keys())[:3]}")
print(f"Total size: {sum(v.numel() * 2 for v in hf_state.values()) / 1e9:.2f} GB")