File size: 959 Bytes
03a8f89 fcae5db 03a8f89 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 | import torch
from transformers import GPT2TokenizerFast
from modeling_perklm import PerkLMConfig, PerkLM
checkpoint = torch.load('perklm.pt', map_location = 'cpu', weights_only = True)
old_config = checkpoint['config']
hf_config = PerkLMConfig(d_model = old_config['model']['d_model'],
n_heads = old_config['model']['n_heads'],
n_layers = old_config['model']['n_layers'],
d_ff = old_config['model']['d_ff'],
maxt = old_config['model']['maxt'],
dropout = old_config['model']['dropout'],
tokenizer_path = 'tokenizer')
hf_config.auto_map = {
"AutoConfig": "modeling_perklm.PerkLMConfig",
"AutoModelForCausalLM": "modeling_perklm.PerkLM"}
model = PerkLM(hf_config)
model.transformer.load_state_dict(checkpoint['model_state_dict'])
model.save_pretrained('perklm_hf')
hf_config.save_pretrained('perklm_hf')
|