Update model.py

model.py

@@ -1,94 +1,88 @@
-from transformers import PreTrainedModel, AutoModel
-import torch
-import torch.nn as nn
-import math
-from .config import BERTMultiAttentionConfig
-
-class MultiHeadAttention(nn.Module):
-    def __init__(self, config):
-        super(MultiHeadAttention, self).__init__()
-        self.hidden_size = config.hidden_size
-        self.num_heads = config.num_heads
-        self.head_dim = config.hidden_size // config.num_heads
-
-        self.query = nn.Linear(config.hidden_size, config.hidden_size)
-        self.key = nn.Linear(config.hidden_size, config.hidden_size)
-        self.value = nn.Linear(config.hidden_size, config.hidden_size)
-        self.out = nn.Linear(config.hidden_size, config.hidden_size)
-
-        self.layer_norm_q = nn.LayerNorm(config.hidden_size)
-        self.layer_norm_k = nn.LayerNorm(config.hidden_size)
-        self.layer_norm_v = nn.LayerNorm(config.hidden_size)
-        self.layer_norm_out = nn.LayerNorm(config.hidden_size)
-
-        self.dropout = nn.Dropout(config.dropout)
-
-    def forward(self, query, key, value):
-        batch_size = query.size(0)
-
-        query = self.layer_norm_q(self.query(query))
-        key = self.layer_norm_k(self.key(key))
-        value = self.layer_norm_v(self.value(value))
-
-        query = query.view(batch_size, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
-        key = key.view(batch_size, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
-        value = value.view(batch_size, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
-
-        attention_scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.head_dim)
-        attention_weights = nn.Softmax(dim=-1)(attention_scores)
-        attention_weights = self.dropout(attention_weights)
-
-        attended_values = torch.matmul(attention_weights, value).permute(0, 2, 1, 3).contiguous()
-        attended_values = attended_values.view(batch_size, -1, self.hidden_size)
-
-        out = self.layer_norm_out(self.out(attended_values))
-        out = self.dropout(out)
-
-        return out
-
-class BERTMultiAttentionModel(PreTrainedModel):
-    config_class = BERTMultiAttentionConfig
-
-    def __init__(self, config):
-        super(BERTMultiAttentionModel, self).__init__(config)
-        self.config = config
-
-        # Initialize the transformer model
-        self.transformer = AutoModel.from_pretrained(config.transformer)
-        self.cross_attention = MultiHeadAttention(config)
-        self.fc1 = nn.Linear(config.hidden_size * 2, 256)
-        self.layer_norm_fc1 = nn.LayerNorm(256)
-        self.dropout1 = nn.Dropout(config.dropout)
-        self.rnn = nn.LSTM(input_size=256, hidden_size=config.rnn_hidden_size, num_layers=config.rnn_num_layers, batch_first=True, bidirectional=config.rnn_bidirectional, dropout=config.dropout)
-        self.layer_norm_rnn = nn.LayerNorm(256)
-        self.dropout2 = nn.Dropout(config.dropout)
-        self.fc_proj = nn.Linear(256, 256)
-        self.layer_norm_proj = nn.LayerNorm(256)
-        self.dropout3 = nn.Dropout(config.dropout)
-        self.fc_final = nn.Linear(256, 1)
-
-    def forward(self, input_ids1, attention_mask1, input_ids2, attention_mask2):
-        output1 = self.transformer(input_ids1, attention_mask=attention_mask1)[0]
-        output2 = self.transformer(input_ids2, attention_mask=attention_mask2)[0]
-
-        attended_output = self.cross_attention(output1, output2, output2)
-        combined_output = torch.cat([output1, attended_output], dim=2)
-        combined_output = torch.mean(combined_output, dim=1)
-
-        combined_output = self.layer_norm_fc1(self.fc1(combined_output))
-        combined_output = self.dropout1(torch.relu(combined_output))
-        combined_output = combined_output.unsqueeze(1)
-
-        _, (hidden_state, _) = self.rnn(combined_output)
-        hidden_state_concat = torch.cat([hidden_state[-2], hidden_state[-1]], dim=-1)
-
-        hidden_state_proj = self.layer_norm_proj(self.fc_proj(hidden_state_concat))
-        hidden_state_proj = self.dropout2(hidden_state_proj)
-
-        final = self.fc_final(hidden_state_proj)
-        final = self.dropout3(final)
-
-        return torch.sigmoid(final)
-
-
-AutoModel.register(BERTMultiAttentionConfig, BERTMultiAttentionModel)
+from transformers import PreTrainedModel, AutoModel
+import torch
+import torch.nn as nn
+import math
+from .config import BERTMultiAttentionConfig
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, config):
+        super(MultiHeadAttention, self).__init__()
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_heads
+        self.head_dim = config.hidden_size // config.num_heads
+
+        self.query = nn.Linear(config.hidden_size, config.hidden_size)
+        self.key = nn.Linear(config.hidden_size, config.hidden_size)
+        self.value = nn.Linear(config.hidden_size, config.hidden_size)
+        self.out = nn.Linear(config.hidden_size, config.hidden_size)
+
+        self.layer_norm_q = nn.LayerNorm(config.hidden_size)
+        self.layer_norm_k = nn.LayerNorm(config.hidden_size)
+        self.layer_norm_v = nn.LayerNorm(config.hidden_size)
+        self.layer_norm_out = nn.LayerNorm(config.hidden_size)
+
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, query, key, value):
+        batch_size = query.size(0)
+
+        query = self.layer_norm_q(self.query(query))
+        key = self.layer_norm_k(self.key(key))
+        value = self.layer_norm_v(self.value(value))
+
+        query = query.view(batch_size, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+        key = key.view(batch_size, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+        value = value.view(batch_size, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+
+        attention_scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        attention_weights = nn.Softmax(dim=-1)(attention_scores)
+        attention_weights = self.dropout(attention_weights)
+
+        attended_values = torch.matmul(attention_weights, value).permute(0, 2, 1, 3).contiguous()
+        attended_values = attended_values.view(batch_size, -1, self.hidden_size)
+
+        out = self.layer_norm_out(self.out(attended_values))
+        out = self.dropout(out)
+
+        return out
+
+class BERTMultiAttentionModel(PreTrainedModel):
+    config_class = BERTMultiAttentionConfig
+
+    def __init__(self, config):
+        super(BERTMultiAttentionModel, self).__init__(config)
+        self.config = config
+
+        # Initialize the transformer model
+        self.transformer = AutoModel.from_pretrained(config.transformer)
+        self.cross_attention = MultiHeadAttention(config)
+        self.fc1 = nn.Linear(config.hidden_size * 2, 256)
+        self.layer_norm_fc1 = nn.LayerNorm(256)
+        self.dropout1 = nn.Dropout(config.dropout)
+        self.rnn = nn.LSTM(input_size=256, hidden_size=config.rnn_hidden_size, num_layers=config.rnn_num_layers, batch_first=True, bidirectional=config.rnn_bidirectional, dropout=config.dropout)
+        self.layer_norm_rnn = nn.LayerNorm(256)
+        self.dropout2 = nn.Dropout(config.dropout)
+        self.fc_proj = nn.Linear(256, 256)
+        self.layer_norm_proj = nn.LayerNorm(256)
+        self.dropout3 = nn.Dropout(config.dropout)
+        self.fc_final = nn.Linear(256, 1)
+
+    def forward(self, input_ids1, attention_mask1, input_ids2, attention_mask2):
+        output1 = self.transformer(input_ids1, attention_mask1)[0]
+        output2 = self.transformer(input_ids2, attention_mask2)[0]
+        attended_output = self.cross_attention(output1, output2, output2)
+        combined_output = torch.cat([output1, attended_output], dim=2)
+        combined_output = torch.mean(combined_output, dim=1)
+        combined_output = self.layer_norm_fc1(self.fc1(combined_output))
+        combined_output = self.dropout1(torch.relu(combined_output))
+        combined_output = combined_output.unsqueeze(1)
+        _, (hidden_state, _) = self.rnn(combined_output)
+        hidden_state_concat = torch.cat([hidden_state[0], hidden_state[1]], dim=-1)
+        hidden_state_proj = self.layer_norm_proj(self.fc_proj(hidden_state_concat))
+        hidden_state_proj = self.dropout2(hidden_state_proj)
+        final = self.fc_final(hidden_state_proj)
+        final = self.dropout3(final)
+        return torch.sigmoid(final)
+
+
+AutoModel.register(BERTMultiAttentionConfig, BERTMultiAttentionModel)