A transformer block, including layer norm, self-attention, another layer norm, and a two-layer MLP. :param embedding_dim: total dimensionality of the self-attention model (equal to `head_size * num_heads`) :param num_heads: number of self-attention heads :param bias: whether
| 208 | |
| 209 | |
| 210 | class TransformerBlock(nn.Module): |
| 211 | """A transformer block, including layer norm, self-attention, another layer norm, |
| 212 | and a two-layer MLP. |
| 213 | |
| 214 | :param embedding_dim: total dimensionality of the self-attention model (equal to |
| 215 | `head_size * num_heads`) |
| 216 | :param num_heads: number of self-attention heads |
| 217 | :param bias: whether to include bias terms; used for layernorms, attention, and MLP |
| 218 | :param dropout: amount of dropout; used for attention, resiudal pathway, and MLP |
| 219 | :param causal: if true, use causal self-attention |
| 220 | :param mlp_expansion: ratio between embedding dimension and side of MLP hidden layer |
| 221 | """ |
| 222 | |
| 223 | def __init__( |
| 224 | self, |
| 225 | embedding_dim: int, |
| 226 | num_heads: int, |
| 227 | causal: bool, |
| 228 | dropout: float, |
| 229 | bias: bool = True, |
| 230 | mlp_expansion: int = 4, |
| 231 | ): |
| 232 | super().__init__() |
| 233 | |
| 234 | self.layernorm1 = LayerNorm(embedding_dim, bias=bias) |
| 235 | self.attention = SelfAttention( |
| 236 | embedding_dim, num_heads, bias=bias, dropout=dropout, causal=causal |
| 237 | ) |
| 238 | self.layernorm2 = LayerNorm(embedding_dim, bias=bias) |
| 239 | |
| 240 | hidden_dim = mlp_expansion * embedding_dim |
| 241 | self.mlp = MLP(embedding_dim, hidden_dim, nn.GELU(), dropout=dropout, bias=bias) |
| 242 | |
| 243 | def forward(self, x): |
| 244 | x = x + self.attention(self.layernorm1(x)) |
| 245 | x = x + self.mlp(self.layernorm2(x)) |
| 246 | return x |
| 247 | |
| 248 | |
| 249 | class LayerNorm(nn.Module): |