load_model.py

import os
from tkinter import NO
from turtle import mode
import numpy as np
import torch
import sys
import pandas as pd
from pathlib import Path

current_directory = Path(__file__).parent.absolute()
parent_dir = current_directory.parent
print(parent_dir)
sys.path.append(str(current_directory))

import numpy as np
from src.mamba.mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
from utils.config.mamba_config import *
from utils.evo.tokenization import Vocab
import utils.data_alphabet


def load_pretrained_model(pretrained_weight):  # emb or eos
    model_config = MambaConfig(
        model_type = "bidirectional_mamba",
        d_model=768,
        n_layer=8,
        vocab_size=71,
        ssm_cfg=dict(layer="Mamba2"),
        rms_norm=True,
        residual_in_fp32=True,
        fused_add_norm=True,
        pad_vocab_size_multiple=1,
    )

    architecture = "CodonMamba"
    rna_alphabet = utils.data_alphabet.data.Alphabet.from_architecture(architecture)
    rna_vocab = Vocab.from_esm_alphabet(rna_alphabet)

    model = MambaLMHeadModel(model_config)

    # ---- load checkpoint ----
    print(f"Loading pretrained weights from: {pretrained_weight}")
    checkpoint = torch.load(
        pretrained_weight,
        map_location=torch.device('cpu'),
        weights_only=False,  # 如果不需要，可以改成普通 torch.load(...)
    )
    state_dict = checkpoint["model_state_dict"]

    new_state_dict = {}
    for key in state_dict.keys():
        # Add "model." prefix if missing
        if not key.startswith("model."):
            new_key = "model." + key
        else:
            new_key = key
        new_state_dict[new_key] = state_dict[key]

    checkpoint["model_state_dict"] = new_state_dict
    model.load_state_dict(checkpoint["model_state_dict"], strict=False)
    
    print("Successfully loaded pretrained model weights.")

    return rna_vocab, model