Spaces:
Sleeping
Sleeping
| """Encoder stack used by the transformer architecture.""" | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| import torch.utils.checkpoint as ckpt | |
| from torch import Tensor | |
| from .attention import MultiHeadAttention | |
| from .feedforward import FeedForwardLayer | |
| __all__ = ["EncoderLayer", "TransformerEncoder"] | |
| class EncoderLayer(nn.Module): | |
| """Self-attention + feed-forward block supporting pre/post layer normalisation.""" | |
| def __init__( | |
| self, | |
| d_model: int, | |
| num_heads: int, | |
| d_ff: int, | |
| dropout_rate: float, | |
| *, | |
| layer_norm_style: str = "post", | |
| ) -> None: | |
| super().__init__() | |
| if not isinstance(d_model, int): | |
| raise TypeError(f"d_model must be an int, got {type(d_model)}") | |
| if not isinstance(num_heads, int): | |
| raise TypeError(f"num_heads must be an int, got {type(num_heads)}") | |
| if not isinstance(d_ff, int): | |
| raise TypeError(f"d_ff must be an int, got {type(d_ff)}") | |
| if not isinstance(dropout_rate, float): | |
| raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}") | |
| if not isinstance(layer_norm_style, str): | |
| raise TypeError(f"layer_norm_style must be a string, got {type(layer_norm_style)}") | |
| if d_model <= 0: | |
| raise ValueError("d_model must be strictly greater than 0") | |
| if num_heads <= 0: | |
| raise ValueError("num_heads must be strictly greater than 0") | |
| if d_ff <= 0: | |
| raise ValueError("d_ff must be strictly greater than 0") | |
| if not 0.0 <= dropout_rate < 1.0: | |
| raise ValueError("dropout_rate must be in [0, 1)") | |
| style = layer_norm_style.lower() | |
| if style not in {"pre", "post"}: | |
| raise ValueError("layer_norm_style must be either 'pre' or 'post' (case-insensitive)") | |
| self.layer_norm_style = style | |
| self.pre_norm = style == "pre" | |
| self.attention_layer = MultiHeadAttention(d_model, num_heads, dropout_rate) | |
| self.feed_forward = FeedForwardLayer(d_model, d_ff, dropout_rate) | |
| self.norm1 = nn.LayerNorm(d_model) | |
| self.dropout1 = nn.Dropout(dropout_rate) | |
| self.norm2 = nn.LayerNorm(d_model) | |
| self.dropout2 = nn.Dropout(dropout_rate) | |
| def forward(self, x: Tensor, src_padding_mask: Tensor) -> Tensor: | |
| if not isinstance(x, torch.Tensor): | |
| raise TypeError("x must be a torch.Tensor") | |
| if x.dim() != 3: | |
| raise ValueError( | |
| f"x must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(x.shape)}" | |
| ) | |
| if not isinstance(src_padding_mask, torch.Tensor): | |
| raise TypeError("src_padding_mask must be a torch.Tensor") | |
| if src_padding_mask.dtype != torch.bool or src_padding_mask.dim() != 4: | |
| raise TypeError( | |
| f"src_padding_mask must be boolean with shape (B, H, 1, S);" | |
| f" got dtype {src_padding_mask.dtype} and shape {tuple(src_padding_mask.shape)}" | |
| ) | |
| if self.pre_norm: | |
| normed = self.norm1(x) | |
| attn_out = self.attention_layer( | |
| normed, normed, normed, src_padding_mask, src_padding_mask | |
| ) | |
| x = x + self.dropout1(attn_out) | |
| ff_out = self.feed_forward(self.norm2(x)) | |
| x = x + self.dropout2(ff_out) | |
| return x | |
| attn_out = self.attention_layer(x, x, x, src_padding_mask, src_padding_mask) | |
| x = self.norm1(x + self.dropout1(attn_out)) | |
| ff_out = self.feed_forward(x) | |
| return self.norm2(x + self.dropout2(ff_out)) | |
| class TransformerEncoder(nn.Module): | |
| """Stack of encoder layers with optional activation checkpointing.""" | |
| def __init__( | |
| self, | |
| d_model: int, | |
| num_heads: int, | |
| d_ff: int, | |
| num_layers: int, | |
| dropout_rate: float, | |
| *, | |
| layer_norm_style: str = "post", | |
| ) -> None: | |
| super().__init__() | |
| if not isinstance(num_layers, int): | |
| raise TypeError(f"num_layers must be an int, got {type(num_layers)}") | |
| if num_layers <= 0: | |
| raise ValueError("num_layers must be strictly greater than 0") | |
| if not isinstance(layer_norm_style, str): | |
| raise TypeError(f"layer_norm_style must be a string, got {type(layer_norm_style)}") | |
| style = layer_norm_style.lower() | |
| if style not in {"pre", "post"}: | |
| raise ValueError("layer_norm_style must be either 'pre' or 'post' (case-insensitive)") | |
| self.layer_norm_style = style | |
| self.layers = nn.ModuleList( | |
| [ | |
| EncoderLayer( | |
| d_model, | |
| num_heads, | |
| d_ff, | |
| dropout_rate, | |
| layer_norm_style=style, | |
| ) | |
| for _ in range(num_layers) | |
| ] | |
| ) | |
| self.use_ckpt = False | |
| def forward(self, x: Tensor, src_padding_mask: Tensor) -> Tensor: | |
| if not isinstance(x, torch.Tensor): | |
| raise TypeError("x must be a torch.Tensor") | |
| if x.dim() != 3: | |
| raise ValueError( | |
| f"x must be a 3D tensor of shape (B, S, D); got shape {tuple(x.shape)}" | |
| ) | |
| for layer in self.layers: | |
| if self.use_ckpt: | |
| def _fn(x_, *, _layer=layer): | |
| return _layer(x_, src_padding_mask) | |
| x = ckpt.checkpoint(_fn, x, use_reentrant=False) | |
| else: | |
| x = layer(x, src_padding_mask) | |
| return x | |