Update app.py

app.py

@@ -11,9 +11,9 @@ import gradio as gr
 with open("config.json") as f:
     config = json.load(f)
 
-pad_id = config["pad_id"]
-bos_id = config["bos_id"]
-eos_id = config["eos_id"]
+padIndex = config["pad_id"]
+BOSIndex = config["bos_id"]
+EOSIndex = config["eos_id"]
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -29,12 +29,14 @@ sp_ar.load("sp_ar.model")
 
 
 # =========================
-# Model Classes (same as training)
+# MODEL (EXACT TRAINING VERSION)
 # =========================
+
 class MultiHeadAttention(nn.Module):
     def __init__(self, d_model, num_heads):
         super().__init__()
         assert d_model % num_heads == 0
+
         self.d_model = d_model
         self.num_heads = num_heads
         self.d_k = d_model // num_heads
@@ -44,35 +46,39 @@ class MultiHeadAttention(nn.Module):
         self.W_v = nn.Linear(d_model, d_model)
         self.W_o = nn.Linear(d_model, d_model)
 
-    def split(self, x):
+    def scaled_dot_product_attention(self, Q, K, V, mask=None):
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
+
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, -1e9)
+
+        attn = torch.softmax(scores, dim=-1)
+        return torch.matmul(attn, V)
+
+    def split_heads(self, x):
         B, T, D = x.size()
         return x.view(B, T, self.num_heads, self.d_k).transpose(1, 2)
 
-    def combine(self, x):
+    def combine_heads(self, x):
         B, H, T, D = x.size()
         return x.transpose(1, 2).contiguous().view(B, T, self.d_model)
 
-    def attention(self, q, k, v, mask=None):
-        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
-        if mask is not None:
-            scores = scores.masked_fill(mask == 0, -1e9)
-        return torch.softmax(scores, dim=-1) @ v
-
-    def forward(self, q, k, v, mask=None):
-        q = self.split(self.W_q(q))
-        k = self.split(self.W_k(k))
-        v = self.split(self.W_v(v))
+    def forward(self, Q, K, V, mask=None):
+        Q = self.split_heads(self.W_q(Q))
+        K = self.split_heads(self.W_k(K))
+        V = self.split_heads(self.W_v(V))
 
-        out = self.attention(q, k, v, mask)
-        return self.W_o(self.combine(out))
+        out = self.scaled_dot_product_attention(Q, K, V, mask)
+        return self.W_o(self.combine_heads(out))
 
 
-class FFN(nn.Module):
-    def __init__(self, d_model, d_ff):
+class PositionWiseFeedForward(nn.Module):
+    def __init__(self, d_model, d_ff, dropout=0.1):
         super().__init__()
         self.net = nn.Sequential(
             nn.Linear(d_model, d_ff),
             nn.ReLU(),
+            nn.Dropout(dropout),
             nn.Linear(d_ff, d_model)
         )
 
@@ -80,96 +86,117 @@ class FFN(nn.Module):
         return self.net(x)
 
 
-class PosEnc(nn.Module):
-    def __init__(self, d_model, max_len):
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_len, dropout=0.1):
         super().__init__()
+
+        self.dropout = nn.Dropout(dropout)
+
         pe = torch.zeros(max_len, d_model)
-        pos = torch.arange(0, max_len).unsqueeze(1)
-        div = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
+        position = torch.arange(0, max_len).unsqueeze(1)
 
-        pe[:, 0::2] = torch.sin(pos * div)
-        pe[:, 1::2] = torch.cos(pos * div)
+        div_term = torch.exp(torch.arange(0, d_model, 2) *
+                             -(math.log(10000.0) / d_model))
 
-        self.pe = pe.unsqueeze(0)
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+
+        self.register_buffer("pe", pe.unsqueeze(0))
 
     def forward(self, x):
