Spaces:
Sleeping
Sleeping
| """Decoder stack used by the encoder-decoder transformer.""" | |
| 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__ = ["DecoderLayer", "TransformerDecoder"] | |
| class DecoderLayer(nn.Module): | |
| """Single decoder block with self-attention, cross-attention, and feed-forward.""" | |
| 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.self_attention_layer = MultiHeadAttention(d_model, num_heads, dropout_rate) | |
| self.cross_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) | |
| self.norm3 = nn.LayerNorm(d_model) | |
| self.dropout3 = nn.Dropout(dropout_rate) | |
| def forward( | |
| self, | |
| x: Tensor, | |
| y: Tensor, | |
| src_padding_mask: Tensor, | |
| tgt_padding_mask: Tensor, | |
| tgt_causal_mask: Tensor | None, | |
| ) -> Tensor: | |
| """Run one decoder layer using the configured layer-normalisation style.""" | |
| 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(y, torch.Tensor): | |
| raise TypeError("y must be a torch.Tensor") | |
| if y.dim() != 3: | |
| raise ValueError( | |
| f"y must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(y.shape)}" | |
| ) | |
| if x.shape[0] != y.shape[0] or x.shape[-1] != y.shape[-1]: | |
| raise ValueError("Encoder memory and decoder input must match in batch and d_model") | |
| for mask_name, mask in ( | |
| ("src_padding_mask", src_padding_mask), | |
| ("tgt_padding_mask", tgt_padding_mask), | |
| ): | |
| if not isinstance(mask, torch.Tensor): | |
| raise TypeError(f"{mask_name} must be a torch.Tensor") | |
| if mask.dtype != torch.bool or mask.dim() != 4: | |
| raise TypeError( | |
| f"{mask_name} must be boolean with shape (B, H, 1, S);" | |
| f" got dtype {mask.dtype} and shape {tuple(mask.shape)}" | |
| ) | |
| if self.pre_norm: | |
| normed = self.norm1(y) | |
| self_attn_out = self.self_attention_layer( | |
| normed, | |
| normed, | |
| normed, | |
| tgt_padding_mask, | |
| tgt_padding_mask, | |
| tgt_causal_mask, | |
| ) | |
| y = y + self.dropout1(self_attn_out) | |
| normed_y = self.norm2(y) | |
| cross_attn_out = self.cross_attention_layer( | |
| normed_y, | |
| x, | |
| x, | |
| tgt_padding_mask, | |
| src_padding_mask, | |
| ) | |
| y = y + self.dropout2(cross_attn_out) | |
| ff_out = self.feed_forward(self.norm3(y)) | |
| y = y + self.dropout3(ff_out) | |
| return y | |
| # Self-attention (post-LN) | |
| self_attn_out = self.self_attention_layer( | |
| y, | |
| y, | |
| y, | |
| tgt_padding_mask, | |
| tgt_padding_mask, | |
| tgt_causal_mask, | |
| ) | |
| y = self.norm1(y + self.dropout1(self_attn_out)) | |
| # Cross-attention (encoder memory as keys/values) | |
| cross_attn_out = self.cross_attention_layer( | |
| y, | |
| x, | |
| x, | |
| tgt_padding_mask, | |
| src_padding_mask, | |
| ) | |
| y = self.norm2(y + self.dropout2(cross_attn_out)) | |
| # Feed-forward block | |
| ff_out = self.feed_forward(y) | |
| return self.norm3(y + self.dropout3(ff_out)) | |
| class TransformerDecoder(nn.Module): | |
| """Stack of decoder 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( | |
| [ | |
| DecoderLayer( | |
| 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, | |
| y: Tensor, | |
| src_padding_mask: Tensor, | |
| tgt_padding_mask: Tensor, | |
| tgt_causal_mask: Tensor | None, | |
| ) -> Tensor: | |
| """Run the decoder stack for all time steps.""" | |
| 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(y, torch.Tensor): | |
| raise TypeError("y must be a torch.Tensor") | |
| if y.dim() != 3: | |
| raise ValueError( | |
| f"y must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(y.shape)}" | |
| ) | |
| if x.shape[0] != y.shape[0] or x.shape[-1] != y.shape[-1]: | |
| raise ValueError("Encoder memory and decoder input must match in batch and d_model") | |
| for layer in self.layers: | |
| if self.use_ckpt: | |
| def _fn(y_, *, _layer=layer): | |
| return _layer(x, y_, src_padding_mask, tgt_padding_mask, tgt_causal_mask) | |
| y = ckpt.checkpoint(_fn, y, use_reentrant=False) | |
| else: | |
| y = layer(x, y, src_padding_mask, tgt_padding_mask, tgt_causal_mask) | |
| return y | |