Update modeling_transhla2.py

modeling_transhla2.py

@@ -1,143 +1,143 @@
-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, **kwargs):
-        super().__init__(**kwargs)
-        self.d_model = d_model
-        # 可加入其它自定义参数
-class TransHLA2Config(nn.Module):
-    config_class = TransHLA2Config
-    def __init__(self):
-        super(TransHLA2Config, self).__init__()
-
-        n_layers = 4
-        n_head = 8
-        d_model = 480
-        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
-        self.model_name_or_path = "facebook/esm2_t12_35M_UR50D"
-        self.tokenizer_name_or_path = "facebook/esm2_t12_35M_UR50D"
-        self.peft_config = LoraConfig(
-            target_modules=['query', 'out_proj', 'value', 'key', 'dense', 'regression'], task_type=TaskType.FEATURE_EXTRACTION, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1)
-        self.esm = EsmModel.from_pretrained(self.model_name_or_path)
-        self.epitope_lora = get_peft_model(self.esm, self.peft_config)
-        self.hla_lora = get_peft_model(self.esm, self.peft_config)
-
-        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)
-        self.epitope_transformer_layers = nn.TransformerEncoderLayer(
-            d_model=d_model, nhead=n_head, dim_feedforward=d_ff, dropout=0.2)
-        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)
-        self.hla_transformer_encoder = nn.TransformerEncoder(
-            self.hla_transformer_layers, num_layers=n_layers)
-        
-        # Cross Attention layers
-        self.cross_attention_epitope_layers = nn.ModuleList(
-            [nn.MultiheadAttention(d_model, n_head, dropout=0.2) for _ in range(4)])
-        self.cross_attention_hla_layers = nn.ModuleList(
-            [nn.MultiheadAttention(d_model, n_head, dropout=0.2) for _ in range(4)])
-
-        self.bn1 = nn.BatchNorm1d(cnn_num_channel)
-        self.bn2 = nn.BatchNorm1d(cnn_num_channel)
-        self.fc_task = nn.Sequential(
-            nn.Linear(2*d_model + 2*cnn_num_channel, 2 * (d_model + cnn_num_channel) // 4),
-            nn.BatchNorm1d(2 * (d_model + cnn_num_channel) // 4),
-            nn.Dropout(0.2),
-            nn.SiLU(),
-            nn.Linear(2 * (d_model + cnn_num_channel) // 4, 96),
-            nn.BatchNorm1d(96),
-        )
-        self.classifier = nn.Linear(96, 2)
-
-    def cnn_block1(self, x):
-        return self.cnn1(self.relu(x))
-
-    def cnn_block2(self, x):
-        x = self.padding2(x)
-        px = self.maxpooling(x)
-        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 forward(self, epitope_in, hla_in):
-        epitope_emb = self.epitope_lora(epitope_in).last_hidden_state
-        hla_emb = self.hla_lora(hla_in).last_hidden_state
-
-        epitope_trans = self.epitope_transformer_encoder(epitope_emb.transpose(0, 1))
-        hla_trans = self.hla_transformer_encoder(hla_emb.transpose(0, 1))
-
-        # Cross Attention layers
-        for cross_attention_epitope, cross_attention_hla in zip(self.cross_attention_epitope_layers, self.cross_attention_hla_layers):
-            epitope_trans, _ = cross_attention_epitope(epitope_trans, hla_trans, hla_trans)
-            hla_trans, _ = cross_attention_hla(hla_trans, epitope_trans, epitope_trans)
-
-        # Mean Pooling
-        epitope_mean = epitope_trans.mean(dim=0)
-        hla_mean = hla_trans.mean(dim=0)
-        
-        epitope_cnn_emb = self.region_cnn1(epitope_emb.transpose(1, 2))
-        epitope_cnn_emb = self.padding1(epitope_cnn_emb)
-        conv = epitope_cnn_emb + self.cnn_block1(self.cnn_block1(epitope_cnn_emb))
-        while conv.size(-1) >= 2:
-            conv = self.cnn_block2(conv)
-        epitope_cnn_out = torch.squeeze(conv, dim=-1)
-        epitope_cnn_out = self.bn1(epitope_cnn_out)
-
-        hla_cnn_emb = self.region_cnn2(hla_emb.transpose(1, 2))
-        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)
-        hla_cnn_out = self.bn2(hla_cnn_out)
-        # CNN Blocks
-
-        # Concatenate and pass through MLP
-        representation = torch.cat((epitope_mean, hla_mean, epitope_cnn_out, hla_cnn_out), dim=1)
-        reduction_feature = self.fc_task(representation)
-        logits_clsf = self.classifier(reduction_feature)
-        logits_clsf = torch.nn.functional.softmax(logits_clsf, dim=1)
-        return logits_clsf, representation
-
-
+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, **kwargs):
+        super().__init__(**kwargs)
+        self.d_model = d_model
+        # 可加入其它自定义参数
+class TransHLA2(PreTrainedModel):
+    config_class = TransHLA2Config
+    def __init__(self):
+        super().__init__(config)
+
+        n_layers = 4
+        n_head = 8
+        d_model = 480
+        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
+        self.model_name_or_path = "facebook/esm2_t12_35M_UR50D"
+        self.tokenizer_name_or_path = "facebook/esm2_t12_35M_UR50D"
+        self.peft_config = LoraConfig(
+            target_modules=['query', 'out_proj', 'value', 'key', 'dense', 'regression'], task_type=TaskType.FEATURE_EXTRACTION, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1)
+        self.esm = EsmModel.from_pretrained(self.model_name_or_path)
+        self.epitope_lora = get_peft_model(self.esm, self.peft_config)
+        self.hla_lora = get_peft_model(self.esm, self.peft_config)
+
+        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)
+        self.epitope_transformer_layers = nn.TransformerEncoderLayer(
+            d_model=d_model, nhead=n_head, dim_feedforward=d_ff, dropout=0.2)
+        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)
+        self.hla_transformer_encoder = nn.TransformerEncoder(
+            self.hla_transformer_layers, num_layers=n_layers)
+        
+        # Cross Attention layers
+        self.cross_attention_epitope_layers = nn.ModuleList(
+            [nn.MultiheadAttention(d_model, n_head, dropout=0.2) for _ in range(4)])
+        self.cross_attention_hla_layers = nn.ModuleList(
+            [nn.MultiheadAttention(d_model, n_head, dropout=0.2) for _ in range(4)])
+
+        self.bn1 = nn.BatchNorm1d(cnn_num_channel)
+        self.bn2 = nn.BatchNorm1d(cnn_num_channel)
+        self.fc_task = nn.Sequential(
+            nn.Linear(2*d_model + 2*cnn_num_channel, 2 * (d_model + cnn_num_channel) // 4),
+            nn.BatchNorm1d(2 * (d_model + cnn_num_channel) // 4),
+            nn.Dropout(0.2),
+            nn.SiLU(),
+            nn.Linear(2 * (d_model + cnn_num_channel) // 4, 96),
+            nn.BatchNorm1d(96),
+        )
+        self.classifier = nn.Linear(96, 2)
+
+    def cnn_block1(self, x):
+        return self.cnn1(self.relu(x))
+
+    def cnn_block2(self, x):
+        x = self.padding2(x)
+        px = self.maxpooling(x)
+        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 forward(self, epitope_in, hla_in):
+        epitope_emb = self.epitope_lora(epitope_in).last_hidden_state
+        hla_emb = self.hla_lora(hla_in).last_hidden_state
+
+        epitope_trans = self.epitope_transformer_encoder(epitope_emb.transpose(0, 1))
+        hla_trans = self.hla_transformer_encoder(hla_emb.transpose(0, 1))
+
+        # Cross Attention layers
+        for cross_attention_epitope, cross_attention_hla in zip(self.cross_attention_epitope_layers, self.cross_attention_hla_layers):
+            epitope_trans, _ = cross_attention_epitope(epitope_trans, hla_trans, hla_trans)
+            hla_trans, _ = cross_attention_hla(hla_trans, epitope_trans, epitope_trans)
+
+        # Mean Pooling
+        epitope_mean = epitope_trans.mean(dim=0)
+        hla_mean = hla_trans.mean(dim=0)
+        
+        epitope_cnn_emb = self.region_cnn1(epitope_emb.transpose(1, 2))
+        epitope_cnn_emb = self.padding1(epitope_cnn_emb)
+        conv = epitope_cnn_emb + self.cnn_block1(self.cnn_block1(epitope_cnn_emb))
+        while conv.size(-1) >= 2:
+            conv = self.cnn_block2(conv)
+        epitope_cnn_out = torch.squeeze(conv, dim=-1)
+        epitope_cnn_out = self.bn1(epitope_cnn_out)
+
+        hla_cnn_emb = self.region_cnn2(hla_emb.transpose(1, 2))
+        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)
+        hla_cnn_out = self.bn2(hla_cnn_out)
+        # CNN Blocks
+
+        # Concatenate and pass through MLP
+        representation = torch.cat((epitope_mean, hla_mean, epitope_cnn_out, hla_cnn_out), dim=1)
+        reduction_feature = self.fc_task(representation)
+        logits_clsf = self.classifier(reduction_feature)
+        logits_clsf = torch.nn.functional.softmax(logits_clsf, dim=1)
+        return logits_clsf, representation
+
+