Update model_structure.py

model_structure.py

@@ -1,30 +1,6 @@
-import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
-
-
-class AttentionFusion(nn.Module):
-    def __init__(self, feature_dims, hidden_dim):
-        super(AttentionFusion, self).__init__()
-        total_dim = sum(feature_dims)
-        self.attention_net = nn.Sequential(
-            nn.Linear(total_dim, hidden_dim),
-            nn.ReLU(),
-            nn.Linear(hidden_dim, len(feature_dims)),
-            nn.Softmax(dim=1)
-        )
-        self.feature_dims = feature_dims
-
-    def forward(self, feature_list):
-        concatenated_features = torch.cat(feature_list, dim=1)
-        attention_weights = self.attention_net(concatenated_features)
-        fused_feature = 0
-        for i, feature in enumerate(feature_list):
-            fused_feature += attention_weights[:, i].unsqueeze(1) * feature
-        return fused_feature
-
 
 class SelfAttention(nn.Module):
     def __init__(self, feature_dim):
@@ -34,33 +10,14 @@ class SelfAttention(nn.Module):
         self.value = nn.Linear(feature_dim, feature_dim)
         self.softmax = nn.Softmax(dim=-1)
 
-    def forward(self, x, mask=None):
-        """
-        x: (batch, seq_len, dim)
-        mask: (batch, seq_len) with 1 for valid tokens and 0 for padding
-        """
+    def forward(self, x): 
         q = self.query(x)
         k = self.key(x)
         v = self.value(x)
-
-        scores = torch.bmm(q, k.transpose(1, 2)) / math.sqrt(k.size(-1))
-
-        if mask is not None:
-            key_mask = mask.unsqueeze(1).expand(-1, scores.size(1), -1)
-            scores = scores.masked_fill(key_mask == 0, -1e9)
-
-        attn = self.softmax(scores)
-        out = torch.bmm(attn, v)  # (batch, seq_len, dim)
-
-        if mask is None:
-            return torch.mean(out, dim=1)
-
-        query_mask = mask.unsqueeze(-1).type_as(out)
-        out = out * query_mask
-        denom = query_mask.sum(dim=1).clamp(min=1.0)
-        pooled = out.sum(dim=1) / denom
-        return pooled
-
+        attention_scores = torch.bmm(q, k.transpose(1, 2))
+        attention_weights = self.softmax(attention_scores / (k.size(-1) ** 0.5))
+        weighted_features = torch.bmm(attention_weights, v)
+        return torch.mean(weighted_features, dim=1)
 
 class ParallelFeatureExtractorWithAttention(nn.Module):
     def __init__(self, input_dim, cnn_out_channels, lstm_hidden_dim, dropout_rate=0.3):
@@ -80,36 +37,22 @@ class ParallelFeatureExtractorWithAttention(nn.Module):
         self.bilstm_attention = SelfAttention(lstm_hidden_dim * 2)
         self.bilstm_branch_output_dim = lstm_hidden_dim * 2
 
-    def forward(self, sequence_embedding, mask=None):
-        # CNN branch
-        cnn_in = sequence_embedding.permute(0, 2, 1)  # (batch, dim, seq_len)
-        cnn_out = F.relu(self.cnn(cnn_in))
-        cnn_out_permuted = cnn_out.permute(0, 2, 1)  # (batch, seq_len, channels)
-        v_cnn = self.cnn_attention(cnn_out_permuted, mask=mask)
-
-        # BiLSTM branch (packed to ignore padding)
-        if mask is not None:
-            lengths = mask.sum(dim=1).to(torch.long).cpu()
-            packed = pack_padded_sequence(sequence_embedding, lengths, batch_first=True, enforce_sorted=False)
-            packed_out, _ = self.bilstm(packed)
-            lstm_out, _ = pad_packed_sequence(
-                packed_out, batch_first=True, total_length=sequence_embedding.size(1)
-            )
-        else:
-            lstm_out, _ = self.bilstm(sequence_embedding)
-
-        v_bilstm = self.bilstm_attention(lstm_out, mask=mask)
+    def forward(self, sequence_embedding):
+        cnn_in = sequence_embedding.permute(0, 2, 1) 
+        cnn_out = self.cnn(cnn_in)
+        cnn_out = F.relu(cnn_out)
+        cnn_out_permuted = cnn_out.permute(0, 2, 1) 
+        v_cnn = self.cnn_attention(cnn_out_permuted)
+        lstm_out, _ = self.bilstm(sequence_embedding) 
+        v_bilstm = self.bilstm_attention(lstm_out)
         return v_cnn, v_bilstm
 
-
 class AVP_Fusion(nn.Module):
     def __init__(self, esm_dim, additional_dim, cnn_out_channels, lstm_hidden_dim, num_classes, dropout_rate=0.42):
         super(AVP_Fusion, self).__init__()
         fused_input_dim = esm_dim + additional_dim
-        self.parallel_extractor = ParallelFeatureExtractorWithAttention(
-            fused_input_dim, cnn_out_channels, lstm_hidden_dim, dropout_rate
-        )
-
+        self.parallel_extractor = ParallelFeatureExtractorWithAttention(fused_input_dim, cnn_out_channels, lstm_hidden_dim, dropout_rate)
+        
         cnn_feature_dim = self.parallel_extractor.cnn_branch_output_dim
         bilstm_feature_dim = self.parallel_extractor.bilstm_branch_output_dim
 
@@ -129,21 +72,14 @@ class AVP_Fusion(nn.Module):
         )
         self.embedding_dim = classifier_input_dim
 
-    def forward(self, esm_sequence_embedding, additional_features, attention_mask=None):
+    def forward(self, esm_sequence_embedding, additional_features):
         seq_len = esm_sequence_embedding.size(1)
-
         expanded_additional_features = additional_features.unsqueeze(1).expand(-1, seq_len, -1)
         fused_sequence_embedding = torch.cat([esm_sequence_embedding, expanded_additional_features], dim=2)
-
-        v_cnn, v_bilstm = self.parallel_extractor(fused_sequence_embedding, mask=attention_mask)
-
+        
+        v_cnn, v_bilstm = self.parallel_extractor(fused_sequence_embedding)
         v_cnn_matched = self.cnn_dim_matcher(v_cnn)
         lambda_gate = self.gating_network(torch.cat([v_cnn, v_bilstm], dim=1))
         final_embedding = lambda_gate * v_cnn_matched + (1 - lambda_gate) * v_bilstm
-
         logits = self.classifier(final_embedding)
         return logits, final_embedding
-
-
-# Backward-compatible alias (do not remove)
-AVP_HNCL_v3 = AVP_Fusion