app.py

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)

padIndex = config["pad_id"]
BOSIndex = config["bos_id"]
EOSIndex = 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 (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

        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 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_heads(self, x):
        B, H, T, D = x.size()
        return x.transpose(1, 2).contiguous().view(B, T, self.d_model)

    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.scaled_dot_product_attention(Q, K, V, mask)
        return self.W_o(self.combine_heads(out))


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)
        )

    def forward(self, x):
        return self.net(x)


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)
        position = torch.arange(0, max_len).unsqueeze(1)

        div_term = torch.exp(torch.arange(0, d_model, 2) *
                             -(math.log(10000.0) / d_model))

        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):
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)


class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()

        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.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, num_heads, d_ff, dropout=0.1):
        super().__init__()

        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, src_vocab, tgt_vocab,
                 d_model=256, num_heads=4, num_layers=3,
                 d_ff=512, max_len=100):

        super().__init__()

        self.d_model = d_model

        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.positional_encoding = PositionalEncoding(d_model, max_len)

        self.encoder_layers = nn.ModuleList([
            EncoderLayer(d_model, num_heads, d_ff)
            for _ in range(num_layers)
        ])

        self.decoder_layers = nn.ModuleList([
            DecoderLayer(d_model, num_heads, d_ff)
            for _ in range(num_layers)
        ])

        self.fc = nn.Linear(d_model, tgt_vocab)

    def generate_mask(self, src, tgt):
        src_mask = (src != 0).unsqueeze(1).unsqueeze(2)

        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.generate_mask(src, 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.encoder_layers:
            enc = layer(enc, src_mask)

        dec = tgt
        for layer in self.decoder_layers:
            dec = layer(dec, enc, src_mask, tgt_mask)

        return self.fc(dec)


# =========================
# Load model
# =========================
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()


# =========================
# Translation
# =========================
def translate(text):

    src = sp_en.encode(text)
    src = [BOSIndex] + src + [EOSIndex]

    src = torch.tensor(src).unsqueeze(0).to(device)

    out = [BOSIndex]

    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 == EOSIndex:
            break

    result = sp_ar.decode([t for t in out if t not in [BOSIndex, EOSIndex, padIndex]])
    return result


# =========================
# UI
# =========================
gr.Interface(
    fn=translate,
    inputs="text",
    outputs="text",
    title="English ↔ Arabic Transformer",
).launch()