File size: 4,236 Bytes
df13d8f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
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
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[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)
|