| 51 | |
| 52 | |
| 53 | class AttentionBlock(nn.Module): |
| 54 | |
| 55 | def __init__(self, dim, num_heads, post_norm, dropout=0.1, eps=1e-5): |
| 56 | super().__init__() |
| 57 | self.dim = dim |
| 58 | self.num_heads = num_heads |
| 59 | self.post_norm = post_norm |
| 60 | self.eps = eps |
| 61 | |
| 62 | # layers |
| 63 | self.attn = SelfAttention(dim, num_heads, dropout, eps) |
| 64 | self.norm1 = nn.LayerNorm(dim, eps=eps) |
| 65 | self.ffn = nn.Sequential( |
| 66 | nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim), |
| 67 | nn.Dropout(dropout)) |
| 68 | self.norm2 = nn.LayerNorm(dim, eps=eps) |
| 69 | |
| 70 | def forward(self, x, mask): |
| 71 | if self.post_norm: |
| 72 | x = self.norm1(x + self.attn(x, mask)) |
| 73 | x = self.norm2(x + self.ffn(x)) |
| 74 | else: |
| 75 | x = x + self.attn(self.norm1(x), mask) |
| 76 | x = x + self.ffn(self.norm2(x)) |
| 77 | return x |
| 78 | |
| 79 | |
| 80 | class XLMRoberta(nn.Module): |