-        return x + self.pe[:, :x.size(1)].to(x.device)
+        x = x + self.pe[:, :x.size(1)]
+        return self.dropout(x)
 
 
 class EncoderLayer(nn.Module):
-    def __init__(self, d_model, heads, d_ff):
+    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
         super().__init__()
-        self.attn = MultiHeadAttention(d_model, heads)
-        self.ffn = FFN(d_model, d_ff)
+
+        self.self_attn = MultiHeadAttention(d_model, num_heads)
+        self.feed_forward = PositionWiseFeedForward(d_model, d_ff, dropout)
+
         self.norm1 = nn.LayerNorm(d_model)
         self.norm2 = nn.LayerNorm(d_model)
 
+        self.dropout = nn.Dropout(dropout)
+
     def forward(self, x, mask):
-        x = self.norm1(x + self.attn(x, x, x, mask))
-        x = self.norm2(x + self.ffn(x))
+        x = self.norm1(x + self.dropout(self.self_attn(x, x, x, mask)))
+        x = self.norm2(x + self.dropout(self.feed_forward(x)))
         return x
 
 
 class DecoderLayer(nn.Module):
-    def __init__(self, d_model, heads, d_ff):
+    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
         super().__init__()
-        self.self_attn = MultiHeadAttention(d_model, heads)
-        self.cross_attn = MultiHeadAttention(d_model, heads)
-        self.ffn = FFN(d_model, d_ff)
-
-        self.n1 = nn.LayerNorm(d_model)
-        self.n2 = nn.LayerNorm(d_model)
-        self.n3 = nn.LayerNorm(d_model)
-
-    def forward(self, x, enc, src_mask, tgt_mask):
-        x = self.n1(x + self.self_attn(x, x, x, tgt_mask))
-        x = self.n2(x + self.cross_attn(x, enc, enc, src_mask))
-        x = self.n3(x + self.ffn(x))
+
+        self.self_attn = MultiHeadAttention(d_model, num_heads)
+        self.cross_attn = MultiHeadAttention(d_model, num_heads)
+        self.feed_forward = PositionWiseFeedForward(d_model, d_ff, dropout)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, enc_out, src_mask, tgt_mask):
+        x = self.norm1(x + self.dropout(self.self_attn(x, x, x, tgt_mask)))
+        x = self.norm2(x + self.dropout(self.cross_attn(x, enc_out, enc_out, src_mask)))
+        x = self.norm3(x + self.dropout(self.feed_forward(x)))
         return x
 
 
 class Transformer(nn.Module):
-    def __init__(self):
+    def __init__(self, src_vocab, tgt_vocab,
+                 d_model=256, num_heads=4, num_layers=3,
+                 d_ff=512, max_len=100):
+
         super().__init__()
 
-        self.d_model = config["d_model"]
+        self.d_model = d_model
 
-        self.enc_emb = nn.Embedding(config["src_vocab_size"], self.d_model, padding_idx=0)
-        self.dec_emb = nn.Embedding(config["tgt_vocab_size"], self.d_model, padding_idx=0)
+        self.encoder_embedding = nn.Embedding(src_vocab, d_model, padding_idx=0)
+        self.decoder_embedding = nn.Embedding(tgt_vocab, d_model, padding_idx=0)
 
-        self.pos = PosEnc(self.d_model, config["max_src_len"])
+        self.positional_encoding = PositionalEncoding(d_model, max_len)
 
