Implemented the following parts for the Transformer model: poisional_encoding, feedfoward, encoder, and skeleton of decoder. As well as Test cases for each class respectively and visualizations

src/models/attention.py

@@ -79,11 +79,35 @@ class ScaledDotProductAttention(nn.Module):
 
         # Mask if provided
         if mask is not None:
-            scores = scores.masked_fill(mask == 0, float('-inf'))
+            # Ensure mask is boolean and on same device as scores
+            mask_bool = mask.to(dtype=torch.bool, device=scores.device)
+            # masked_fill expects broadcastable mask: True means keep, False means mask out
+            scores = scores.masked_fill(~mask_bool, float("-1e9"))
         # Applying Softmax to get attention weights
         attention_weights = F.softmax(scores, dim=-1)
         
-        return torch.matmul(attention_weights, value), attention_weights
+        # Softmax to get attention probabilities
+        p_attn = F.softmax(scores, dim=-1)
+        
+        # If mask was provided, ensure masked positions are exactly zero (and handle all-masked rows)
+        if mask is not None:
+            # Convert mask to same dtype as p_attn for multiplication
+            mask_float = mask.to(dtype=p_attn.dtype, device=p_attn.device)
+            # Broadcast-multiply (zero out masked key positions)
+            p_attn = p_attn * mask_float
+            # Replace any NaNs (can occur when a row was entirely -inf prior to softmax) with 0.0
+            # torch.nan_to_num is efficient and handles negative/positive inf as well
+            p_attn = torch.nan_to_num(p_attn, nan=0.0, posinf=0.0, neginf=0.0)
+
+            # re-normalize rows that still have non-zero sum, this is not strictly necessary
+            # if mask is correct, but safe to avoid tiny numerical issues:
+            row_sums = p_attn.sum(dim=-1, keepdim=True)
+            # Avoid division by zero; only divide where row_sums > 0
+            nonzero_rows = row_sums > 0
+            p_attn = torch.where(nonzero_rows, p_attn / (row_sums + 1e-12), p_attn)
+        
+        output = torch.matmul(p_attn, value)
+        return output, p_attn
     
 # --------------- Multi-Head Attention ---------------
 

src/models/decoder.py

