I was following this well-known paper and wanted to generate the attention matrices demonstrated on the last three pages of this paper. The data and model I built are based on this extensive tutorial. However, I am having difficulty extracting one transformer block and generating an attention matrix between the words and the class tokens for some of the correctly and incorrectly predicted reviews. Here is the part I believe needs significant modifications (even tried to get rid of a pre-trained model and manually coded one, but failed).
class BertBinaryClassifier(nn.Module):
def __init__(self, dropout=0.1):
super(BertBinaryClassifier, self).__init__()
self.bert = BertModel.from_pretrained('bert-base-uncased')
self.dropout = nn.Dropout(dropout)
self.linear = nn.Linear(768, 1)
self.sigmoid = nn.Sigmoid()
def forward(self, tokens, masks=None):
_, pooled_output = self.bert(tokens, attention_mask=masks, output_all_encoded_layers=False)
dropout_output = self.dropout(pooled_output)
linear_output = self.linear(dropout_output)
proba = self.sigmoid(linear_output)
return proba
Could someone offer some guidance or insights on the implementation?
