| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| from collections import OrderedDict | |
| import torch | |
| model_path = 'model_path' | |
| out_path = 'out_path' | |
| model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, torch_dtype=torch.bfloat16) | |
| tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False, trust_remote_code=True) | |
| new_dict = OrderedDict() | |
| for k,v in model.state_dict().items(): | |
| if not 'self_attn.W_pack' in k: | |
| new_dict[k] = v | |
| continue | |
| name_base = k[:k.find('W_pack.weight')] | |
| q,k,v = [v[model.config.hidden_size*i:model.config.hidden_size*(i+1),:] for i in range(3)] | |
| new_dict[name_base + 'q_proj.weight'] = q | |
| new_dict[name_base + 'k_proj.weight'] = k | |
| new_dict[name_base + 'v_proj.weight'] = v | |
| model.save_pretrained(out_path, state_dict=new_dict) | |
| tokenizer.save_pretrained(out_path) | |