import torch
from transformers import GPT2Config, GPT2LMHeadModel

ckpt = torch.load("ckpt.pt", map_location="cpu")
state = ckpt["model"] if "model" in ckpt else ckpt

# Remove DDP prefix if present
state = {
    k.replace("_orig_mod.", ""): v
    for k, v in state.items()
}

config = GPT2Config(
    vocab_size=100277,
    n_positions=64,
    n_ctx=64,
    n_embd=128,
    n_layer=4,
    n_head=4,
    bos_token_id=100257,
    eos_token_id=100257,
)

model = GPT2LMHeadModel(config)

new_state = {}

transpose = [
    "attn.c_attn.weight",
    "attn.c_proj.weight",
    "mlp.c_fc.weight",
    "mlp.c_proj.weight",
]

for k, v in state.items():
    hf = k

    hf = hf.replace("transformer.wte", "transformer.wte")
    hf = hf.replace("transformer.wpe", "transformer.wpe")
    hf = hf.replace("transformer.h", "transformer.h")
    hf = hf.replace("transformer.ln_f", "transformer.ln_f")

    if any(hf.endswith(x) for x in transpose):
        v = v.t()

    new_state[hf] = v

missing, unexpected = model.load_state_dict(new_state, strict=False)

print("Missing:", missing)
print("Unexpected:", unexpected)

model.save_pretrained("hf_model", safe_serialization=True)