-        self.enc_layers = nn.ModuleList([
-            EncoderLayer(self.d_model, config["num_heads"], config["d_ff"])
-            for _ in range(config["num_layers"])
+        self.encoder_layers = nn.ModuleList([
+            EncoderLayer(d_model, num_heads, d_ff)
+            for _ in range(num_layers)
         ])
 
-        self.dec_layers = nn.ModuleList([
-            DecoderLayer(self.d_model, config["num_heads"], config["d_ff"])
-            for _ in range(config["num_layers"])
+        self.decoder_layers = nn.ModuleList([
+            DecoderLayer(d_model, num_heads, d_ff)
+            for _ in range(num_layers)
         ])
 
-        self.fc = nn.Linear(self.d_model, config["tgt_vocab_size"])
+        self.fc = nn.Linear(d_model, tgt_vocab)
 
-    def masks(self, src, tgt):
+    def generate_mask(self, src, tgt):
         src_mask = (src != 0).unsqueeze(1).unsqueeze(2)
-        tgt_mask = (tgt != 0).unsqueeze(1).unsqueeze(3)
-        size = tgt.size(1)
-        causal = torch.tril(torch.ones(size, size)).bool().to(tgt.device)
-        return src_mask, tgt_mask & causal
+
+        tgt_pad = (tgt != 0).unsqueeze(1).unsqueeze(3)
+        T = tgt.size(1)
+
+        causal = torch.tril(torch.ones(T, T)).bool().to(tgt.device)
+
+        tgt_mask = tgt_pad & causal
+
+        return src_mask, tgt_mask
 
     def forward(self, src, tgt):
-        src_mask, tgt_mask = self.masks(src, tgt)
+        src_mask, tgt_mask = self.generate_mask(src, tgt)
 
-        src = self.pos(self.enc_emb(src))
-        tgt = self.pos(self.dec_emb(tgt))
+        src = self.positional_encoding(self.encoder_embedding(src) * math.sqrt(self.d_model))
+        tgt = self.positional_encoding(self.decoder_embedding(tgt) * math.sqrt(self.d_model))
 
         enc = src
-        for layer in self.enc_layers:
+        for layer in self.encoder_layers:
             enc = layer(enc, src_mask)
 
         dec = tgt
-        for layer in self.dec_layers:
+        for layer in self.decoder_layers:
             dec = layer(dec, enc, src_mask, tgt_mask)
 
         return self.fc(dec)
@@ -178,22 +205,31 @@ class Transformer(nn.Module):
 # =========================
 # Load model
 # =========================
-model = Transformer().to(device)
+model = Transformer(
+    config["src_vocab_size"],
+    config["tgt_vocab_size"],
+    config["d_model"],
+    config["num_heads"],
+    config["num_layers"],
+    config["d_ff"],
+    max_len=max(config["max_src_len"], config["max_tgt_len"])
+).to(device)
+
 model.load_state_dict(torch.load("best_model.pt", map_location=device))
 model.eval()
 
 
 # =========================
-# Inference
+# Translation
 # =========================
-def translate(sentence):
+def translate(text):
 
-    tokens = sp_en.encode(sentence)
-    tokens = [bos_id] + tokens + [eos_id]
+    src = sp_en.encode(text)
+    src = [BOSIndex] + src + [EOSIndex]
 
-    src = torch.tensor(tokens).unsqueeze(0).to(device)
+    src = torch.tensor(src).unsqueeze(0).to(device)
 
-    out = [bos_id]
+    out = [BOSIndex]
 
     for _ in range(50):
 
@@ -203,25 +239,21 @@ def translate(sentence):
             pred = model(src, tgt)
 
         next_token = pred[0, -1].argmax().item()
-
         out.append(next_token)
 
-        if next_token == eos_id:
+        if next_token == EOSIndex:
             break
 
-    result = sp_ar.decode([t for t in out if t not in [bos_id, eos_id, pad_id]])
+    result = sp_ar.decode([t for t in out if t not in [BOSIndex, EOSIndex, padIndex]])
     return result
 
 
 # =========================
 # UI
 # =========================
-demo = gr.Interface(
+gr.Interface(
     fn=translate,
     inputs="text",
     outputs="text",
-    title="Arabic ↔ English Translator (Transformer)",
-    description="Enter English sentence and get Arabic translation"
-)
-
-demo.launch()
+    title="English ↔ Arabic Transformer",
+).launch()