import torch from torch import Tensor import torch.nn as nn from safetensors.torch import load_model from jaxtyping import Bool, Int, Float from huggingface_hub import hf_hub_download from embedding import InputEmbeddings, PositionalEncoding from modules import Encoder, Decoder import config class Generator(nn.Module): """ Implements the final Linear (Projection) layer and Softmax. This module takes the final output of the Decoder stack (B, T, D) and projects it onto the vocabulary space (B, T, vocab_size) to produce the logits. (This layer's weights can be tied with the target embedding layer, which we will handle in the main 'Transformer' model class). """ def __init__(self, d_model: int, vocab_size: int) -> None: """ Initializes the Generator (Output Projection) layer. Args: d_model (int): The dimension of the model (D). vocab_size (int): The size of the target vocabulary. """ super().__init__() self.proj: nn.Linear = nn.Linear(d_model, vocab_size, bias=False) def forward( self, x: Float[Tensor, "B T_tgt D"] ) -> Float[Tensor, "B T_tgt vocab_size"]: """ Forward pass for the Generator. Args: x (Tensor): The final output tensor from the Decoder stack. Returns: Tensor: The output logits over the vocabulary. """ # (B, T_tgt, D) -> (B, T_tgt, vocab_size) logits = self.proj(x) return logits class Transformer(nn.Module): """ The main Transformer model architecture, combining the Encoder and Decoder stacks, as described in "Attention Is All You Need". This implementation follows modern best practices (Pre-LN) and is designed for a sequence-to-sequence task (e.g., translation). """ def __init__( self, src_vocab_size: int, tgt_vocab_size: int, d_model: int, n_heads: int, n_layers: int, # N=6 in the paper d_ff: int, dropout: float = 0.1, max_seq_len: int = 512, # Max length for positional encoding ) -> None: """ Initializes the full Transformer model. Args: src_vocab_size (int): Vocabulary size for the source language. tgt_vocab_size (int): Vocabulary size for the target language. d_model (int): The dimension of the model (D). n_heads (int): The number of attention heads (H). n_layers (int): The number of Encoder/Decoder layers (N). d_ff (int): The inner dimension of the Feed-Forward Network (D_FF). dropout (float): The dropout rate. max_seq_len (int): The maximum sequence length for positional encoding. """ super().__init__() self.d_model = d_model # --- 1. Source (Encoder) Embeddings --- # We create two separate embedding layers self.src_embed: InputEmbeddings = InputEmbeddings(d_model, src_vocab_size) # --- 2. Target (Decoder) Embeddings --- self.tgt_embed: InputEmbeddings = InputEmbeddings(d_model, tgt_vocab_size) # --- 3. Positional Encoding --- # We use "one" PositionalEncoding module # and share it for both source and target. self.pos_enc: PositionalEncoding = PositionalEncoding( d_model, max_seq_len, dropout ) # --- 4. Encoder Stack --- self.encoder: Encoder = Encoder(d_model, n_heads, d_ff, n_layers, dropout) # --- 5. Decoder Stack --- self.decoder: Decoder = Decoder(d_model, n_heads, d_ff, n_layers, dropout) # --- 6. Final Output Projection (Generator) --- self.generator: Generator = Generator(d_model, tgt_vocab_size) # --- Weight Typing --- # We tie the weights of the target embedding and the generator. # This saves parameters and improves performance. self.generator.proj.weight = self.tgt_embed.token_emb.weight # --- Initialize weights --- # This is crucial for stable training. self.apply(self._init_weights) def _init_weights(self, module: nn.Module): """ Applies Xavier/Glorot uniform initialization to linear layers. This is a common and effective initialization strategy. """ if isinstance(module, nn.Linear): nn.init.xavier_uniform_(module.weight) if module.bias is not None: nn.init.constant_(module.bias, 0) elif isinstance(module, nn.Embedding): # Initialize embeddings (e.g., from a normal distribution) nn.init.normal_(module.weight, mean=0, std=self.d_model**-0.5) def forward( self, src: Int[Tensor, "B T_src"], # Source token IDs (e.g., English) tgt: Int[Tensor, "B T_tgt"], # Target token IDs (e.g., Vietnamese) src_mask: Bool[Tensor, "B 1 1 T_src"], # Source padding mask tgt_mask: Bool[Tensor, "B 1 T_tgt T_tgt"], # Target combined mask ) -> Float[Tensor, "B T_tgt vocab_size"]: """ Defines the main forward pass of the Transformer model. Args: src (Tensor): Source sequence token IDs. tgt (Tensor): Target sequence token IDs (shifted right). src_mask (Tensor): Padding mask for the source sequence. tgt_mask (Tensor): Combined padding and look-ahead mask for the target sequence. Returns: Tensor: The output logits from the model (B, T_tgt, vocab_size). """ # 1. Encode the source sequence # (B, T_src) -> (B, T_scr, D) src_embeded = self.src_embed(src) src_with_pos = self.pos_enc(src_embeded) # (B, T_src, D) -> (B, T_src, D) # This 'memory' will be used by every DecoderLayer enc_output: Tensor = self.encoder(src_with_pos, src_mask) # 2. Decode the target sequence # (B, T_tgt) -> (B, T_tgt, D) tgt_embeded = self.tgt_embed(tgt) tgt_with_pos = self.pos_enc(tgt_embeded) # (B, T_tgt, D) -> (B, T_tgt, D) dec_output: Tensor = self.decoder(tgt_with_pos, enc_output, src_mask, tgt_mask) # 3. Generate final logits # (B, T_tgt, D) -> (B, T_tgt, vocab_size) logits: Tensor = self.generator(dec_output) return logits def load_trained_model( config_obj, checkpoint_path, device: torch.device ) -> Transformer: print("Downloading safetensors from Hub...") model_path = hf_hub_download(repo_id=config.REPO_ID, filename=config.FILENAME) print("Instantiating the Transformer model...") model = Transformer( src_vocab_size=config_obj.VOCAB_SIZE, tgt_vocab_size=config_obj.VOCAB_SIZE, d_model=config_obj.D_MODEL, n_heads=config_obj.N_HEADS, n_layers=config_obj.N_LAYERS, d_ff=config_obj.D_FF, dropout=config_obj.DROPOUT, max_seq_len=config_obj.MAX_SEQ_LEN, ).to(device) # print(f"Loading model from: {checkpoint_path}") # load_model(model, filename=checkpoint_path) print(f"Loading model from: {model_path}") load_model(model, filename=model_path) print(f"Successfully loaded trained weights from {model_path}") return model