| 236 | |
| 237 | |
| 238 | class Block(nn.Module): |
| 239 | def __init__(self, head_size, mlp_hidden_size, num_heads, sequence_length): |
| 240 | """A decoder layer. |
| 241 | |
| 242 | This module combines a Multi-headed Attention module and an MLP to |
| 243 | create a layer of the transformer, with normalization and skip-connections. |
| 244 | See Lecture 06, slide 33. |
| 245 | """ |
| 246 | super(Block, self).__init__() |
| 247 | self.head_size = head_size |
| 248 | self.mlp_hidden_size = mlp_hidden_size |
| 249 | self.num_heads = num_heads |
| 250 | self.sequence_length = sequence_length |
| 251 | self.hidden_size = num_heads * head_size |
| 252 | |
| 253 | self.attention = MultiHeadedAttention(head_size, num_heads, sequence_length) |
| 254 | self.norm1 = nn.LayerNorm(self.hidden_size) |
| 255 | self.mlp = nn.Sequential( |
| 256 | nn.Linear(self.hidden_size, mlp_hidden_size), |
| 257 | nn.GELU(), |
| 258 | nn.Linear(mlp_hidden_size, self.hidden_size), |
| 259 | ) |
| 260 | self.norm2 = nn.LayerNorm(self.hidden_size) |
| 261 | |
| 262 | def forward(self, hidden_states): |
| 263 | attention_outputs = self.attention(hidden_states) |
| 264 | attention_outputs = self.norm1(attention_outputs + hidden_states) |
| 265 | outputs = self.mlp(attention_outputs) |
| 266 | outputs = self.norm2(outputs + attention_outputs) |
| 267 | return outputs |
| 268 | |
| 269 | |
| 270 | class MiniGPT1(nn.Module): |