Upload 14 files

app.py

@@ -0,0 +1,31 @@
+import gradio as gr
+from src.inference import load_model_and_tokenizer, translate
+from src.ui import build_demo
+
+tokenizer_path = "tokenizer/bpe_tokenizer.json"
+model_checkpoint_path = "model/transformer_nmt_model_params.pt"
+
+model, tokenizer = load_model_and_tokenizer(tokenizer_path , model_checkpoint_path)
+
+def translate_fn(src_text, max_len):
+    return translate(model, tokenizer, [src_text], max_len=max_len, device=None)[0]
+
+inputs = [
+    gr.Textbox(label="📝 English Text", lines=3),
+    gr.Slider(10, 100, value=50, step=5, label="📏 Max Translated Length"),
+]
+
+outputs = [gr.Textbox(label="🌎 Spanish Translation", lines=5, interactive=False)]
+
+demo = build_demo(
+    translate_fn,
+    inputs,
+    outputs,
+    english_title = "# 🌐✨ TransformerTorch: Transformer-Based Neural Machine Translation 🚀",
+    persian_title = "# 🌐✨ مترجم هوشمند انگلیسی به اسپانیایی مبتنی بر معماری ترنسفورمر 🚀",
+    assets_dir = "assets",
+    app_title = "🌐 TransformerTorch 🌟"
+)
+
+if __name__ == "__main__":
+    demo.launch()

assets/css/custom.css

@@ -0,0 +1,79 @@
+/* Background */
+.gradio-container {
+    background: linear-gradient(135deg, #fdfbfb, #ebedee) !important;
+    font-family: 'Inter', 'Segoe UI', sans-serif !important;
+}
+.dark .gradio-container {
+    background: linear-gradient(135deg, #1e1a5e, #2a0a3a) !important;
+}
+
+/* Buttons */
+button {
+    border-radius: 14px !important;
+    padding: 10px 18px !important;
+    font-weight: 600 !important;
+    background: linear-gradient(90deg, #6a11cb, #2575fc) !important;
+    color: white !important;
+    box-shadow: 0 4px 10px rgba(0,0,0,0.15) !important;
+    transition: transform 0.15s ease-in-out;
+}
+button:hover {
+    transform: translateY(-2px);
+    box-shadow: 0 6px 14px rgba(0,0,0,0.25) !important;
+}
+
+/* Title */
+#title {
+    font-size: 2.8em !important;
+    font-weight: 700 !important;
+    color: #1e3a8a;
+    text-align: center;
+    margin-top: 28px;
+    margin-bottom: 12px;
+    text-shadow: 1px 2px 6px rgba(0,0,0,0.1);
+}
+.dark #title {
+    color: #e0f7fa !important;
+    text-shadow: 1px 2px 6px rgba(0,0,0,0.4);
+}
+
+/* Summary / Description */
+#summary {
+    color: #374151;
+    background: rgba(255,255,255,0.7);
+    padding: 18px;
+    border-radius: 16px;
+    box-shadow: 0 4px 12px rgba(0,0,0,0.08);
+    margin-bottom: 16px;
+    text-align: justify !important;
+}
+.dark #summary {
+    color: #d1d5db !important;
+    background: rgba(30, 30, 46, 0.6) !important;
+}
+
+/* Help / Info Box */
+#help_text {
+    color: #1f2937;
+    background: rgba(240, 249, 255, 0.9);
+    padding: 16px;
+    border-left: 5px solid #3b82f6;
+    border-radius: 14px;
+    box-shadow: 0 4px 10px rgba(0,0,0,0.05);
+    margin-top: 12px;
+    text-align: justify !important;
+}
+.dark #help_text {
+    color: #d1d5db !important;
+    background: rgba(30, 30, 46, 0.7) !important;
+    border-left: 5px solid #60a5fa !important;
+}
+
+/* RTL Support */
+.persian {
+    direction: rtl;
+    text-align: right;
+}
+#summary.persian, #help_text.persian {
+    text-align: justify !important;
+}

assets/markdown/english_summary.md

