Create app.py

app.py

@@ -0,0 +1,227 @@
+import torch
+import torch.nn as nn
+import math
+import json
+import sentencepiece as spm
+import gradio as gr
+
+# =========================
+# Load config
+# =========================
+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"]
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+# =========================
+# SentencePiece
+# =========================
+sp_en = spm.SentencePieceProcessor()
+sp_en.load("sp_en.model")
+
+sp_ar = spm.SentencePieceProcessor()
+sp_ar.load("sp_ar.model")
+
+
+# =========================
+# Model Classes (same as training)
+# =========================
+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
+
+        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)
+
+    def split(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):
+        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))
+
+        out = self.attention(q, k, v, mask)
+        return self.W_o(self.combine(out))
+
+
+class FFN(nn.Module):
+    def __init__(self, d_model, d_ff):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(d_model, d_ff),
+            nn.ReLU(),
+            nn.Linear(d_ff, d_model)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class PosEnc(nn.Module):
+    def __init__(self, d_model, max_len):
+        super().__init__()
+        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))
+
+        pe[:, 0::2] = torch.sin(pos * div)
+        pe[:, 1::2] = torch.cos(pos * div)
+
+        self.pe = pe.unsqueeze(0)
+
+    def forward(self, x):
+        return x + self.pe[:, :x.size(1)].to(x.device)
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, d_model, heads, d_ff):
+        super().__init__()
+        self.attn = MultiHeadAttention(d_model, heads)
+        self.ffn = FFN(d_model, d_ff)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+
+    def forward(self, x, mask):
+        x = self.norm1(x + self.attn(x, x, x, mask))
+        x = self.norm2(x + self.ffn(x))
+        return x
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, d_model, heads, d_ff):
+        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))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.d_model = config["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.pos = PosEnc(self.d_model, config["max_src_len"])
+
+        self.enc_layers = nn.ModuleList([
+            EncoderLayer(self.d_model, config["num_heads"], config["d_ff"])
+            for _ in range(config["num_layers"])
+        ])
+
+        self.dec_layers = nn.ModuleList([
+            DecoderLayer(self.d_model, config["num_heads"], config["d_ff"])
+            for _ in range(config["num_layers"])
+        ])
+
+        self.fc = nn.Linear(self.d_model, config["tgt_vocab_size"])
+
+    def masks(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
+
+    def forward(self, src, tgt):
+        src_mask, tgt_mask = self.masks(src, tgt)
+
+        src = self.pos(self.enc_emb(src))
+        tgt = self.pos(self.dec_emb(tgt))
+
+        enc = src
+        for layer in self.enc_layers:
+            enc = layer(enc, src_mask)
+
+        dec = tgt
+        for layer in self.dec_layers:
+            dec = layer(dec, enc, src_mask, tgt_mask)
+
+        return self.fc(dec)
+
+
+# =========================
+# Load model
+# =========================
+model = Transformer().to(device)
+model.load_state_dict(torch.load("best_model.pt", map_location=device))
+model.eval()
+
+
+# =========================
+# Inference
+# =========================
+def translate(sentence):
+
+    tokens = sp_en.encode(sentence)
+    tokens = [bos_id] + tokens + [eos_id]
+
+    src = torch.tensor(tokens).unsqueeze(0).to(device)
+
+    out = [bos_id]
+
+    for _ in range(50):
+
+        tgt = torch.tensor(out).unsqueeze(0).to(device)
+
+        with torch.no_grad():
+            pred = model(src, tgt)
+
+        next_token = pred[0, -1].argmax().item()
+
+        out.append(next_token)
+
+        if next_token == eos_id:
+            break
+
+    result = sp_ar.decode([t for t in out if t not in [bos_id, eos_id, pad_id]])
+    return result
+
+
+# =========================
+# UI
+# =========================
+demo = gr.Interface(
+    fn=translate,
+    inputs="text",
+    outputs="text",
+    title="Arabic ↔ English Translator (Transformer)",
+    description="Enter English sentence and get Arabic translation"
+)
+
+demo.launch()