| 186 | |
| 187 | |
| 188 | class BertSelfAttention(nn.Module): |
| 189 | def __init__(self, config): |
| 190 | super().__init__() |
| 191 | if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"): |
| 192 | raise ValueError( |
| 193 | "The hidden size (%d) is not a multiple of the number of attention " |
| 194 | "heads (%d)" % (config.hidden_size, config.num_attention_heads) |
| 195 | ) |
| 196 | |
| 197 | self.num_attention_heads = config.num_attention_heads |
| 198 | self.attention_head_size = int(config.hidden_size / config.num_attention_heads) |
| 199 | self.all_head_size = self.num_attention_heads * self.attention_head_size |
| 200 | |
| 201 | self.query = nn.Linear(config.hidden_size, self.all_head_size) |
| 202 | self.key = nn.Linear(config.hidden_size, self.all_head_size) |
| 203 | self.value = nn.Linear(config.hidden_size, self.all_head_size) |
| 204 | |
| 205 | self.dropout = nn.Dropout(config.attention_probs_dropout_prob) |
| 206 | |
| 207 | def transpose_for_scores(self, x): |
| 208 | new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size) |
| 209 | x = x.view(*new_x_shape) |
| 210 | return x.permute(0, 2, 1, 3) |
| 211 | |
| 212 | def forward( |
| 213 | self, |
| 214 | hidden_states, |
| 215 | attention_mask=None, |
| 216 | head_mask=None, |
| 217 | encoder_hidden_states=None, |
| 218 | encoder_attention_mask=None, |
| 219 | output_attentions=False, |
| 220 | ): |
| 221 | mixed_query_layer = self.query(hidden_states) |
| 222 | |
| 223 | # If this is instantiated as a cross-attention module, the keys |
| 224 | # and values come from an encoder; the attention mask needs to be |
| 225 | # such that the encoder's padding tokens are not attended to. |
| 226 | if encoder_hidden_states is not None: |
| 227 | mixed_key_layer = self.key(encoder_hidden_states) |
| 228 | mixed_value_layer = self.value(encoder_hidden_states) |
| 229 | attention_mask = encoder_attention_mask |
| 230 | else: |
| 231 | mixed_key_layer = self.key(hidden_states) |
| 232 | mixed_value_layer = self.value(hidden_states) |
| 233 | |
| 234 | query_layer = self.transpose_for_scores(mixed_query_layer) |
| 235 | key_layer = self.transpose_for_scores(mixed_key_layer) |
| 236 | value_layer = self.transpose_for_scores(mixed_value_layer) |
| 237 | |
| 238 | # Take the dot product between "query" and "key" to get the raw attention scores. |
| 239 | attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) |
| 240 | attention_scores = attention_scores / math.sqrt(self.attention_head_size) |
| 241 | if attention_mask is not None: |
| 242 | # Apply the attention mask is (precomputed for all layers in BertModel forward() function) |
| 243 | attention_scores = attention_scores + attention_mask |
| 244 | |
| 245 | # Normalize the attention scores to probabilities. |