@@ -0,0 +1,9 @@
+**🌐 TransformerTorch** is a Transformer‑based Neural Machine Translation system trained on **220K English–Spanish sentence pairs**. It leverages advanced techniques to improve efficiency, including **Mixed Precision training, Weight Tying, a shared vocabulary and embedding space, and BPE tokenization**. The model uses a custom **greedy decoder** that computes the encoder memory once and then decodes autoregressively with causal and padding masks, reusing that memory at each step for efficient inference.
+
+Originally, for this project I implemented the Transformer architecture from scratch with PyTorch, which you can explore here: [GitHub – TransformerTorch](https://github.com/HooM4N/TransformerTorch)  
+
+**✍️ Improve Translation Quality**  
+To improve translation quality, include proper punctuation in the English source text:  
+- End **declarative sentences** with a period (`.`)  
+- End **questions** with a question mark (`?`)  
+- Use **exclamation marks** (`!`) where appropriate  

assets/markdown/persian_summary.md

@@ -0,0 +1,13 @@
+
+TransformerTorch یک سامانه‌ی ترجمه‌ی ماشینی انگلیسی به اسپانیایی است که با معماری ترنسفورمر در PyTorch پیاده‌سازی شده و روی **۲۲۰ هزار جمله‌ی موازی انگلیسی–اسپانیایی** آموزش دیده است. 
+در طراحی این مدل، علاوه بر ساختار اصلی ترنسفورمر، از تکنیک‌های بهینه‌سازی پیشرفته مانند آموزش با Mixed Precision، اشتراک گذاری پارامترهای امبدینگ و خروجی (Weight Tying)، امبدینگ و Vocabulary مشترک و توکانیزشن زیرواژه (BPE) برای افزایش سرعت آموزش و کاهش مصرف حافظه کارت گرافیگ بهره گرفته شده.
+
+فرایند استنتاج بصورت بهینه صورت میگیرد: در هنگام ترجمه خروجی انکودر یکبار محاسبه شده و در تمام گامهای زمانی تولید جمله اسپانیایی به دیکودر داده میشود ، رویکردی که باعث کاهش محاسبات تکراری و افزایش کارایی در زمان استنتاج شده است.   
+
+ضمن این پروژه معماری ترنسفومر از پایه با پایتورچ پیاده سازی شده است، برای اطلاعات بیشتر به گیتهاب پروژه مراجعه کنید: [GitHub – TransformerTorch](https://github.com/HooM4N/TransformerTorch)  
+
+**✍️ بهبود کیفیت ترجمه**  
+برای بهبود کیفیت ترجمه، متن انگلیسی ورودی را با نشانه‌گذاری درست بنویسید:  
+- جمله‌های خبری را با نقطه (`.`) پایان دهید  
+- جمله‌های پرسشی را با علامت سؤال (`?`) تمام کنید  
+- و در صورت نیاز از علامت تعجب (`!`) استفاده کنید  

model/transformer_nmt_model_params.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0dd9c9d916b68637804763c51787bc86a7b136494b80b32439ef75a75bae1244
+size 84552837

requirements.txt

@@ -0,0 +1,5 @@
+pandas
+tokenizers>=0.22.1
+torchmetrics>=1.8.2
+torch>=2.8.0
+gradio

src/config.py

@@ -0,0 +1,32 @@
+from dataclasses import dataclass, field
+
+@dataclass
+class HPARAMS:
+    vocab_size = 12500
+    max_seq_len = 32
+    batch_size = 128
+
+    model_hparams: dict = field(default_factory=lambda: {
+    "d_model" : 512,
+    "nhead" : 8,
+    "num_encoder_layers" : 2,
+    "num_decoder_layers" : 2,
+    "dim_feedforward" : 2048,
+    "dropout" : 0.1,
+    "padding_idx" : 0,
+    })
+
+    optimizer_hparams: dict = field(default_factory=lambda: {
+        "lr": 1e-3,
+        "weight_decay": 2e-5
+    })
+
+
+    trainer_hparams: dict = field(default_factory=lambda: {
+    "n_epochs": 20,
+    "enable_mixed_precision": True,
+    "restore_best_model" : False,
+    "use_early_stopping" : True,
+    "early_stopping_patience" : 3,
+    "grad_clip_value" : None
+    })

src/inference.py

@@ -0,0 +1,95 @@
+import torch, tokenizers
+import torch.nn as nn
+from .config import HPARAMS
+from .model import TransformerNMT
+
+@torch.no_grad()
+def greedy_decode(model, src_ids, pad_id, bos_id, eos_id, max_len, device):
+    """
+    Greedy decoding for Transformer model: computes encoder memory once, then
+    iteratively generates target tokens using prior decoder outputs and memory.
+    Supports batched inference, stops at EOS or max_len, and builds its own
+    padding and causal masks.
+    """
+    batch_size = src_ids.size(0)
+    model.eval()
+    src_ids = src_ids.to(device)
+    src_key_padding_mask = (src_ids == pad_id).to(device) # (N, S)
+
+    # compute encoder memory
+    src_emb = model.positional_embedding(model.shared_embedding(src_ids)) # (N, S, E)
+    memory = model.transformer.encoder(src = src_emb,
+                                       src_key_padding_mask = src_key_padding_mask) # (N, S, E)
+
+    # prepare initial decoder input
+    current_tokens = torch.full((batch_size, 1), bos_id, dtype=torch.long).to(device) # (N, 1)
+    finished = torch.zeros(batch_size, dtype=torch.bool).to(device)
+    outputs = [[] for _ in range(batch_size)]
+
+    # decoding
+    for step in range(max_len):
+        # target embedding & masks (causal/padding)
+        tgt_emb = model.positional_embedding(model.shared_embedding(current_tokens)).to(device) # (N, L, E)
+        tgt_key_padding_mask = (current_tokens == pad_id).to(device) # usually false (N ,L)
+        causal_mask = nn.Transformer.generate_square_subsequent_mask(tgt_emb.size(1), dtype=torch.bool).to(device) # (L, L)
+
+        # decoder outputs
+        decoder_outputs = model.transformer.decoder(tgt = tgt_emb, memory = memory, tgt_mask = causal_mask, 
+                                                    tgt_key_padding_mask = tgt_key_padding_mask,
+                                                    memory_key_padding_mask = src_key_padding_mask) # (N, L, E)
+
+        next_logits = model.output(decoder_outputs)[:, -1, :] # (N, vocab_size)
+        next_tokens = next_logits.argmax(dim=-1) # (N,)
+
+        # update current decoded tokens
+        current_tokens = torch.cat([current_tokens, next_tokens.unsqueeze(1)], dim=1) # (N, L+1)
+
+        # store output tokens & stop if EOS token found
+        for i in range(batch_size):
+            if not finished[i]:
+                outputs[i].append(int(next_tokens[i].item()))
+                if next_tokens[i] == eos_id:
+                    finished[i] = True
+
+        if finished.all():
+            break
+
+    return outputs
+
+def translate(model, tokenizer, src_list, max_len=64, device=None):
+    """
+    args:
+        src_list (List[str]): Source sentences to translate.
+        max_len (int): maximum length of generated output sequence.
+        device (torch.device, optional)
+    returns:
+        List[str]: translated target sentences.
+    """
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    pad_id, bos_id, eos_id = [tokenizer.token_to_id(i) for i in ["[PAD]", "[BOS]", "[EOS]"]]    
+    src_ids = torch.tensor([enc.ids for enc in tokenizer.encode_batch(src_list)],
+                               dtype=torch.long) # (N, S)
+
+    outputs = greedy_decode(model, src_ids, pad_id, bos_id, eos_id, max_len, device)
+    return tokenizer.decode_batch(outputs)
+
+def load_model_and_tokenizer(tokenizer_path, model_checkpoint_path, device=None):
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Torch Device: {device}")
+    hp = HPARAMS()
+
+    try:
+        tokenizer = tokenizers.Tokenizer.from_file(tokenizer_path)
+        tokenizer.enable_truncation(hp.max_seq_len)
+        tokenizer.enable_padding(pad_id = 0, pad_token = "[PAD]")
+        
+        model = TransformerNMT(tokenizer.get_vocab_size(), hp.max_seq_len, **hp.model_hparams).to(device)
+        state_dict = torch.load(model_checkpoint_path, map_location=device, weights_only=True)
+        model.load_state_dict(state_dict)
+        return model, tokenizer
+    except Exception as e:
+        print(f"Error loading model/tokenizer: {e}")
+        return None, None

src/model.py

@@ -0,0 +1,72 @@
+import torch
+import torch.nn as nn
+
+class PositionalEmbedding(nn.Module):
+    """
+    Positional Embedding
+
+    shapes:
+        N: batch size
+        L: seq len (max seq len of batch)
+        E: embedding dim
+        max_seq_len: max seq len across all samples
+
+    forward args:
+        X: batch of semantic embeddings (N, L, E)
+    """
+    def __init__(self, emb_dim, max_seq_len, dropout_p=0.1):
+        super().__init__()
+
+        # full embedding matrix with shape (maximum_sample_lenght, embedding_dim)
+        self.pos_embedding = nn.Parameter(torch.randn(max_seq_len, emb_dim) * 0.01)
+        self.dropout = nn.Dropout(dropout_p)
+
+    def forward(self, X):
+
+        # sliced for current batch max sequence lenght
+        emb_matrix = self.pos_embedding[:X.size(1)].unsqueeze(0)  # (1, L, E)         
+        return self.dropout(X + emb_matrix) # (N, L, E) 
+class TransformerNMT(nn.Module):
+    """
+    forward args:
+        src_ids: (N, S) token ids
+        tgt_ids: (N, L) token ids
+        src_key_padding_mask: (N, S) bool, True=PAD (ignored)
+        tgt_key_padding_mask: (N, L) bool, True=PAD (ignored)
+    """
+    def __init__(self, vocab_size, max_seq_len, d_model=512, nhead=4,
+                 num_encoder_layers=2, num_decoder_layers=2,
+                 dim_feedforward=2048, dropout=0.1, padding_idx=0):
+        super().__init__()
+
+        self.shared_embedding = nn.Embedding(vocab_size, d_model, padding_idx = padding_idx)
+        self.positional_embedding = PositionalEmbedding(d_model, max_seq_len)
+
+        self.transformer = nn.Transformer(d_model, nhead,
+                                          num_encoder_layers, num_decoder_layers,
+                                          dim_feedforward, dropout,
+                                          activation="relu", batch_first=True, 
+                                          norm_first=False, bias=True)
+
+        self.output = nn.Linear(d_model, vocab_size, bias=False)
+
+        # weight tying
+        self.output.weight = self.shared_embedding.weight
+
+    def forward(self, src_ids, tgt_ids, src_key_padding_mask, tgt_key_padding_mask):
+
+        src = self.positional_embedding(self.shared_embedding(src_ids)) # (N, S, E)
+        tgt = self.positional_embedding(self.shared_embedding(tgt_ids)) # (N, L, E)
+
+        # create target causal mask
+        L = tgt.size(1)
+        causal_mask = nn.Transformer.generate_square_subsequent_mask(L, dtype=torch.bool, device = tgt.device)
+
+        out = self.transformer(src = src , tgt = tgt,
+                               src_key_padding_mask = src_key_padding_mask,
+                               tgt_key_padding_mask = tgt_key_padding_mask,
+                               memory_key_padding_mask = src_key_padding_mask,
+                               tgt_mask = causal_mask
+                              ) # (N, L, E)
+
+        return self.output(out).transpose(-2,-1) # (N, vocab_size, L)

src/raw_data_builder.py

@@ -0,0 +1,17 @@
+from datasets import load_dataset
+import pandas as pd
+
+def prepare_dataset():
+    dataset = load_dataset("ageron/tatoeba_mt_train", "eng-spa")
+
+    df = pd.concat([
+        dataset["validation"].to_pandas(), dataset["test"].to_pandas()
+    ], axis=0)\
+           .sample(frac=1, random_state=42)\
+           .reset_index(drop=True)
+    
+    df[["source_text", "target_text"]].to_parquet("eng_spa.parquet")
+    print("Data saved to eng_spa.parquet")
+    
+if __name__ == "__main__":
+    prepare_dataset()

src/ui.py

@@ -0,0 +1,78 @@
+import os
+import gradio as gr
+
+def read_file(path: str, default_content: str = "") -> str:
+    """
+    Ensure file exists (with default_content if missing) and return its contents.
+    """
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    if not os.path.exists(path):
+        with open(path, "w", encoding="utf-8") as f:
+            f.write(default_content)
+    with open(path, "r", encoding="utf-8") as f:
+        return f.read()
+        
+def build_demo(
+    generation_fn,
+    inputs,
+    outputs,
+    english_title: str,
+    persian_title: str,
+    assets_dir: str = "assets",
+    app_title: str = "Demo"
+):
+    """
+    args:
+        generation_fn: callable for inference
+        inputs: list of Gradio input components
+        outputs: list of Gradio output components
+    """
+    md_dir = os.path.join(assets_dir, "markdown")
+    css_dir = os.path.join(assets_dir, "css")
+    english_md = os.path.join(md_dir, "english_summary.md")
+    persian_md = os.path.join(md_dir, "persian_summary.md")
+    english_summary = read_file(english_md)
+    persian_summary = read_file(persian_md)
+
+    css_file = os.path.join(css_dir, "custom.css")
+    css = read_file(css_file, "/* Custom CSS overrides */\n")
+
+    with gr.Blocks(css=css, title=app_title) as demo:
+        title_md = gr.Markdown(english_title, elem_id="title")
+
+        with gr.Row():
+            english_btn = gr.Button("English")
+            persian_btn = gr.Button("فارسی (Persian)")
+
+        summary_md = gr.Markdown(english_summary, elem_id="summary")
+
+        # generation panel
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1, variant="panel"):
+                for inp in inputs:
+                    inp.render()
+                generate_btn = gr.Button("✨ Translate", variant="primary")
+
+            with gr.Column(scale=1, variant="panel"):
+                for out in outputs:
+                    out.render()
+
+        # events
+        generate_btn.click(generation_fn, inputs=inputs, outputs=outputs)
+
+        def set_english():
+            return (
+                gr.update(value=english_title, elem_classes=[]),
+                gr.update(value=english_summary, elem_classes=[]),
+            )
+
+        def set_persian():
+            return (
+                gr.update(value=persian_title, elem_classes=["persian"]),
+                gr.update(value=persian_summary, elem_classes=["persian"]),
+            )
+
+        english_btn.click(set_english, outputs=[title_md, summary_md])
+        persian_btn.click(set_persian, outputs=[title_md, summary_md])
+
+    return demo

src/utils.py

@@ -0,0 +1,29 @@
+import matplotlib.pyplot as plt
+
+def plot_training_logs(train_logs):
+    fig, ax = plt.subplots(1, 3, figsize=(14, 4))
+
+    # Loss
+    ax[0].plot(train_logs['train_loss'], label="train")
+    ax[0].plot(train_logs['val_loss'], label="val")
+    ax[0].set_title("Loss")
+    ax[0].set_xlabel("Epoch")
+    ax[0].set_ylabel("Loss")
+    ax[0].legend()
+    ax[0].grid(True)
+
+    # Validation metric
+    ax[1].plot(train_logs['val_metric'], label="val metric", color="tab:orange")
+    ax[1].set_title("Validation Metric")
+    ax[1].set_xlabel("Epoch")
+    ax[1].set_ylabel("Metric")
+    ax[1].grid(True)
+
+    # Learning rate
+    ax[2].plot(train_logs['lr'], label="lr", color="tab:green")
+    ax[2].set_title("Learning Rate")
+    ax[2].set_xlabel("Epoch")
+    ax[2].set_ylabel("LR")
+    ax[2].grid(True)
+
+    plt.tight_layout();