Update app/model_def.py

app/model_def.py

@@ -2,6 +2,7 @@ import torch
 import torch.nn as nn
 import math
 
+
 class InputEmbedding(nn.Module):
     def __init__(self, d_model: int, vocab_size: int):
         super().__init__()
@@ -12,6 +13,7 @@ class InputEmbedding(nn.Module):
     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__()
@@ -20,16 +22,19 @@ class PositionalEncoding(nn.Module):
         self.dropout = nn.Dropout(dropout)
         pe = torch.zeros(seq_len, d_model)
         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))
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
+        )
         pe[:, 0::2] = torch.sin(position * div_term)
         pe[:, 1::2] = torch.cos(position * div_term)
         pe = pe.unsqueeze(0)
-        self.register_buffer('pe', pe)
+        self.register_buffer("pe", pe)
 
     def forward(self, x):
-        x = x + (self.pe[:, :x.shape[1], :]).requires_grad_(False)
+        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__()
@@ -42,6 +47,7 @@ class LayerNorm(nn.Module):
         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__()
@@ -52,6 +58,7 @@ class FeedForward(nn.Module):
     def forward(self, x):
         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__()
@@ -80,13 +87,18 @@ class MHA(nn.Module):
         query = self.w_q(q)
         key = self.w_k(k)
         value = self.w_v(v)
-        query = query.view(query.shape[0], query.shape[1], self.h, self.d_k).transpose(1, 2)
+        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)
+        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)
         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__()
@@ -96,20 +108,26 @@ class SkipConnection(nn.Module):
     def forward(self, x, sublayer):
         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):
+    def __init__(
+        self, self_attention: MHA, ffn: FeedForward, d_model: int, dropout: float
+    ):
         super().__init__()
         # Name required by the saved model file
         self.attention = self_attention
         self.ffn = ffn
         # Name required by the saved model file
-        self.residual = nn.ModuleList([SkipConnection(d_model, dropout) for _ in range(2)])
+        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__()
@@ -121,22 +139,36 @@ class Encoder(nn.Module):
             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):
+    def __init__(
+        self,
+        self_attention: MHA,
+        cross_attention: MHA,
+        ffn: FeedForward,
+        d_model: int,
+        dropout: float,
+    ):
         super().__init__()
         # Name required by the saved model file
         self.self_attention = self_attention
         self.cross_attention = cross_attention
         self.ffn = ffn
         # Name required by the saved model file
-        self.residual = nn.ModuleList([SkipConnection(d_model, dropout) for _ in range(3)])
+        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.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[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__()
@@ -148,6 +180,7 @@ class Decoder(nn.Module):
             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__()
@@ -156,8 +189,18 @@ class Output(nn.Module):
     def forward(self, x):
         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):
+    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
@@ -188,8 +231,6 @@ class Transformer(nn.Module):
         for block in self.decoder.layers:
             attn_scores = block.cross_attention.attention_scores
             if attn_scores is not None:
-                # Shape: [batch, heads, tgt_len, src_len]
-                # Squeeze batch, transpose to [heads, src_len, tgt_len]
                 weights.append(attn_scores.squeeze(0).transpose(-2, -1))
         return weights
 
@@ -216,7 +257,18 @@ class Transformer(nn.Module):
             "cross_attn": cross_weights,
         }
 
-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:
+
+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:
     src_embed = InputEmbedding(d_model, src_vocab_size)
     trg_embed = InputEmbedding(d_model, trg_vocab_size)
     src_pos = PositionalEncoding(d_model, src_seq_len, dropout)
@@ -232,13 +284,17 @@ def BuildTransformer(src_vocab_size: int, trg_vocab_size: int, src_seq_len: int,
         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_block = DecoderBlock(
+            decoder_self_attention, cross_attention, ffn, d_model, dropout
+        )
         decoder_blocks.append(decoder_block)
     encoder = Encoder(d_model, nn.ModuleList(encoder_blocks))
     decoder = Decoder(d_model, nn.ModuleList(decoder_blocks))
     projection = Output(d_model, trg_vocab_size)
-    transformer = Transformer(encoder, decoder, src_embed, trg_embed, src_pos, trg_pos, projection)
+    transformer = Transformer(
+        encoder, decoder, src_embed, trg_embed, src_pos, trg_pos, projection
+    )
     for p in transformer.parameters():
         if p.dim() > 1:
             nn.init.xavier_uniform_(p)
-    return transformer
+    return transformer