@@ -0,0 +1,300 @@
+"""
+Transformer Decoder layout (Pre-LN)
+
+Contents:
+- create_causal_mask: utility to build a causal (subsequent) mask
+- TransformerDecoderLayer: one decoder block (masked self-attn, cross-attn, FFN)
+- TransformerDecoder: embedding/pos-encoding + stack of decoder layers + generation helpers
+
+Notes / conventions:
+- Pre-LN (LayerNorm before each sublayer) is assumed for stability (consistent with your encoder).
+- MultiHeadAttention, FeedForward, PositionalEncoding are expected to live in src/models
+  (you already implemented them).
+- Masks use boolean semantics: True = allowed, False = masked.
+- The decoder API supports:
+    - inputs: token ids (LongTensor, (B, T)) or embeddings ((B, T, d_model))
+    - memory: encoder outputs (B, S, d_model)
+    - mask arguments: tgt_mask (causal/padding), memory_mask (encoder padding)
+    - collect_attn: return attention maps per layer if requested
+- Generation helpers (greedy) are skeletons that you can extend to beam search or caching.
+
+TODO status keys:
+- [IMPLEMENT] : core implementation required
+- [OPTIONAL]  : useful enhancement (caching, beam search, advanced scheduling)
+"""
+
+from typing import Optional, Tuple, List, Union, Dict
+import math
+import torch
+import torch.nn as nn
+
+from .attention import MultiHeadAttention
+from .feedforward import FeedForward
+from .positional_encoding import PositionalEncoding
+
+
+def create_causal_mask(seq_len: int, device: Optional[torch.device] = None) -> torch.Tensor:
+    """
+    Create a square causal mask of shape (seq_len, seq_len).
+    True indicates allowed positions; False indicates masked (future) positions.
+
+    Returns:
+        mask: torch.BoolTensor of shape (seq_len, seq_len)
+    """
+    # return a mask with True on and below diagonal, False above diagonal 
+    # The torch.trui function does masking, which is the idea of zeroing all the values in a matrix below a certain diagonal
+    mask = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1)
+    # mask has True above diagonal (to be masked). Want True for allowed, so invert:
+    return ~mask # shape (seq_len, seq_len) or (T, T)
+
+
+class TransformerDecoderLayer(nn.Module):
+    """
+    One decoder layer with:
+      - Masked self-attention (query/key/value = tgt)
+      - Encoder-Decoder cross-attention (query = tgt, key/value = memory)
+      - Position-wise FeedForward
+      - Pre-LN + residuals + dropout
+
+    Args:
+      d_model: model hidden size
+      num_heads: number of attention heads
+      d_ff: ff intermediate size
+      dropout: dropout for residuals / FFN
+    """
+
+    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float = 0.1):
+        super().__init__()
+        # NOTE: instantiate internal MHA with dropout=0.0 and manage dropout at layer-level
+        self.self_attn = MultiHeadAttention(d_model=d_model, num_heads=num_heads, dropout=0.0)
+        self.cross_attn = MultiHeadAttention(d_model=d_model, num_heads=num_heads, dropout=0.0)
+        self.ffn = FeedForward(d_model=d_model, d_ff=d_ff, dropout=dropout)
+
+        # LayerNorms (Pre-LN)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+
+        # Dropouts applied after sublayers (on sublayer outputs before residual add)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        tgt: torch.Tensor,
+        memory: torch.Tensor,
+        tgt_mask: Optional[torch.Tensor] = None,
+        memory_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        """
+        Forward pass for one decoder layer.
+
+        Args:
+            tgt: (batch, tgt_len, d_model)
+            memory: (batch, src_len, d_model)  -- encoder outputs
+            tgt_mask: optional (batch, tgt_len, tgt_len) or (batch, 1, tgt_len, tgt_len)
+            memory_mask: optional (batch, src_len, src_len) or (batch, 1, tgt_len, src_len)
+
+        Returns:
+            output: (batch, tgt_len, d_model)
+            attn_maps: dict with keys 'self' and 'cross' containing attention tensors
+                       shapes: (batch, num_heads, tgt_len, tgt_len) and (batch, num_heads, tgt_len, src_len)
+        """
+        # TODO [IMPLEMENT] Self-attention (Pre-LN)
+        # x_norm = self.norm1(tgt)
+        # self_out, self_attn = self.self_attn(x_norm, x_norm, x_norm, tgt_mask)
+        # tgt = tgt + self.dropout1(self_out)
+
+        # TODO [IMPLEMENT] Cross-attention (Pre-LN)
+        # x_norm = self.norm2(tgt)
+        # cross_out, cross_attn = self.cross_attn(x_norm, memory, memory, memory_mask)
+        # tgt = tgt + self.dropout2(cross_out)
+
+        # TODO [IMPLEMENT] Feed-forward (Pre-LN)
+        # x_norm = self.norm3(tgt)
+        # ffn_out = self.ffn(x_norm)
+        # tgt = tgt + self.dropout3(ffn_out)
+
+        # TODO [RETURN] Return (tgt, {"self": self_attn, "cross": cross_attn})
+        raise NotImplementedError("TransformerDecoderLayer.forward: implement Pre-LN pipeline")
+
+
+class TransformerDecoder(nn.Module):
+    """
+    Full decoder: token embedding + positional encoding + stack of decoder layers.
+    Also supports simple greedy decoding.
+
+    Args:
+        vocab_size: for token embeddings (if using token ids)
+        d_model, num_layers, num_heads, d_ff, dropout, max_len, pad_token_id: same semantics as encoder
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        d_model: int = 512,
+        num_layers: int = 6,
+        num_heads: int = 8,
+        d_ff: int = 2048,
+        dropout: float = 0.1,
+        max_len: int = 512,
+        pad_token_id: Optional[int] = None,
+    ):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.pad_token_id = pad_token_id
+
+        # Token embedding (used if inputs are token ids)
+        self.embedding = nn.Embedding(vocab_size, d_model)
+
+        # Positional encoding
+        self.pos_encoder = PositionalEncoding(d_model=d_model, max_len=max_len, dropout=dropout)
+
+        # Decoder layers
+        self.layers = nn.ModuleList(
+            [
+                TransformerDecoderLayer(d_model=d_model, num_heads=num_heads, d_ff=d_ff, dropout=dropout)
+                for _ in range(num_layers)
+            ]
+        )
+
+        # Final layer norm for Pre-LN stacks
+        self.final_norm = nn.LayerNorm(d_model)
+
+        # Output projection to vocabulary (logits)
+        self.output_projection = nn.Linear(d_model, vocab_size)
+
+        # Input dropout (after pos encoding)
+        self.input_dropout = nn.Dropout(dropout)
+
+    def _build_padding_mask_from_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """
+        Build (batch, seq, seq) boolean mask from input ids and pad_token_id.
+        True = allowed, False = masked.
+        """
+        assert self.pad_token_id is not None, "pad_token_id must be set to build mask from ids"
+        pad_mask = (input_ids != self.pad_token_id)  # (B, S)
+        attn_mask = pad_mask.unsqueeze(1) & pad_mask.unsqueeze(2)  # (B, S, S)
+        return attn_mask
+
+    def forward(
+        self,
+        inputs: torch.Tensor,
+        memory: torch.Tensor,
+        tgt_mask: Optional[torch.Tensor] = None,
+        memory_mask: Optional[torch.Tensor] = None,
+        collect_attn: bool = False,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[Dict[str, torch.Tensor]]]]:
+        """
+        Forward pass for the decoder stack.
+
+        Args:
+            inputs: token ids (B, T) or embeddings (B, T, d_model)
+            memory: encoder outputs (B, S, d_model)
+            tgt_mask: optional mask for decoder self-attention. If None, a causal mask will be created.
+                      Mask shapes: (B, T, T) or (B, 1, T, T)
+            memory_mask: optional mask over memory (B, S, S) or (B, 1, T, S)
+            collect_attn: if True returns (logits/outputs, [per-layer attn dicts])
+
+        Returns:
+            logits: (B, T, vocab_size) or (B, T, d_model) if you prefer returning hidden states
+            or (logits, attn_list) if collect_attn True
+        """
+        # Inputs: if token ids, embed and scale; else assume embeddings
+        if inputs.dim() == 2:  # token ids
+            x = self.embedding(inputs) * math.sqrt(self.d_model)
+        elif inputs.dim() == 3:
+            x = inputs
+        else:
+            raise ValueError("inputs must be (B, T) token ids or (B, T, d_model) embeddings")
+
+        # Positional encoding + dropout
+        x = self.pos_encoder(x)
+        x = self.input_dropout(x)
+
+        # Build tgt_mask if not provided: combine causal mask and padding mask if available
+        seq_len = x.size(1)
+        if tgt_mask is None:
+            # base causal mask (T, T)
+            causal = create_causal_mask(seq_len, device=x.device)  # [TODO implement]
+            # expand to batch dim later if padding present
+            if inputs.dim() == 2 and self.pad_token_id is not None:
+                padding_mask = self._build_padding_mask_from_ids(inputs)  # (B, T, T)
+                # combine: True only where both causal and padding allow attention
+                # TODO: ensure shapes align; broadcast causal to (1, T, T) then & with padding_mask
+                raise NotImplementedError("tgt_mask construction: combine causal + padding_mask")
+            else:
+                # TODO: Broadcast causal to (1, T, T) or (B, 1, T, T) depending on downstream expectations
+                raise NotImplementedError("tgt_mask construction: broadcast causal mask for batch")
+
+        # Ensure memory_mask is boolean on correct device if provided
+        if memory_mask is not None:
+            memory_mask = memory_mask.to(dtype=torch.bool, device=x.device)
+
+        attn_list: List[Dict[str, torch.Tensor]] = []
+
+        # Pass through layers
+        for layer in self.layers:
+            x, attn = layer(x, memory, tgt_mask=tgt_mask, memory_mask=memory_mask)
+            if collect_attn:
+                attn_list.append(attn)
+
+        x = self.final_norm(x)  # Pre-LN final normalization
+
+        logits = self.output_projection(x)  # (B, T, vocab)
+        if collect_attn:
+            return logits, attn_list
+        return logits
+
+    # ---------------------------------------------------------------------
+    # Generation / inference helpers (skeletons)
+    # ---------------------------------------------------------------------
+    def greedy_decode(
+        self,
+        memory: torch.Tensor,
+        max_len: int,
+        start_token_id: int,
+        end_token_id: Optional[int] = None,
+        device: Optional[torch.device] = None,
+    ) -> torch.LongTensor:
+        """
+        Greedy autoregressive decoding using the decoder stack.
+
+        Args:
+            memory: encoder outputs (B, S, d_model)
+            max_len: maximum target length to generate
+            start_token_id: BOS token id
+            end_token_id: optional EOS token id to stop early
+        Returns:
+            generated: (B, T_out) long tensor of token ids
+        """
+        # TODO [IMPLEMENT]:
+        #  - Start with tensor of shape (B, 1) filled with start_token_id
+        #  - Repeatedly call decoder.forward in incremental mode (or full forward with causal mask)
+        #  - At each step pick argmax over logits and append to sequence
+        #  - Stop if all sequences produced end_token_id or max_len reached
+        raise NotImplementedError("greedy_decode: implement autoregressive generation loop")
+
+    # Optional: incremental step method with caching of past keys/values for speed
+    def step(
+        self,
+        last_token_ids: torch.LongTensor,
+        memory: torch.Tensor,
+        cache: Optional[Dict] = None,
+    ) -> Tuple[torch.Tensor, Dict]:
+        """
+        Single-step decoder that returns logits for the next token given last_token_ids.
+
+        Args:
+            last_token_ids: (B, 1) tokens at current time step
+            memory: encoder outputs
+            cache: optional dict storing per-layer cached keys/values
+
+        Returns:
+            logits: (B, vocab_size)
+            new_cache: updated cache
+        """
+        # TODO [OPTIONAL]: implement fast incremental decoding caching keys/values per layer
+        raise NotImplementedError("step: incremental decoding (optional optimization)")

src/models/encoder.py

@@ -0,0 +1,203 @@
+"""
+Transformer encoder implementation (Pre-LN).
+
+Contains:
+- TransformerEncoderLayer: one encoder block (self-attention + FFN with residuals + LayerNorm)
+- TransformerEncoder: embedding + positional encoding + stack of encoder layers
+
+Design choices:
+- Pre-LN (LayerNorm before each sublayer) for stable training.
+- The FeedForward module is position-wise and does NOT include residuals or normalization.
+- MultiHeadAttention handles mask broadcasting from (B, S, S) -> (B, 1, S, S) internally.
+- The encoder accepts either token ids (LongTensor) or precomputed embeddings (FloatTensor).
+  If you pass token ids, provide vocab_size when constructing the encoder and optionally pad_token_id.
+- Optionally collect attention weights by passing collect_attn=True to forward().
+"""
+
+from typing import Optional, Tuple, List, Union
+
+import math
+import torch
+import torch.nn as nn
+
+from .attention import MultiHeadAttention
+from .feedforward import FeedForward
+from .positional_encoding import PositionalEncoding
+
+
+class TransformerEncoderLayer(nn.Module):
+    """
+    Single Transformer encoder layer (Pre-LN).
+
+    Args:
+        d_model: model hidden size
+        num_heads: number of attention heads
+        d_ff: hidden dimension of the position-wise feed-forward network
+        dropout: dropout probability applied to sublayer outputs
+    """
+
+    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float = 0.1):
+        super().__init__()
+        self.self_attn = MultiHeadAttention(d_model=d_model, num_heads=num_heads, dropout=0.0)
+        # set MHA internal dropout to 0.0 and use dropout1/dropout2 in the layer
+        self.ffn = FeedForward(d_model=d_model, d_ff=d_ff, dropout=dropout)
+        
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Forward pass for the encoder layer.
+
+        Args:
+            x: (batch, seq_len, d_model) - input embeddings / representations
+            mask: optional attention mask, shape either (batch, seq_q, seq_k) or (batch, 1, seq_q, seq_k)
+
+        Returns:
+            x: (batch, seq_len, d_model)
+            If you want attention weights, set collect_attn externally (the encoder stack can collect them).
+        """
+        # Self-attention sublayer (Pre-LN)
+        x_norm = self.norm1(x)  # Pre-LN
+        # self_attn expects query, key, value; for encoder they are the same
+        attn_out, attn_weights = self.self_attn(x_norm, x_norm, x_norm, mask)
+        x = x + self.dropout1(attn_out)
+
+        # Feed-forward sublayer (Pre-LN)
+        x_norm = self.norm2(x)
+        ffn_out = self.ffn(x_norm)
+        x = x + self.dropout2(ffn_out)
+
+        # Return output (and optionally attn_weights if caller wants to collect them)
+        return x, attn_weights
+
+
+class TransformerEncoder(nn.Module):
+    """
+    Full encoder: token embedding + positional encoding + N encoder layers.
+
+    Args:
+        vocab_size: vocabulary size (ignored if you always pass embeddings)
+        d_model: model hidden size
+        num_layers: number of encoder layers to stack
+        num_heads: number of attention heads
+        d_ff: hidden dimension in FFN
+        dropout: dropout probability (applied in positional encoding & residuals)
+        max_len: maximum sequence length for positional encoding
+        pad_token_id: optional token id for padding; if provided and input is token ids,
+                      a padding mask will be constructed automatically
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        d_model: int = 512,
+        num_layers: int = 6,
+        num_heads: int = 8,
+        d_ff: int = 2048,
+        dropout: float = 0.1,
+        max_len: int = 512,
+        pad_token_id: Optional[int] = None,
+    ):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.pad_token_id = pad_token_id
+
+        # Token embedding (only used if forward receives token ids)
+        self.embedding = nn.Embedding(vocab_size, d_model)
+
+        # Positional encoding (adds dropout internally)
+        self.pos_encoder = PositionalEncoding(d_model=d_model, max_len=max_len, dropout=dropout)
+
+        # Encoder layers stack
+        self.layers = nn.ModuleList(
+            [TransformerEncoderLayer(d_model=d_model, num_heads=num_heads, d_ff=d_ff, dropout=dropout)
+             for _ in range(num_layers)]
+        )
+
+        # Final LayerNorm for Pre-LN stacks (recommended)
+        self.final_norm = nn.LayerNorm(d_model)
+
+        # Dropout applied after embedding + positional encoding (paper uses this)
+        self.input_dropout = nn.Dropout(dropout)
+
+    def _build_padding_mask(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """
+        Build a 3D attention mask (batch, seq, seq) from input_ids and pad_token_id.
+        True indicates valid positions; False indicates masked (pad).
+        """
+        assert self.pad_token_id is not None, "pad_token_id must be set to build padding mask from ids."
+        # mask shape: (batch, seq) where True = token kept (non-pad)
+        pad_mask = (input_ids != self.pad_token_id)
+        # Convert to (batch, seq_q, seq_k) by outer product broadcasting
+        # We want positions that are valid as both query and key
+        attn_mask = pad_mask.unsqueeze(1) & pad_mask.unsqueeze(2)
+        # attn_mask dtype should be bool
+        return attn_mask
+
+    def forward(
+        self,
+        inputs: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        collect_attn: bool = False,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[torch.Tensor]]]:
+        """
+        Forward through the encoder.
+
+        Args:
+            inputs: either
+                - token ids: LongTensor of shape (batch, seq)
+                - embeddings: FloatTensor of shape (batch, seq, d_model)
+            mask: optional attention mask. If None and pad_token_id is set and inputs are token ids,
+                  a padding mask will be created automatically with shape (batch, seq, seq).
+                  The mask should be boolean where True indicates allowed attention.
+            collect_attn: if True, returns (output, [attn_weights_per_layer]) where each entry is (batch, num_heads, seq, seq)
+
+        Returns:
+            output: (batch, seq, d_model)
+            or (output, attn_list) if collect_attn True
+        """
+        # If inputs are token ids, embed them; otherwise assume they are embeddings
+        if inputs.dim() == 2:  # token ids
+            if self.embedding is None:
+                raise ValueError("Encoder was not constructed with an embedding layer.")
+            x = self.embedding(inputs) * math.sqrt(self.d_model)
+        elif inputs.dim() == 3:  # already embeddings
+            x = inputs
+        else:
+            raise ValueError("inputs must be (batch, seq) token ids or (batch, seq, d_model) embeddings")
+
+        # Positional encoding + dropout
+        x = self.pos_encoder(x)
+        x = self.input_dropout(x)
+
+        # Build mask if needed
+        if mask is None and inputs.dim() == 2 and self.pad_token_id is not None:
+            mask = self._build_padding_mask(inputs)
+
+        # Ensure mask is boolean and on the same device
+        if mask is not None:
+            mask = mask.to(dtype=torch.bool, device=x.device)
+
+        attn_weights_per_layer: List[torch.Tensor] = []
+
+        # Pass through each encoder layer (optionally collect attn)
+        for layer in self.layers:
+            x, attn = layer(x, mask=mask)
+            if collect_attn:
+                attn_weights_per_layer.append(attn)
+
+        # Final normalization (Pre-LN stack)
+        x = self.final_norm(x)
+
+        if collect_attn:
+            return x, attn_weights_per_layer
+        return x

src/models/feedforward.py

@@ -0,0 +1,40 @@
+"""
+Position-wise Feed-Forward Network.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+from typing import Literal
+
+class FeedForward(nn.Module):
+    """
+    FFN(x) = max(0, xW₁ + b₁)W₂ + b₂
+    
+    Or with GELU: FFN(x) = GELU(xW₁ + b₁)W₂ + b₂
+    """
+    
+    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1, activation: Literal["gelu", "relu"] = "gelu"):
+        super().__init__()
+        self.linear1 = nn.Linear(d_model, d_ff) # w_1
+        self.activation = nn.GELU() if activation == 'gelu' else nn.ReLU()
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(d_ff, d_model) # w_2
+        
+        # Weight Initialization
+        init.xavier_uniform_(self.linear1.weight)
+        init.zeros_(self.linear1.bias)
+        init.xavier_uniform_(self.linear2.weight)
+        init.zeros_(self.linear2.bias)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x: (batch, seq_len, d_model)
+        returns: (batch, seq_len, d_model)
+        """
+        x = self.linear1(x)       # (batch, seq_len, d_ff)
+        x = self.activation(x)    # activation
+        x = self.dropout(x)       # dropout
+        x = self.linear2(x)       # (batch, seq_len, d_model)
+        return x
+    

src/models/positional_encoding.py

@@ -0,0 +1,79 @@
+# src/models/positional_encoding.py
+
+"""
+Positional Encoding for Transformer models.
+
+Injects information about the position of tokens in a sequence, since
+self-attention has no inherent notion of token order.
+"""
+
+import torch
+import torch.nn as nn
+import math
+
+class PositionalEncoding(nn.Module):
+    """
+    Implements the sinusoidal positional encoding from "Attention Is All You Need".
+    
+    Formula:
+        PE(pos, 2i)   = sin(pos / 10000^(2i/d_model))
+        PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
+    
+    Where:
+        pos: position in sequence (0 to max_len-1)
+        i: dimension index (0 to d_model/2)
+    
+    Args:
+        d_model: Dimension of the model embeddings
+        max_len: Maximum sequence length to pre-compute
+        dropout: Dropout probability to apply after adding positional encoding
+    
+    Shape:
+        Input: (batch, seq_len, d_model)
+        Output: (batch, seq_len, d_model)
+    
+    Example:
+        >>> pos_enc = PositionalEncoding(d_model=512, max_len=5000)
+        >>> x = torch.randn(32, 100, 512)  # (batch, seq, d_model)
+        >>> output = pos_enc(x)
+        >>> output.shape
+        torch.Size([32, 100, 512])
+    """
+    
+    def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        # Create a tensor of positions: [0, 1, 2, ..., max_len-1]
+        # Create a tensor of dimension indices: [0, 1, 2, ..., d_model-1]
+        # Compute the division term: 10000^(2i/d_model)
+        # Apply sin to even indices, cos to odd indices
+        # Register as buffer (not a parameter, but part of state_dict)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
+        )
+        pe = torch.zeros(max_len, d_model)
+        pe[:, 0::2] = torch.sin(position * div_term)  # Even indices
+        pe[:, 1::2] = torch.cos(position * div_term)  # Odd indices
+        pe = pe.unsqueeze(0)
+        self.register_buffer("pe", pe)
+    
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Add positional encoding to input embeddings.
+        
+        Args:
+            x: Input embeddings (batch, seq_len, d_model)
+            
+        Returns:
+            x with positional encoding added (batch, seq_len, d_model)
+        """
+        # Get sequence length from input
+        # Add the appropriate slice of positional encoding
+        # Apply dropout
+        # Return result
+        x = x + self.pe[:, : x.size(1)].requires_grad_(False)
+        # self.pe contains pre-computed encodings for all positions
+        # just need to add the first seq_len positions to x
+        return self.dropout(x)

