tar2bin.py

import torch
from collections import OrderedDict


ckpt_path = "./CSAT_RCKD.pth.tar"
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)

# 1) state_dict 꺼내기
#   - 보통 {'state_dict': ...} 형태니까 먼저 이걸 시도하고,
#   - 아니면 그냥 ckpt 전체가 state_dict인 경우도 있어서 fallback
state_dict = ckpt.get("state_dict", ckpt)
# 2) DataParallel 썼으면 key 앞에 'module.' 붙어있을 수 있어서 제거
new_state_dict = OrderedDict()
for k, v in state_dict.items():
    if k.startswith("module."):
        new_k = k[len("module."):]
    else:
        new_k = k
    new_state_dict[new_k] = v

# 3) HuggingFace 관례대로 파일명 저장
torch.save(new_state_dict, "CSAT_RCKD.bin")
print("saved to pytorch_model.bin")