| 53 | |
| 54 | |
| 55 | class BertLayer(nn.Module): |
| 56 | def __init__(self, hidden_size: int, num_heads: int, intermediate_size: int, dropout: float = 0.0): |
| 57 | super().__init__() |
| 58 | self.attention = BertSelfAttention(hidden_size, num_heads, dropout) |
| 59 | self.norm1 = nn.LayerNorm(hidden_size) |
| 60 | self.mlp = BertMLP(hidden_size, intermediate_size, dropout) |
| 61 | self.norm2 = nn.LayerNorm(hidden_size) |
| 62 | |
| 63 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 64 | x = self.norm1(x + self.attention(x)) |
| 65 | x = self.norm2(x + self.mlp(x)) |
| 66 | return x |
| 67 | |
| 68 | |
| 69 | class BertModel(nn.Module): |