tests/test_models/test_encoder.py

@@ -0,0 +1,176 @@
+import math
+import torch
+import pytest
+from src.models.encoder import TransformerEncoder
+
+
+def test_encoder_token_ids_and_padding_mask_and_grad():
+    """
+    Test using token ids as input, automatic padding mask creation when pad_token_id is provided,
+    output shape, and that gradients flow through the model.
+    """
+    torch.manual_seed(0)
+    vocab_size = 50
+    pad_token_id = 0
+    d_model = 64
+    num_layers = 3
+    num_heads = 8
+    d_ff = 128
+    batch_size = 2
+    seq_len = 12
+
+    encoder = TransformerEncoder(
+        vocab_size=vocab_size,
+        d_model=d_model,
+        num_layers=num_layers,
+        num_heads=num_heads,
+        d_ff=d_ff,
+        dropout=0.1,
+        max_len=seq_len,
+        pad_token_id=pad_token_id,
+    )
+
+    # create inputs with some padding at the end
+    input_ids = torch.randint(1, vocab_size, (batch_size, seq_len), dtype=torch.long)
+    input_ids[0, -3:] = pad_token_id  # first sample has last 3 tokens as padding
+    input_ids[1, -1:] = pad_token_id  # second sample has last token as padding
+
+    # Forward pass (token ids)
+    out = encoder(input_ids)  # default collect_attn=False
+    assert out.shape == (batch_size, seq_len, d_model)
+
+    # Check gradients flow
+    loss = out.sum()
+    loss.backward()
+    grads = [p.grad for p in encoder.parameters() if p.requires_grad]
+    assert any(g is not None for g in grads), "No gradients found on any parameter"
+
+
+def test_encoder_embeddings_input_and_collect_attn():
+    """
+    Test passing pre-computed embeddings to the encoder, collecting attention weights,
+    and verify shapes of attention maps per layer.
+    """
+    torch.manual_seed(1)
+    vocab_size = 100  # not used in this test
+    d_model = 48
+    num_layers = 4
+    num_heads = 6
+    d_ff = 128
+    batch_size = 1
+    seq_len = 10
+
+    encoder = TransformerEncoder(
+        vocab_size=vocab_size,
+        d_model=d_model,
+        num_layers=num_layers,
+        num_heads=num_heads,
+        d_ff=d_ff,
+        dropout=0.0,
+        max_len=seq_len,
+        pad_token_id=None,
+    )
+
+    # Create random embeddings directly
+    embeddings = torch.randn(batch_size, seq_len, d_model)
+
+    out, attn_list = encoder(embeddings, mask=None, collect_attn=True)
+    assert out.shape == (batch_size, seq_len, d_model)
+    assert isinstance(attn_list, list)
+    assert len(attn_list) == num_layers
+
+    # Each attention weight tensor should have shape (batch, num_heads, seq, seq)
+    for attn in attn_list:
+        assert attn.shape == (batch_size, num_heads, seq_len, seq_len)
+
+
+def test_mask_accepts_3d_and_4d_and_broadcasts():
+    """
+    Test that a provided 3D mask (batch, seq, seq) and an equivalent 4D mask
+    (batch, 1, seq, seq) produce outputs of the same shape and do not error.
+    """
+    torch.manual_seed(2)
+    vocab_size = 40
+    d_model = 32
+    num_layers = 2
+    num_heads = 4
+    d_ff = 64
+    batch_size = 2
+    seq_len = 7
+
+    encoder = TransformerEncoder(
+        vocab_size=vocab_size,
+        d_model=d_model,
+        num_layers=num_layers,
+        num_heads=num_heads,
+        d_ff=d_ff,
+        dropout=0.0,
+        max_len=seq_len,
+        pad_token_id=None,
+    )
+
+    # Create dummy embeddings
+    embeddings = torch.randn(batch_size, seq_len, d_model)
+
+    # 3D mask: True indicates allowed attention
+    mask3 = torch.ones(batch_size, seq_len, seq_len, dtype=torch.bool)
+    mask3[:, :, -2:] = False  # mask out last two keys
+
+    # 4D mask equivalent
+    mask4 = mask3.unsqueeze(1)  # (B, 1, S, S)
+
+    out3 = encoder(embeddings, mask=mask3)
+    out4 = encoder(embeddings, mask=mask4)
+
+    assert out3.shape == (batch_size, seq_len, d_model)
+    assert out4.shape == (batch_size, seq_len, d_model)
+    # Outputs should be finite and not NaN
+    assert torch.isfinite(out3).all()
+    assert torch.isfinite(out4).all()
+
+
+def test_train_eval_determinism_and_dropout_effect():
+    """
+    Validate that in train mode with dropout enabled, repeated forwards differ,
+    and in eval mode they are equal (deterministic).
+    """
+    torch.manual_seed(3)
+    vocab_size = 60
+    pad_token_id = 0
+    d_model = 64
+    num_layers = 2
+    num_heads = 8
+    d_ff = 128
+    batch_size = 2
+    seq_len = 9
+
+    encoder = TransformerEncoder(
+        vocab_size=vocab_size,
+        d_model=d_model,
+        num_layers=num_layers,
+        num_heads=num_heads,
+        d_ff=d_ff,
+        dropout=0.4,
+        max_len=seq_len,
+        pad_token_id=pad_token_id,
+    )
+
+    # token ids with occasional padding
+    input_ids = torch.randint(1, vocab_size, (batch_size, seq_len), dtype=torch.long)
+    input_ids[0, -2:] = pad_token_id
+
+    # Training mode: randomness due to dropout -> outputs should likely differ
+    encoder.train()
+    out1 = encoder(input_ids)
+    out2 = encoder(input_ids)
+    assert not torch.allclose(out1, out2), "Outputs identical in train mode despite dropout"
+
+    # Eval mode: deterministic
+    encoder.eval()
+    out3 = encoder(input_ids)
+    out4 = encoder(input_ids)
+    assert torch.allclose(out3, out4), "Outputs differ in eval mode"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-q"])

