Upload convert.py

convert.py

@@ -0,0 +1,53 @@
+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)