modeling_transhla2.py

import torch
import torch.nn as nn
from transformers import PreTrainedModel, PretrainedConfig
from peft import LoraConfig, get_peft_model, TaskType
from transformers import EsmModel


class TransHLA2Config(PretrainedConfig):
    model_type = "transhla2"

    def __init__(
        self,
        d_model=480,
        n_layers=4,
        n_head=8,
        d_ff=64,
        cnn_num_channel=256,
        region_embedding_size=3,
        cnn_kernel_size=3,
        cnn_padding_size=1,
        cnn_stride=1,
        pooling_size=2,
        esm_model_name="facebook/esm2_t12_35M_UR50D",
        lora_r=8,
        lora_alpha=32,
        lora_dropout=0.1,
        lora_inference_mode=False,
        target_modules=None,
        return_prob=True,  # 是否在 forward 返回概率（softmax），否则返回 logits
        pad_token_id=1,     # ESM 默认 pad id
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_head = n_head
        self.d_ff = d_ff
        self.cnn_num_channel = cnn_num_channel
        self.region_embedding_size = region_embedding_size
        self.cnn_kernel_size = cnn_kernel_size
        self.cnn_padding_size = cnn_padding_size
        self.cnn_stride = cnn_stride
        self.pooling_size = pooling_size

        self.esm_model_name = esm_model_name

        self.lora_r = lora_r
        self.lora_alpha = lora_alpha
        self.lora_dropout = lora_dropout
        self.lora_inference_mode = lora_inference_mode
        self.target_modules = target_modules or ['query', 'out_proj', 'value', 'key', 'dense', 'regression']

        self.return_prob = return_prob
        self.pad_token_id = pad_token_id


class TransHLA2(PreTrainedModel):
    config_class = TransHLA2Config

    def __init__(self, config: TransHLA2Config):
        super().__init__(config)
        self.config = config

        d_model = config.d_model
        n_layers = config.n_layers
        n_head = config.n_head
        d_ff = config.d_ff
        cnn_num_channel = config.cnn_num_channel
        region_embedding_size = config.region_embedding_size
        cnn_kernel_size = config.cnn_kernel_size
        cnn_padding_size = config.cnn_padding_size
        cnn_stride = config.cnn_stride
        pooling_size = config.pooling_size

        # Backbone + LoRA
        self.esm = EsmModel.from_pretrained(config.esm_model_name)
        self.peft_config = LoraConfig(
            target_modules=config.target_modules,
            task_type=TaskType.FEATURE_EXTRACTION,
            inference_mode=config.lora_inference_mode,
            r=config.lora_r,
            lora_alpha=config.lora_alpha,
            lora_dropout=config.lora_dropout,
        )
        # 两套 LoRA 头，分别用于 epitope 和 hla 分支
        self.epitope_lora = get_peft_model(self.esm, self.peft_config)
        self.hla_lora = get_peft_model(self.esm, self.peft_config)

        # CNN branches
        self.region_cnn1 = nn.Conv1d(d_model, cnn_num_channel, region_embedding_size)
        self.region_cnn2 = nn.Conv1d(d_model, cnn_num_channel, region_embedding_size)
        self.padding1 = nn.ConstantPad1d((1, 1), 0)
        self.padding2 = nn.ConstantPad1d((0, 1), 0)
        self.relu = nn.SiLU()
        self.cnn1 = nn.Conv1d(
            cnn_num_channel, cnn_num_channel,
            kernel_size=cnn_kernel_size, padding=cnn_padding_size, stride=cnn_stride
        )
        self.cnn2 = nn.Conv1d(
            cnn_num_channel, cnn_num_channel,
            kernel_size=cnn_kernel_size, padding=cnn_padding_size, stride=cnn_stride
        )
        self.maxpooling = nn.MaxPool1d(kernel_size=pooling_size)

        # Transformer encoders (expect shape [S, B, D])
        self.epitope_transformer_layers = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=n_head, dim_feedforward=d_ff, dropout=0.2, batch_first=False
        )
        self.epitope_transformer_encoder = nn.TransformerEncoder(
            self.epitope_transformer_layers, num_layers=n_layers
        )
        self.hla_transformer_layers = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=n_head, dim_feedforward=d_ff, dropout=0.2, batch_first=False
        )
        self.hla_transformer_encoder = nn.TransformerEncoder(
            self.hla_transformer_layers, num_layers=n_layers
        )

        # Cross Attention layers (expect [S, B, D])
        self.cross_attention_epitope_layers = nn.ModuleList(
            [nn.MultiheadAttention(d_model, n_head, dropout=0.2, batch_first=False) for _ in range(4)]
        )
        self.cross_attention_hla_layers = nn.ModuleList(
            [nn.MultiheadAttention(d_model, n_head, dropout=0.2, batch_first=False) for _ in range(4)]
        )

        self.bn1 = nn.BatchNorm1d(cnn_num_channel)
        self.bn2 = nn.BatchNorm1d(cnn_num_channel)

        fused_dim = 2 * d_model + 2 * cnn_num_channel
        hidden_dim = 2 * (d_model + cnn_num_channel) // 4
        self.fc_task = nn.Sequential(
            nn.Linear(fused_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.Dropout(0.2),
            nn.SiLU(),
            nn.Linear(hidden_dim, 96),
            nn.BatchNorm1d(96),
        )
        self.classifier = nn.Linear(96, 2)

    def cnn_block1(self, x):
        # x: (B, C, L)
        return self.cnn1(self.relu(x))

    def cnn_block2(self, x):
        # x: (B, C, L)
        x = self.padding2(x)        # pad right by 1
        px = self.maxpooling(x)     # downsample
        x = self.relu(px)
        x = self.cnn1(x)
        x = self.relu(x)
        x = self.cnn1(x)
        x = px + x
        return x

    def structure_block1(self, x):
        return self.cnn2(self.relu(x))

    def structure_block2(self, x):
        x = self.padding2(x)
        px = self.maxpooling(x)
        x = self.relu(px)
        x = self.cnn2(x)
        x = self.relu(x)
        x = self.cnn2(x)
        x = px + x
        return x

    def _ensure_mapping_input(self, x):
        # 允许两种输入形式：
        # 1) 字典: {"input_ids": ..., "attention_mask": ...}
        # 2) 直接的 input_ids 张量: (B, L)
        if isinstance(x, torch.Tensor):
            # 仅用 input_ids；如需自动构造 attention_mask，可解除注释：
            # pad_id = self.config.pad_token_id
            # return {"input_ids": x, "attention_mask": (x != pad_id).long()}
            return {"input_ids": x}
        elif isinstance(x, dict):
            return x
        else:
            raise TypeError(f"Unsupported input type: {type(x)}; expected Tensor or dict.")

    def forward(self, epitope_in, hla_in, return_dict=None):
        # 兼容张量或字典输入
        epitope_in = self._ensure_mapping_input(epitope_in)
        hla_in = self._ensure_mapping_input(hla_in)

        epitope_outputs = self.epitope_lora(**epitope_in)
        hla_outputs = self.hla_lora(**hla_in)
        # last_hidden_state: (B, L, D)
        epitope_emb = epitope_outputs.last_hidden_state
        hla_emb = hla_outputs.last_hidden_state

        # Transformer encoder path (expects [S, B, D])
        epitope_trans = self.epitope_transformer_encoder(epitope_emb.transpose(0, 1))  # (L, B, D)
        hla_trans = self.hla_transformer_encoder(hla_emb.transpose(0, 1))              # (L, B, D)

        # Cross Attention
        for ca_e, ca_h in zip(self.cross_attention_epitope_layers, self.cross_attention_hla_layers):
            epitope_trans, _ = ca_e(epitope_trans, hla_trans, hla_trans)  # (L, B, D)
            hla_trans, _ = ca_h(hla_trans, epitope_trans, epitope_trans)  # (L, B, D)

        # Mean Pooling over sequence length
        epitope_mean = epitope_trans.mean(dim=0)  # (B, D)
        hla_mean = hla_trans.mean(dim=0)          # (B, D)

        # CNN branches expect (B, C, L). Convert ESM embeddings to (B, D, L)
        epitope_cnn_emb = epitope_emb.transpose(1, 2)  # (B, D, L)
        epitope_cnn_emb = self.region_cnn1(epitope_cnn_emb)  # (B, C, L')
        epitope_cnn_emb = self.padding1(epitope_cnn_emb)
        conv = epitope_cnn_emb + self.cnn_block1(self.cnn_block1(epitope_cnn_emb))
        # 迭代收缩长度直到 < 2
        while conv.size(-1) >= 2:
            conv = self.cnn_block2(conv)
        epitope_cnn_out = torch.squeeze(conv, dim=-1)  # (B, C)
        epitope_cnn_out = self.bn1(epitope_cnn_out)

        hla_cnn_emb = hla_emb.transpose(1, 2)  # (B, D, L)
        hla_cnn_emb = self.region_cnn2(hla_cnn_emb)  # (B, C, L')
        hla_cnn_emb = self.padding1(hla_cnn_emb)
        hla_conv = hla_cnn_emb + self.structure_block1(self.structure_block1(hla_cnn_emb))
        while hla_conv.size(-1) >= 2:
            hla_conv = self.structure_block2(hla_conv)
        hla_cnn_out = torch.squeeze(hla_conv, dim=-1)  # (B, C)
        hla_cnn_out = self.bn2(hla_cnn_out)

        # Fuse and classify
        representation = torch.cat((epitope_mean, hla_mean, epitope_cnn_out, hla_cnn_out), dim=1)  # (B, 2D+2C)
        features = self.fc_task(representation)  # (B, 96)
        logits = self.classifier(features)       # (B, 2)

        if self.config.return_prob:
            probs = torch.softmax(logits, dim=1)
            return probs, representation
        else:
            return logits, representation