tests/test_models/test_encoder_layer.py

@@ -0,0 +1,78 @@
+import torch
+import pytest
+from src.models.encoder import TransformerEncoderLayer
+
+
+def test_output_shape_and_grad():
+    """
+    The encoder layer should preserve the input shape (batch, seq_len, d_model)
+    and gradients should flow to parameters.
+    """
+    d_model, num_heads, d_ff = 64, 8, 256
+    batch_size, seq_len = 2, 10
+
+    layer = TransformerEncoderLayer(d_model=d_model, num_heads=num_heads, d_ff=d_ff, dropout=0.0)
+    x = torch.randn(batch_size, seq_len, d_model, requires_grad=True)
+
+    out = layer(x)  # should accept mask=None by default
+    assert out.shape == (batch_size, seq_len, d_model)
+
+    # simple backward to ensure gradients propagate
+    loss = out.sum()
+    loss.backward()
+
+    grads = [p.grad for p in layer.parameters() if p.requires_grad]
+    assert any(g is not None for g in grads), "No gradients found on any parameter"
+
+
+def test_dropout_behavior_train_vs_eval():
+    """
+    With dropout > 0, the outputs should differ between two forward calls in train mode
+    and be identical in eval mode.
+    """
+    torch.manual_seed(0)
+    d_model, num_heads, d_ff = 64, 8, 256
+    batch_size, seq_len = 2, 10
+
+    layer = TransformerEncoderLayer(d_model=d_model, num_heads=num_heads, d_ff=d_ff, dropout=0.5)
+    x = torch.randn(batch_size, seq_len, d_model)
+
+    layer.train()
+    out1 = layer(x)
+    out2 = layer(x)
+    # Training mode with dropout: outputs usually differ
+    assert not torch.allclose(out1, out2), "Outputs identical in train mode despite dropout"
+
+    layer.eval()
+    out3 = layer(x)
+    out4 = layer(x)
+    # Eval mode deterministic: outputs should be identical
+    assert torch.allclose(out3, out4), "Outputs differ in eval mode"
+
+
+def test_mask_broadcasting_accepts_3d_and_4d_mask():
+    """
+    The encoder layer should accept a 3D mask (batch, seq_q, seq_k) and a 4D mask
+    (batch, 1, seq_q, seq_k) and handle broadcasting across heads without error.
+    """
+    d_model, num_heads, d_ff = 64, 8, 256
+    batch_size, seq_len = 2, 7
+
+    layer = TransformerEncoderLayer(d_model=d_model, num_heads=num_heads, d_ff=d_ff, dropout=0.0)
+    x = torch.randn(batch_size, seq_len, d_model)
+
+    # 3D mask: (batch, seq, seq)
+    mask3 = torch.ones(batch_size, seq_len, seq_len, dtype=torch.bool)
+    mask3[:, :, -2:] = False  # mask out last two key positions
+    out3 = layer(x, mask=mask3)  # should not raise
+    assert out3.shape == (batch_size, seq_len, d_model)
+
+    # 4D mask: (batch, 1, seq, seq) already including head dim for broadcasting
+    mask4 = mask3.unsqueeze(1)
+    out4 = layer(x, mask=mask4)
+    assert out4.shape == (batch_size, seq_len, d_model)
+
+
+if __name__ == "__main__":
+    # Run tests interactively if needed
+    pytest.main([__file__, "-q"])

