Spaces:
Running
Running
| # coding: utf-8 | |
| """ | |
| Fixed and modernized transformer/attention layers. | |
| Key fixes and improvements: | |
| - Robust mask broadcasting for both causal/attention masks and padding masks. | |
| - Apply both key padding masks and optional attention masks to scores BEFORE softmax | |
| (using a large negative value instead of -inf to avoid NaNs). | |
| - Consistent tensor shapes and comments. | |
| - Corrected positional encoding shapes and sinusoidal implementation. | |
| - Defensive programming with clear error messages when masks cannot be broadcast. | |
| """ | |
| import math | |
| from typing import Optional | |
| import torch | |
| import torch.nn as nn | |
| from torch import Tensor | |
| from constants import TARGET_PAD | |
| _LARGE_NEG = -1e9 # used instead of -inf to avoid NaNs when entire row is masked | |
| def _prepare_mask_for_scores(mask: Tensor, scores: Tensor) -> Tensor: | |
| """ | |
| Turn a mask (which may be 2D: [B, Lk] or [B, Lq] or 3D [B, Lq, Lk]) into a boolean | |
| tensor that can be broadcasted against `scores` (shape [B, num_heads, Lq, Lk]). | |
| Returned mask has shape that can broadcast to scores; True means **keep**. | |
| """ | |
| if mask is None: | |
| return None | |
| # Ensure boolean | |
| mask = mask.bool() | |
| # scores shape: [B, H, Lq, Lk] | |
| _, _, Lq, Lk = scores.shape | |
| # If mask is 2D (B, Lk) -> key padding mask: expand to (B, 1, 1, Lk) | |
| if mask.dim() == 2 and mask.size(1) == Lk: | |
| return mask.unsqueeze(1).unsqueeze(2) | |
| # If mask is 2D (B, Lq) -> query mask: expand to (B, 1, Lq, 1) | |
| if mask.dim() == 2 and mask.size(1) == Lq: | |
| return mask.unsqueeze(1).unsqueeze(-1) | |
| # If mask is 3D and matches (B, Lq, Lk) | |
| if mask.dim() == 3 and mask.size(1) == Lq and mask.size(2) == Lk: | |
| return mask.unsqueeze(1) # -> (B,1,Lq,Lk) | |
| # If mask is 4D and already matches scores shape (B,H,Lq,Lk) return as-is | |
| if mask.dim() == 4 and mask.shape == scores.shape: | |
| return mask | |
| # try to be flexible: add singleton dims on the left until it can broadcast | |
| while mask.dim() < 4: | |
| mask = mask.unsqueeze(1) | |
| # now check if last two dims can broadcast with scores' last two dims | |
| if not (mask.size(-2) in (1, Lq) and mask.size(-1) in (1, Lk)): | |
| raise RuntimeError(f"Cannot broadcast mask of shape {mask.shape} to attention scores shape {scores.shape}") | |
| return mask | |
| class MultiHeadedAttention(nn.Module): | |
| def __init__(self, num_heads: int, size: int, dropout: float = 0.1): | |
| super(MultiHeadedAttention, self).__init__() | |
| assert size % num_heads == 0, "`size` must be divisible by `num_heads`" | |
| self.head_size = size // num_heads | |
| self.model_size = size | |
| self.num_heads = num_heads | |
| self.k_layer = nn.Linear(size, num_heads * self.head_size) | |
| self.v_layer = nn.Linear(size, num_heads * self.head_size) | |
| self.q_layer = nn.Linear(size, num_heads * self.head_size) | |
| self.output_layer = nn.Linear(size, size) | |
| self.softmax = nn.Softmax(dim=-1) | |
| self.dropout = nn.Dropout(dropout) | |
| self.target_pad = TARGET_PAD | |
| def forward(self, | |
| k: Tensor, | |
| v: Tensor, | |
| q: Tensor, | |
| mask: Optional[Tensor] = None, | |
| padding_mask: Optional[Tensor] = None) -> Tensor: | |
| """ | |
| k: [B, Lk, D] | |
| v: [B, Lk, D] | |
| q: [B, Lq, D] | |
| mask: optional attention mask (causal or custom). Typical shapes: | |
| - [B, Lq, Lk] OR | |
| - [B, Lq] OR | |
| - [Lq, Lk] (will be broadcasted across batch) | |
| padding_mask: optional key padding mask with shape [B, Lk] (True where token is VALID) | |
| Returns: context vectors of shape [B, Lq, D] | |
| """ | |
| batch_size = q.size(0) | |
| # project to multi-head space | |
| k_proj = self.k_layer(k) # [B, Lk, H*hs] | |
| v_proj = self.v_layer(v) | |
| q_proj = self.q_layer(q) | |
| # reshape -> [B, H, L, hs] | |
| def _reshape(x): | |
| B, L, _ = x.size() | |
| return x.view(B, L, self.num_heads, self.head_size).transpose(1, 2) | |
| k_heads = _reshape(k_proj) | |
| v_heads = _reshape(v_proj) | |
| q_heads = _reshape(q_proj) | |
| # scale queries | |
| q_heads = q_heads / math.sqrt(self.head_size) | |
| # compute attention scores [B, H, Lq, Lk] | |
| scores = torch.matmul(q_heads, k_heads.transpose(2, 3)) | |
| # Prepare masks for scores and apply before softmax to zero out on softmax | |
| # Convert both mask (causal/custom) and padding_mask (keys) into broadcastable masks | |
| combined_mask = None | |
| if padding_mask is not None: | |
| # padding_mask: True where token is VALID -> we want keep True | |
| key_mask = _prepare_mask_for_scores(padding_mask, scores) | |
| combined_mask = key_mask if combined_mask is None else (combined_mask & key_mask) | |
| if mask is not None: | |
| att_mask = _prepare_mask_for_scores(mask, scores) | |
| combined_mask = att_mask if combined_mask is None else (combined_mask & att_mask) | |
| if combined_mask is not None: | |
| # combined_mask True means KEEP, so invert for masked_fill | |
| scores = scores.masked_fill(~combined_mask, _LARGE_NEG) | |
| # softmax | |
| attention = self.softmax(scores) | |
| attention = self.dropout(attention) | |
| # guard against any NaNs that could arise when rows are all -LARGE_NEG | |
| if torch.isnan(attention).any(): | |
| attention = torch.nan_to_num(attention, nan=0.0, posinf=0.0, neginf=0.0) | |
| # compute context [B, H, Lq, hs] | |
| context = torch.matmul(attention, v_heads) | |
| # reshape back to [B, Lq, H*hs] | |
| context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.head_size) | |
| output = self.output_layer(context) | |
| return output | |
| class PositionwiseFeedForward(nn.Module): | |
| def __init__(self, input_size, ff_size, dropout=0.1): | |
| super(PositionwiseFeedForward, self).__init__() | |
| self.layer_norm = nn.LayerNorm(input_size, eps=1e-6) | |
| self.pwff_layer = nn.Sequential( | |
| nn.Linear(input_size, ff_size), | |
| nn.ReLU(), | |
| nn.Dropout(dropout), | |
| nn.Linear(ff_size, input_size), | |
| nn.Dropout(dropout), | |
| ) | |
| def forward(self, x: Tensor) -> Tensor: | |
| x_norm = self.layer_norm(x) | |
| return self.pwff_layer(x_norm) + x | |
| class SinusoidalPositionEmbeddings(nn.Module): | |
| """Sinusoidal timestep embeddings for diffusion/time embeddings.""" | |
| def __init__(self, dim: int): | |
| super().__init__() | |
| self.dim = dim | |
| def forward(self, time: Tensor) -> Tensor: | |
| # time: [B] or [B, 1] | |
| device = time.device | |
| half_dim = self.dim // 2 | |
| emb = math.log(10000) / (half_dim - 1) | |
| exponents = torch.exp(torch.arange(half_dim, device=device, dtype=torch.float) * -emb) | |
| # time may be [B] -> make [B,1] | |
| time = time.float().unsqueeze(1) | |
| args = time * exponents.unsqueeze(0) # [B, half_dim] | |
| emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1) | |
| return emb | |
| class PositionalEncoding(nn.Module): | |
| def __init__(self, size: int = 0, max_len: int = 200000, mask_count: bool = False): | |
| super(PositionalEncoding, self).__init__() | |
| if size % 2 != 0: | |
| raise ValueError("Cannot use sin/cos positional encoding with odd dim") | |
| pe = torch.zeros(max_len, size) | |
| position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) | |
| div_term = torch.exp(torch.arange(0, size, 2, dtype=torch.float) * -(math.log(10000.0) / size)) | |
| pe[:, 0::2] = torch.sin(position * div_term) | |
| pe[:, 1::2] = torch.cos(position * div_term) | |
| pe = pe.unsqueeze(0) # [1, max_len, size] | |
| self.register_buffer('pe', pe) | |
| self.dim = size | |
| self.mask_count = mask_count | |
| def forward(self, emb: Tensor) -> Tensor: | |
| # emb: [B, L, D] | |
| L = emb.size(1) | |
| return emb + self.pe[:, :L, :] | |
| class TransformerEncoderLayer(nn.Module): | |
| def __init__(self, size: int = 0, ff_size: int = 0, num_heads: int = 0, dropout: float = 0.1): | |
| super(TransformerEncoderLayer, self).__init__() | |
| self.layer_norm = nn.LayerNorm(size, eps=1e-6) | |
| self.src_src_att = MultiHeadedAttention(num_heads, size, dropout=dropout) | |
| self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size) | |
| self.dropout = nn.Dropout(dropout) | |
| self.size = size | |
| def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor: | |
| x_norm = self.layer_norm(x) | |
| h = self.src_src_att(x_norm, x_norm, x_norm, mask=mask) | |
| h = self.dropout(h) + x | |
| o = self.feed_forward(h) | |
| return o | |
| class TransformerDecoderLayer(nn.Module): | |
| def __init__(self, size: int = 0, ff_size: int = 0, num_heads: int = 0, dropout: float = 0.1, decoder_trg_trg: bool = True): | |
| super(TransformerDecoderLayer, self).__init__() | |
| self.size = size | |
| self.trg_trg_att = MultiHeadedAttention(num_heads, size, dropout=dropout) | |
| self.src_trg_att = MultiHeadedAttention(num_heads, size, dropout=dropout) | |
| self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size) | |
| self.x_layer_norm = nn.LayerNorm(size, eps=1e-6) | |
| self.dec_layer_norm = nn.LayerNorm(size, eps=1e-6) | |
| self.dropout = nn.Dropout(dropout) | |
| self.decoder_trg_trg = decoder_trg_trg | |
| def forward(self, | |
| x: Tensor = None, | |
| memory: Tensor = None, | |
| src_mask: Optional[Tensor] = None, | |
| trg_mask: Optional[Tensor] = None, | |
| padding_mask: Optional[Tensor] = None) -> (Tensor, Tensor): | |
| # decoder/target self-attention | |
| h1 = self.x_layer_norm(x) | |
| # Target-Target Self Attention (causal + padding handled by masks) | |
| if self.decoder_trg_trg: | |
| h1 = self.trg_trg_att(h1, h1, h1, mask=trg_mask, padding_mask=padding_mask) | |
| h1 = self.dropout(h1) + x | |
| # Source-Target Attention: keys/values from memory (encoder output), queries from h1 | |
| h1_norm = self.dec_layer_norm(h1) | |
| h2 = self.src_trg_att(memory, memory, h1_norm, mask=src_mask, padding_mask=None) | |
| # final position-wise feed-forward layer | |
| o = self.feed_forward(self.dropout(h2) + h1) | |
| return o, h2 | |