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')