Spaces:
Sleeping
Sleeping
| """Multi-head attention building block with explicit validation helpers.""" | |
| from __future__ import annotations | |
| from typing import cast | |
| import torch | |
| import torch.nn as nn | |
| from torch import Tensor | |
| from ..utils import ( | |
| calculate_attention, | |
| combine_masks, | |
| create_qk_padding_mask, | |
| join_heads, | |
| split_heads, | |
| ) | |
| __all__ = ["MultiHeadAttention"] | |
| class MultiHeadAttention(nn.Module): | |
| """ | |
| Multi-head attention: linear projections -> split heads -> scaled dot-product | |
| attention (via utils) -> merge heads -> output projection (+ dropout). | |
| Args: | |
| d_model (int): Model dimension (>0). Must be divisible by num_heads. | |
| num_heads (int): Number of attention heads (>0). | |
| dropout_rate (float): Dropout probability in (0,1). | |
| Inputs: | |
| query, key, value: (B, S, D) with D == d_model | |
| mask (optional): Boolean tensor broadcastable to (B, H, S_q, S_k) where True entries are masked. | |
| Returns: | |
| Tensor: (B, S_q, D) | |
| """ | |
| def __init__(self, d_model: int, num_heads: int, dropout_rate: float): | |
| super().__init__() | |
| # ---- type checks | |
| if not isinstance(num_heads, int): | |
| raise TypeError(f"num_heads must be an int, got {type(num_heads)}") | |
| if not isinstance(d_model, int): | |
| raise TypeError(f"d_model must be an int, got {type(d_model)}") | |
| if not isinstance(dropout_rate, float): | |
| raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}") | |
| # ---- value checks | |
| if num_heads <= 0: | |
| raise ValueError(f"num_heads must be strictly greater than 0, got {num_heads}") | |
| if d_model <= 0: | |
| raise ValueError(f"d_model must be strictly greater than 0, got {d_model}") | |
| if d_model % num_heads != 0: | |
| raise ValueError("d_model must be divisible by num_heads") | |
| if not (0 <= dropout_rate < 1): | |
| raise ValueError(f"dropout_rate must be between 0 and 1 excluded, got {dropout_rate}") | |
| self.d_model = d_model | |
| self.num_heads = num_heads | |
| self.d_head = d_model // num_heads | |
| self.dropout_rate = dropout_rate | |
| self.query_linear = nn.Linear(d_model, d_model, bias=False) | |
| self.key_linear = nn.Linear(d_model, d_model, bias=False) | |
| self.value_linear = nn.Linear(d_model, d_model, bias=False) | |
| self.output_linear = nn.Linear(d_model, d_model, bias=True) | |
| def forward( | |
| self, | |
| query: Tensor, | |
| key: Tensor, | |
| value: Tensor, | |
| q_mask: Tensor, | |
| k_mask: Tensor, | |
| causal_mask: Tensor | None = None, | |
| ) -> Tensor: | |
| """Run multi-head attention with boolean padding and optional causal masks.""" | |
| # ---- basic type checks | |
| if not isinstance(query, torch.Tensor): | |
| raise TypeError(f"query must be a torch.Tensor, got {type(query)}") | |
| if not isinstance(key, torch.Tensor): | |
| raise TypeError(f"key must be a torch.Tensor, got {type(key)}") | |
| if not isinstance(value, torch.Tensor): | |
| raise TypeError(f"value must be a torch.Tensor, got {type(value)}") | |
| if not isinstance(q_mask, torch.Tensor): | |
| raise TypeError(f"q_mask must be a torch.Tensor, got {type(q_mask)}") | |
| if k_mask is None or not isinstance(k_mask, torch.Tensor): | |
| raise TypeError(f"k_mask must be a torch.Tensor, got {type(k_mask)}") | |
| if query.dim() != 3 or key.dim() != 3 or value.dim() != 3: | |
| raise ValueError( | |
| "query/key/value must be 3D tensors of shape (B, S, D); " | |
| f"got q={tuple(query.shape)}, k={tuple(key.shape)}, v={tuple(value.shape)}" | |
| ) | |
| Bq, Sq, Dq = query.shape | |
| Bk, Sk, Dk = key.shape | |
| Bv, Sv, Dv = value.shape | |
| if not (Dq == Dk == Dv == self.d_model): | |
| raise ValueError( | |
| f"Last dimension must equal d_model={self.d_model}; got Dq={Dq}, Dk={Dk}, Dv={Dv}" | |
| ) | |
| if not (Bq == Bk == Bv): | |
| raise ValueError(f"Batch size mismatch: q={Bq}, k={Bk}, v={Bv}") | |
| if Sk != Sv: | |
| raise ValueError(f"Key/Value seq length mismatch: Sk={Sk} vs Sv={Sv}") | |
| # ---- padding mask validation ---- | |
| if q_mask.dtype != torch.bool: | |
| raise TypeError(f"q_mask must be boolean, got {q_mask.dtype}") | |
| if k_mask.dtype != torch.bool: | |
| raise TypeError(f"k_mask must be boolean, got {k_mask.dtype}") | |
| if q_mask.dim() != 4 or k_mask.dim() != 4: | |
| raise ValueError( | |
| "q_mask and k_mask must be 4D tensors shaped (B, H, 1, S); " | |
| f"got {tuple(q_mask.shape)} and {tuple(k_mask.shape)}" | |
| ) | |
| # ---- project and split into heads -> (B, H, S, Dh) | |
| q = split_heads(self.query_linear(query), self.num_heads) | |
| k = split_heads(self.key_linear(key), self.num_heads) | |
| v = split_heads(self.value_linear(value), self.num_heads) | |
| pad_mask = create_qk_padding_mask(q_mask, k_mask) | |
| combined_mask = combine_masks(causal_mask, pad_mask) | |
| p = self.dropout_rate if self.training else 0.0 | |
| attn = cast(Tensor, calculate_attention(q, k, v, combined_mask, dropout_p=p)) | |
| out = join_heads(attn) # (B, Sq, D) | |
| return self.output_linear(out) | |