tests/test_models/test_feedforward.py

@@ -0,0 +1,57 @@
+import torch
+import pytest
+from src.models.feedforward import FeedForward
+
+
+class TestFeedForward:
+    def test_output_shape(self):
+        d_model, d_ff = 512, 2048
+        batch_size, seq_len = 2, 10
+
+        ffn = FeedForward(d_model=d_model, d_ff=d_ff, dropout=0.0)
+        x = torch.randn(batch_size, seq_len, d_model)
+        out = ffn(x)
+
+        assert out.shape == (batch_size, seq_len, d_model)
+
+    def test_dropout_changes_output(self):
+        torch.manual_seed(0)
+        d_model, d_ff = 128, 512
+        x = torch.randn(2, 8, d_model)
+
+        ffn = FeedForward(d_model=d_model, d_ff=d_ff, dropout=0.5)
+        ffn.train()
+        out1 = ffn(x)
+        out2 = ffn(x)
+        # With dropout in train mode, outputs should differ (most likely)
+        assert not torch.allclose(out1, out2)
+
+        ffn.eval()
+        out3 = ffn(x)
+        out4 = ffn(x)
+        # In eval mode (no dropout), outputs should be identical for same input
+        assert torch.allclose(out3, out4)
+
+    def test_parameter_count_and_grads(self):
+        d_model, d_ff = 64, 256
+        ffn = FeedForward(d_model=d_model, d_ff=d_ff, dropout=0.0)
+
+        # Parameter existence
+        param_names = [name for name, _ in ffn.named_parameters()]
+        assert any('linear1' in name for name in param_names)
+        assert any('linear2' in name for name in param_names)
+
+        # Parameter shapes
+        shapes = {name: p.shape for name, p in ffn.named_parameters()}
+        assert shapes.get('linear1.weight') == (d_ff, d_model)
+        assert shapes.get('linear2.weight') == (d_model, d_ff)
+        assert shapes.get('linear1.bias') == (d_ff,)
+        assert shapes.get('linear2.bias') == (d_model,)
+
+        # ensure gradients flow
+        x = torch.randn(3, 5, d_model)
+        out = ffn(x)
+        loss = out.sum()
+        loss.backward()
+        for _, p in ffn.named_parameters():
+            assert p.grad is not None

