Update app/model_def.py

app/model_def.py

@@ -1,243 +1,243 @@
-import torch
-import torch.nn as nn
-import math
-
-class InputEmbedding(nn.Module):
-    def __init__(self, d_model: int, vocab_size: int):
-        super().__init__()
-        self.d_model = d_model
-        self.vocab_size = vocab_size
-        self.embed = nn.Embedding(vocab_size, d_model)
-
-    def forward(self, x):
-        return self.embed(x) * math.sqrt(self.d_model)
-
-class PositionalEncoding(nn.Module):
-    def __init__(self, d_model: int, seq_len: int, dropout: float):
-        super().__init__()
-        self.d_model = d_model
-        self.seq_len = seq_len
-        self.dropout = nn.Dropout(dropout)
-
-        # Create a matrix of shape (seq_len, d_model)
-        pe = torch.zeros(seq_len, d_model)
-        
-        # Create a vector of shape (seq_len, 1)
-        position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)
-        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
-        
-        # Apply sin to even indices
-        pe[:, 0::2] = torch.sin(position * div_term)
-        # Apply cos to odd indices
-        pe[:, 1::2] = torch.cos(position * div_term)
-
-        # Add a batch dimension
-        pe = pe.unsqueeze(0) # (1, seq_len, d_model)
-
-        # Register 'pe' as a buffer, so it's not a model parameter
-        self.register_buffer('pe', pe)
-
-    def forward(self, x):
-        x = x + (self.pe[:, :x.shape[1], :]).requires_grad_(False)
-        return self.dropout(x)
-
-class LayerNorm(nn.Module):
-    def __init__(self, d_model: int, epsilon: float = 1e-6):
-        super().__init__()
-        self.epsilon = epsilon
-        self.gamma = nn.Parameter(torch.ones(d_model)) # Multiplicative
-        self.beta = nn.Parameter(torch.zeros(d_model))  # Additive
-
-    def forward(self, x):
-        mean = x.mean(dim=-1, keepdim=True)
-        std = x.std(dim=-1, keepdim=True)
-        return self.gamma * (x - mean) / (std + self.epsilon) + self.beta
-
-class FeedForward(nn.Module):
-    def __init__(self, d_model: int, d_ff: int, dropout: float):
-        super().__init__()
-        self.layer1 = nn.Linear(d_model, d_ff)
-        self.layer2 = nn.Linear(d_ff, d_model)
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(self, x):
-        # (Batch, Seq_Len, d_model) -> (Batch, Seq_Len, d_ff) -> (Batch, Seq_Len, d_model)
-        return self.layer2(self.dropout(torch.relu(self.layer1(x))))
-
-class MHA(nn.Module):
-    def __init__(self, d_model: int, h: int, dropout: float):
-        super().__init__()
-        self.d_model = d_model
-        self.h = h
-        assert d_model % h == 0, "d_model must be divisible by h"
-
-        self.d_k = d_model // h
-        self.w_q = nn.Linear(d_model, d_model)
-        self.w_k = nn.Linear(d_model, d_model)
-        self.w_v = nn.Linear(d_model, d_model)
-        self.w_o = nn.Linear(d_model, d_model)
-        self.dropout = nn.Dropout(dropout)
-
-    @staticmethod
-    def attention(query, key, value, mask, dropout: nn.Dropout):
-        d_k = query.shape[-1]
-        attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k)
-        if mask is not None:
-            attention_scores = attention_scores.masked_fill(mask == 0, -1e9)
-        
-        attention_scores = attention_scores.softmax(dim=-1)
-        if dropout is not None:
-            attention_scores = dropout(attention_scores)
-        
-        return (attention_scores @ value), attention_scores
-
-    def forward(self, q, k, v, mask):
-        query = self.w_q(q) # (Batch, Seq_Len, d_model)
-        key = self.w_k(k)   # (Batch, Seq_Len, d_model)
-        value = self.w_v(v) # (Batch, Seq_Len, d_model)
-
-        # (Batch, Seq_Len, d_model) -> (Batch, Seq_Len, h, d_k) -> (Batch, h, Seq_Len, d_k)
-        query = query.view(query.shape[0], query.shape[1], self.h, self.d_k).transpose(1, 2)
-        key = key.view(key.shape[0], key.shape[1], self.h, self.d_k).transpose(1, 2)
-        value = value.view(value.shape[0], value.shape[1], self.h, self.d_k).transpose(1, 2)
-
-        x, self.attention_scores = MHA.attention(query, key, value, mask, self.dropout)
-
-        # (Batch, h, Seq_Len, d_k) -> (Batch, Seq_Len, h, d_k) -> (Batch, Seq_Len, d_model)
-        x = x.transpose(1, 2).contiguous().view(x.shape[0], -1, self.h * self.d_k)
-
-        return self.w_o(x)
-
-class SkipConnection(nn.Module):
-    def __init__(self, d_model: int, dropout: float):
-        super().__init__()
-        self.dropout = nn.Dropout(dropout)
-        self.norm = LayerNorm(d_model)
-
-    def forward(self, x, sublayer):
-        # Pre-Norm architecture
-        return x + self.dropout(sublayer(self.norm(x)))
-
-class EncoderBlock(nn.Module):
-    def __init__(self, self_attention: MHA, ffn: FeedForward, d_model: int, dropout: float):
-        super().__init__()
-        self.self_attention = self_attention
-        self.ffn = ffn
-        self.skip_connections = nn.ModuleList([SkipConnection(d_model, dropout) for _ in range(2)])
-
-    def forward(self, x, src_mask):
-        x = self.skip_connections[0](x, lambda x: self.self_attention(x, x, x, src_mask))
-        x = self.skip_connections[1](x, self.ffn)
-        return x
-
-class Encoder(nn.Module):
-    def __init__(self, d_model: int, layers: nn.ModuleList):
-        super().__init__()
-        self.layers = layers
-        self.norm = LayerNorm(d_model)
-
-    def forward(self, x, mask):
-        for layer in self.layers:
-            x = layer(x, mask)
-        return self.norm(x)
-
-class DecoderBlock(nn.Module):
-    def __init__(self, self_attention: MHA, cross_attention: MHA, ffn: FeedForward, d_model: int, dropout: float):
-        super().__init__()
-        self.self_attention = self_attention
-        self.cross_attention = cross_attention
-        self.ffn = ffn
-        self.skip_connections = nn.ModuleList([SkipConnection(d_model, dropout) for _ in range(3)])
-
-    def forward(self, x, encoder_output, src_mask, trg_mask):
-        x = self.skip_connections[0](x, lambda x: self.self_attention(x, x, x, trg_mask))
-        x = self.skip_connections[1](x, lambda x: self.cross_attention(x, encoder_output, encoder_output, src_mask))
-        x = self.skip_connections[2](x, self.ffn)
-        return x
-
-class Decoder(nn.Module):
-    def __init__(self, d_model: int, layers: nn.ModuleList):
-        super().__init__()
-        self.layers = layers
-        self.norm = LayerNorm(d_model)
-
-    def forward(self, x, encoder_output, src_mask, trg_mask):
-        for layer in self.layers:
-            x = layer(x, encoder_output, src_mask, trg_mask)
-        return self.norm(x)
-
-class Output(nn.Module):
-    def __init__(self, d_model: int, vocab_size: int):
-        super().__init__()
-        self.proj = nn.Linear(d_model, vocab_size)
-
-    def forward(self, x):
-        # (Batch, Seq_Len, d_model) -> (Batch, Seq_Len, vocab_size)
-        return self.proj(x)
-
-class Transformer(nn.Module):
-    def __init__(self, encoder: Encoder, decoder: Decoder, src_embed: InputEmbedding, trg_embed: InputEmbedding, src_pos: PositionalEncoding, trg_pos: PositionalEncoding, output: Output):
-        super().__init__()
-        self.encoder = encoder
-        self.decoder = decoder
-        self.src_embed = src_embed
-        self.trg_embed = trg_embed
-        self.src_pos = src_pos
-        self.trg_pos = trg_pos
-        self.output_layer = output
-
-    def encode(self, src, src_mask):
-        src = self.src_embed(src)
-        src = self.src_pos(src)
-        return self.encoder(src, src_mask)
-
-    def decode(self, encoder_output, src_mask, trg, trg_mask):
-        trg = self.trg_embed(trg)
-        trg = self.trg_pos(trg)
-        return self.decoder(trg, encoder_output, src_mask, trg_mask)
-
-    def project(self, x):
-        return self.output_layer(x)
-
-def BuildTransformer(src_vocab_size: int, trg_vocab_size: int, src_seq_len: int, trg_seq_len: int, d_model: int = 512, N: int = 6, h: int = 8, dropout: float = 0.1, d_ff: int = 2048) -> Transformer:
-    # Create the embedding layers
-    src_embed = InputEmbedding(d_model, src_vocab_size)
-    trg_embed = InputEmbedding(d_model, trg_vocab_size)
-
-    # Create the positional encoding layers
-    src_pos = PositionalEncoding(d_model, src_seq_len, dropout)
-    trg_pos = PositionalEncoding(d_model, trg_seq_len, dropout)
-
-    # Create the encoder blocks
-    encoder_blocks = []
-    for _ in range(N):
-        encoder_self_attention = MHA(d_model, h, dropout)
-        ffn = FeedForward(d_model, d_ff, dropout)
-        encoder_block = EncoderBlock(encoder_self_attention, ffn, d_model, dropout)
-        encoder_blocks.append(encoder_block)
-
-    # Create the decoder blocks
-    decoder_blocks = []
-    for _ in range(N):
-        decoder_self_attention = MHA(d_model, h, dropout)
-        cross_attention = MHA(d_model, h, dropout)
-        ffn = FeedForward(d_model, d_ff, dropout)
-        decoder_block = DecoderBlock(decoder_self_attention, cross_attention, ffn, d_model, dropout)
-        decoder_blocks.append(decoder_block)
-    
-    # Create the encoder and decoder
-    encoder = Encoder(d_model, nn.ModuleList(encoder_blocks))
-    decoder = Decoder(d_model, nn.ModuleList(decoder_blocks))
-
-    # Create the projection layer
-    projection = Output(d_model, trg_vocab_size)
-
-    # Create the transformer
-    transformer = Transformer(encoder, decoder, src_embed, trg_embed, src_pos, trg_pos, projection)
-
-    # Initialize parameters
-    for p in transformer.parameters():
-        if p.dim() > 1:
-            nn.init.xavier_uniform_(p)
-    
+import torch
+import torch.nn as nn
+import math
+
+class InputEmbedding(nn.Module):
+    def __init__(self, d_model: int, vocab_size: int):
+        super().__init__()
+        self.d_model = d_model
+        self.vocab_size = vocab_size
+        self.embed = nn.Embedding(vocab_size, d_model)
+
+    def forward(self, x):
+        return self.embed(x) * math.sqrt(self.d_model)
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model: int, seq_len: int, dropout: float):
+        super().__init__()
+        self.d_model = d_model
+        self.seq_len = seq_len
+        self.dropout = nn.Dropout(dropout)
+
+        # Create a matrix of shape (seq_len, d_model)
+        pe = torch.zeros(seq_len, d_model)
+        
+        # Create a vector of shape (seq_len, 1)
+        position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        
+        # Apply sin to even indices
+        pe[:, 0::2] = torch.sin(position * div_term)
+        # Apply cos to odd indices
+        pe[:, 1::2] = torch.cos(position * div_term)
+
+        # Add a batch dimension
+        pe = pe.unsqueeze(0) # (1, seq_len, d_model)
+
+        # Register 'pe' as a buffer, so it's not a model parameter
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        x = x + (self.pe[:, :x.shape[1], :]).requires_grad_(False)
+        return self.dropout(x)
+
+class LayerNorm(nn.Module):
+    def __init__(self, d_model: int, epsilon: float = 1e-6):
+        super().__init__()
+        self.epsilon = epsilon
+        self.gamma = nn.Parameter(torch.ones(d_model)) # Multiplicative
+        self.beta = nn.Parameter(torch.zeros(d_model))  # Additive
+
+    def forward(self, x):
+        mean = x.mean(dim=-1, keepdim=True)
+        std = x.std(dim=-1, keepdim=True)
+        return self.gamma * (x - mean) / (std + self.epsilon) + self.beta
+
+class FeedForward(nn.Module):
+    def __init__(self, d_model: int, d_ff: int, dropout: float):
+        super().__init__()
+        self.layer1 = nn.Linear(d_model, d_ff)
+        self.layer2 = nn.Linear(d_ff, d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        # (Batch, Seq_Len, d_model) -> (Batch, Seq_Len, d_ff) -> (Batch, Seq_Len, d_model)
+        return self.layer2(self.dropout(torch.relu(self.layer1(x))))
+
+class MHA(nn.Module):
+    def __init__(self, d_model: int, h: int, dropout: float):
+        super().__init__()
+        self.d_model = d_model
+        self.h = h
+        assert d_model % h == 0, "d_model must be divisible by h"
+
+        self.d_k = d_model // h
+        self.w_q = nn.Linear(d_model, d_model)
+        self.w_k = nn.Linear(d_model, d_model)
+        self.w_v = nn.Linear(d_model, d_model)
+        self.w_o = nn.Linear(d_model, d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    @staticmethod
+    def attention(query, key, value, mask, dropout: nn.Dropout):
+        d_k = query.shape[-1]
+        attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k)
+        if mask is not None:
+            attention_scores = attention_scores.masked_fill(mask == 0, -1e9)
+        
+        attention_scores = attention_scores.softmax(dim=-1)
+        if dropout is not None:
+            attention_scores = dropout(attention_scores)
+        
+        return (attention_scores @ value), attention_scores
+
+    def forward(self, q, k, v, mask):
+        query = self.w_q(q) # (Batch, Seq_Len, d_model)
+        key = self.w_k(k)   # (Batch, Seq_Len, d_model)
+        value = self.w_v(v) # (Batch, Seq_Len, d_model)
+
+        # (Batch, Seq_Len, d_model) -> (Batch, Seq_Len, h, d_k) -> (Batch, h, Seq_Len, d_k)
+        query = query.view(query.shape[0], query.shape[1], self.h, self.d_k).transpose(1, 2)
+        key = key.view(key.shape[0], key.shape[1], self.h, self.d_k).transpose(1, 2)
+        value = value.view(value.shape[0], value.shape[1], self.h, self.d_k).transpose(1, 2)
+
+        x, self.attention_scores = MHA.attention(query, key, value, mask, self.dropout)
+
+        # (Batch, h, Seq_Len, d_k) -> (Batch, Seq_Len, h, d_k) -> (Batch, Seq_Len, d_model)
+        x = x.transpose(1, 2).contiguous().view(x.shape[0], -1, self.h * self.d_k)
+
+        return self.w_o(x)
+
+class SkipConnection(nn.Module):
+    def __init__(self, d_model: int, dropout: float):
+        super().__init__()
+        self.dropout = nn.Dropout(dropout)
+        self.norm = LayerNorm(d_model)
+
+    def forward(self, x, sublayer):
+        # Pre-Norm architecture
+        return x + self.dropout(sublayer(self.norm(x)))
+
+class EncoderBlock(nn.Module):
+    def __init__(self, self_attention: MHA, ffn: FeedForward, d_model: int, dropout: float):
+        super().__init__()
+        self.attention = self_attention 
+        self.ffn = ffn
+        self.residual = nn.ModuleList([SkipConnection(d_model, dropout) for _ in range(2)]) 
+    
+    def forward(self, x, src_mask):
+        x = self.residual[0](x, lambda x: self.attention(x, x, x, src_mask)) 
+        x = self.residual[1](x, self.ffn) 
+        return x
+
+class Encoder(nn.Module):
+    def __init__(self, d_model: int, layers: nn.ModuleList):
+        super().__init__()
+        self.layers = layers
+        self.norm = LayerNorm(d_model)
+
+    def forward(self, x, mask):
+        for layer in self.layers:
+            x = layer(x, mask)
+        return self.norm(x)
+
+class DecoderBlock(nn.Module):
+    def __init__(self, self_attention: MHA, cross_attention: MHA, ffn: FeedForward, d_model: int, dropout: float):
+        super().__init__()
+        self.attention = self_attention 
+        self.cross_attention = cross_attention
+        self.ffn = ffn
+        self.residual = nn.ModuleList([SkipConnection(d_model, dropout) for _ in range(3)]) 
+    
+    def forward(self, x, encoder_output, src_mask, trg_mask):
+        x = self.residual[0](x, lambda x: self.attention(x, x, x, trg_mask)) 
+        x = self.residual[1](x, lambda x: self.cross_attention(x, encoder_output, encoder_output, src_mask)) 
+        x = self.residual[2](x, self.ffn) 
+        return x
+
+class Decoder(nn.Module):
+    def __init__(self, d_model: int, layers: nn.ModuleList):
+        super().__init__()
+        self.layers = layers
+        self.norm = LayerNorm(d_model)
+
+    def forward(self, x, encoder_output, src_mask, trg_mask):
+        for layer in self.layers:
+            x = layer(x, encoder_output, src_mask, trg_mask)
+        return self.norm(x)
+
+class Output(nn.Module):
+    def __init__(self, d_model: int, vocab_size: int):
+        super().__init__()
+        self.proj = nn.Linear(d_model, vocab_size)
+
+    def forward(self, x):
+        # (Batch, Seq_Len, d_model) -> (Batch, Seq_Len, vocab_size)
+        return self.proj(x)
+
+class Transformer(nn.Module):
+    def __init__(self, encoder: Encoder, decoder: Decoder, src_embed: InputEmbedding, trg_embed: InputEmbedding, src_pos: PositionalEncoding, trg_pos: PositionalEncoding, output: Output):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.src_embed = src_embed
+        self.trg_embed = trg_embed
+        self.src_pos = src_pos
+        self.trg_pos = trg_pos
+        self.output_layer = output
+
+    def encode(self, src, src_mask):
+        src = self.src_embed(src)
+        src = self.src_pos(src)
+        return self.encoder(src, src_mask)
+
+    def decode(self, encoder_output, src_mask, trg, trg_mask):
+        trg = self.trg_embed(trg)
+        trg = self.trg_pos(trg)
+        return self.decoder(trg, encoder_output, src_mask, trg_mask)
+
+    def project(self, x):
+        return self.output_layer(x)
+
+def BuildTransformer(src_vocab_size: int, trg_vocab_size: int, src_seq_len: int, trg_seq_len: int, d_model: int = 512, N: int = 6, h: int = 8, dropout: float = 0.1, d_ff: int = 2048) -> Transformer:
+    # Create the embedding layers
+    src_embed = InputEmbedding(d_model, src_vocab_size)
+    trg_embed = InputEmbedding(d_model, trg_vocab_size)
+
+    # Create the positional encoding layers
+    src_pos = PositionalEncoding(d_model, src_seq_len, dropout)
+    trg_pos = PositionalEncoding(d_model, trg_seq_len, dropout)
+
+    # Create the encoder blocks
+    encoder_blocks = []
+    for _ in range(N):
+        encoder_self_attention = MHA(d_model, h, dropout)
+        ffn = FeedForward(d_model, d_ff, dropout)
+        encoder_block = EncoderBlock(encoder_self_attention, ffn, d_model, dropout)
+        encoder_blocks.append(encoder_block)
+
+    # Create the decoder blocks
+    decoder_blocks = []
+    for _ in range(N):
+        decoder_self_attention = MHA(d_model, h, dropout)
+        cross_attention = MHA(d_model, h, dropout)
+        ffn = FeedForward(d_model, d_ff, dropout)
+        decoder_block = DecoderBlock(decoder_self_attention, cross_attention, ffn, d_model, dropout)
+        decoder_blocks.append(decoder_block)
+    
+    # Create the encoder and decoder
+    encoder = Encoder(d_model, nn.ModuleList(encoder_blocks))
+    decoder = Decoder(d_model, nn.ModuleList(decoder_blocks))
+
+    # Create the projection layer
+    projection = Output(d_model, trg_vocab_size)
+
+    # Create the transformer
+    transformer = Transformer(encoder, decoder, src_embed, trg_embed, src_pos, trg_pos, projection)
+
+    # Initialize parameters
+    for p in transformer.parameters():
+        if p.dim() > 1:
+            nn.init.xavier_uniform_(p)
+    
     return transformer