1"""Pre-norm Transformer encoder (nn.TransformerEncoderLayer, norm_first=True) with key padding masks."""2from __future__ import annotations34import torch5import torch.nn as nn678class Encoder(nn.Module):9 def __init__(self, d_model: int, n_layers: int, n_heads: int, d_ff: int, dropout: float = 0.1):10 super().__init__()11 layer = nn.TransformerEncoderLayer(d_model, n_heads, d_ff, dropout, activation="gelu", batch_first=True,12 norm_first=True)13 self.layers = nn.TransformerEncoder(layer, n_layers, enable_nested_tensor=False)14 self.norm = nn.LayerNorm(d_model)1516 def forward(self, tokens: torch.Tensor, key_padding_mask: torch.Tensor | None = None) -> torch.Tensor:17 # key_padding_mask: True = ignore. Guard against all-masked rows (would produce NaN attention).18 if key_padding_mask is not None:19 all_masked = key_padding_mask.all(dim=1)20 if all_masked.any():21 key_padding_mask = key_padding_mask.clone()22 key_padding_mask[all_masked, 0] = False23 return self.norm(self.layers(tokens, src_key_padding_mask=key_padding_mask))24