tests/test_models/test_positional_encoding.py

@@ -0,0 +1,103 @@
+# tests/test_models/test_positional_encoding.py
+
+"""
+Tests for positional encoding.
+"""
+
+import pytest
+import torch
+import matplotlib.pyplot as plt
+import seaborn as sns
+from src.models.positional_encoding import PositionalEncoding
+
+
+class TestPositionalEncoding:
+    """Test suite for PositionalEncoding."""
+    
+    def test_output_shape(self):
+        """Test that output shape matches input shape."""
+        d_model, max_len = 512, 5000
+        batch_size, seq_len = 2, 100
+        
+        pos_enc = PositionalEncoding(d_model, max_len, dropout=0.0)
+        x = torch.randn(batch_size, seq_len, d_model)
+        
+        output = pos_enc(x)
+        assert output.shape == (batch_size, seq_len, d_model)
+    
+    def test_different_sequence_lengths(self):
+        """Test with various sequence lengths."""
+        pos_enc = PositionalEncoding(d_model=256, max_len=1000, dropout=0.0)
+        
+        for seq_len in [10, 50, 100, 500]:
+            x = torch.randn(1, seq_len, 256)
+            output = pos_enc(x)
+            assert output.shape == (1, seq_len, 256)
+    
+    def test_dropout_changes_output(self):
+        """Test that dropout is applied during training."""
+        torch.manual_seed(42)
+        pos_enc = PositionalEncoding(d_model=128, dropout=0.5)
+        pos_enc.train()
+        
+        x = torch.randn(2, 10, 128)
+        
+        output1 = pos_enc(x)
+        output2 = pos_enc(x)
+        
+        # Should be different due to dropout
+        assert not torch.allclose(output1, output2)
+        
+        # In eval mode, should be deterministic
+        pos_enc.eval()
+        output3 = pos_enc(x)
+        output4 = pos_enc(x)
+        assert torch.allclose(output3, output4)
+    
+    def test_encoding_properties(self):
+        """Test mathematical properties of encoding."""
+        pos_enc = PositionalEncoding(d_model=128, max_len=100, dropout=0.0)
+        
+        # Get the raw encoding (without dropout)
+        pe = pos_enc.pe[0]  # Remove batch dimension
+        
+        # Each row should have values in [-1, 1] (sin/cos range)
+        assert (pe >= -1).all() and (pe <= 1).all()
+        
+        # Different positions should have different encodings
+        assert not torch.allclose(pe[0], pe[1])
+        assert not torch.allclose(pe[0], pe[50])
+
+
+def test_visualize_positional_encoding():
+    """
+    Visualize the positional encoding pattern.
+    Creates heatmap showing encoding values.
+    """
+    pos_enc = PositionalEncoding(d_model=128, max_len=100, dropout=0.0)
+    
+    # Get encoding matrix
+    pe = pos_enc.pe.squeeze(0).numpy()  # (max_len, d_model)
+    
+    # Plot first 50 positions and 64 dimensions
+    plt.figure(figsize=(12, 8))
+    sns.heatmap(
+        pe[:50, :64].T,
+        cmap='RdBu_r',
+        center=0,
+        xticklabels=5,
+        yticklabels=8,
+        cbar_kws={'label': 'Encoding Value'}
+    )
+    plt.xlabel('Position in Sequence')
+    plt.ylabel('Embedding Dimension')
+    plt.title('Positional Encoding Pattern\n(Notice the wave patterns with different frequencies)')
+    plt.tight_layout()
+    plt.savefig('outputs/positional_encoding_heatmap.png', dpi=150)
+    print("✅ Saved to outputs/positional_encoding_heatmap.png")
+
+
+if __name__ == "__main__":
+    import os
+    os.makedirs('outputs', exist_ok=True)
+    test_visualize_positional_encoding()