| import os |
| import json |
| import torch |
|
|
| |
| model_dir = "./" |
|
|
| |
| index_file = os.path.join(model_dir, "pytorch_model.bin.index.json") |
| with open(index_file, "r", encoding="utf-8") as f: |
| index_data = json.load(f) |
|
|
| |
| new_weight_map = {} |
| for k, v in index_data["weight_map"].items(): |
| new_k = k.replace("module.", "", 1) if k.startswith("module.") else k |
| new_weight_map[new_k] = v |
|
|
| index_data["weight_map"] = new_weight_map |
|
|
| |
| with open(index_file, "w", encoding="utf-8") as f: |
| json.dump(index_data, f, indent=2, ensure_ascii=False) |
|
|
| print(f"已处理 {index_file}") |
|
|
| |
| for fname in os.listdir(model_dir): |
| if fname.endswith(".bin"): |
| bin_path = os.path.join(model_dir, fname) |
| print(f"处理 {bin_path} ...") |
|
|
| state_dict = torch.load(bin_path, map_location="cpu") |
|
|
| new_state_dict = {} |
| for k, v in state_dict.items(): |
| new_k = k.replace("module.", "", 1) if k.startswith("module.") else k |
| new_state_dict[new_k] = v |
|
|
| |
| torch.save(new_state_dict, bin_path) |
|
|
| print(f"完成 {fname}") |
|
|
| print("所有文件已完成 module. 前缀删除 ✅") |
