Upload 6 files

app_2.py

@@ -0,0 +1,130 @@
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
+
+# Load the model and tokenizer
+model_name = "jbochi/madlad400-3b-mt"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16,
+    device_map="auto"
+)
+
+def translate_text(text, source_lang, target_lang):
+    """
+    Translate text between English and Persian using MADLAD-400-3B
+    """
+    # Define language codes for the model
+    lang_codes = {
+        "English": "en",
+        "Persian": "fa"
+    }
+    
+    source_code = lang_codes[source_lang]
+    target_code = lang_codes[target_lang]
+    
+    # Create the translation prompt in the format the model expects
+    prompt = f"<2{target_code}> {text}"
+    
+    try:
+        # Tokenize input
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+        
+        # Move inputs to the same device as model
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
+        
+        # Generate translation
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_length=512,
+                num_beams=5,
+                early_stopping=True,
+                no_repeat_ngram_size=3,
+                length_penalty=1.0
+            )
+        
+        # Decode the output
+        translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        return translated_text
+    
+    except Exception as e:
+        return f"Error during translation: {str(e)}"
+
+# Create the Gradio interface
+with gr.Blocks(title="English-Persian Translator") as demo:
+    gr.Markdown(
+        """
+        # 🌍 English-Persian Translator
+        **Powered by MADLAD-400-3B Model**
+        
+        Translate text between English and Persian using the state-of-the-art MADLAD-400 model.
+        """
+    )
+    
+    with gr.Row():
+        with gr.Column():
+            source_lang = gr.Dropdown(
+                choices=["English", "Persian"],
+                value="English",
+                label="Source Language"
+            )
+            input_text = gr.Textbox(
+                lines=5,
+                placeholder="Enter text to translate...",
+                label="Input Text"
+            )
+            translate_btn = gr.Button("Translate", variant="primary")
+        
+        with gr.Column():
+            target_lang = gr.Dropdown(
+                choices=["Persian", "English"],
+                value="Persian",
+                label="Target Language"
+            )
+            output_text = gr.Textbox(
+                lines=5,
+                label="Translated Text",
+                interactive=False
+            )
+    
+    # Examples
+    gr.Examples(
+        examples=[
+            ["Hello, how are you today?", "English", "Persian"],
+            ["What is your name?", "English", "Persian"],
+            ["سلام، حالتون چطوره؟", "Persian", "English"],
+            ["امروز هوا خوب است", "Persian", "English"]
+        ],
+        inputs=[input_text, source_lang, target_lang],
+        outputs=output_text,
+        fn=translate_text,
+        cache_examples=False
+    )
+    
+    # Connect the button
+    translate_btn.click(
+        fn=translate_text,
+        inputs=[input_text, source_lang, target_lang],
+        outputs=output_text
+    )
+    
+    # Auto-update target language based on source selection
+    def update_target_lang(source_lang):
+        return "Persian" if source_lang == "English" else "English"
+    
+    source_lang.change(
+        fn=update_target_lang,
+        inputs=source_lang,
+        outputs=target_lang
+    )
+
+if __name__ == "__main__":
+    # Launch the app
+    demo.launch(
+        server_name="0.0.0.0",  # Allow external access
+        share=False,  # Set to True to get a public URL
+        debug=True
+    )

config.py

@@ -0,0 +1,42 @@
+import torch
+class Config:
+    # Data
+    dataset_name = "ParsBench/parsinlu-machine-translation-fa-en-alpaca-style"
+    source_lang = "instruction"   # English
+    target_lang = "output"        # Persian
+    max_length = 32
+    batch_size = 24
+
+    # Model
+    input_dim = 5000             # Vocabulary size for English
+    output_dim = 5000            # Vocabulary size for Persian 
+    embedding_dim = 64           # Word vector dimensions
+    hidden_dim = 128              # LSTM hidden state size
+    num_layers = 1                # Stacked LSTM layers
+    dropout = 0.1                 # Regularization to prevent overfitting
+
+    # Training
+    learning_rate = 0.001
+    num_epochs = 5
+    teacher_forcing_ratio = 0.7   # Mix of ground truth vs model predictions
+
+
+    # Optimization
+    gradient_accumulation_steps = 1
+    use_amp = True                 # Mixed precision for speed
+    use_gradient_clipping = True
+    max_grad_norm = 1.0
+
+
+    # Checkpoint Configuration =====
+    checkpoint_interval = 1        # Save every 2 epochs
+    save_best_only = True          # Only save when model improves
+
+    # Device
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    # Paths
+    model_save_path = "models/seq2seq_model.pth"
+    tokenizer_save_path = "models/tokenizers/"
+    checkpoint_path = "models/checkpoint.pth"
+    best_model_path = "models/best_model.pth"

inference.py

@@ -0,0 +1,37 @@
+import torch
+from config import Config
+from models.seq2seq import Encoder, Decoder, Seq2Seq
+from utils.tokenizer import build_vocab
+from datasets import load_from_disk
+
+def translate_sentence(sentence, model, src_tokenizer, src_vocab, trg_vocab, device, max_len=30):
+    model.eval()
+    tokens = src_tokenizer(sentence.lower())
+    src_tensor = torch.tensor([src_vocab["<sos>"]] + [src_vocab[t] for t in tokens] + [src_vocab["<eos>"]]).unsqueeze(1).to(device)
+    with torch.no_grad():
+        hidden = model.encoder(src_tensor)
+    trg_indexes = [trg_vocab["<sos>"]]
+
+    for _ in range(max_len):
+        trg_tensor = torch.tensor([trg_indexes[-1]]).to(device)
+        with torch.no_grad():
+            output, hidden = model.decoder(trg_tensor, hidden)
+        pred_token = output.argmax(1).item()
+        trg_indexes.append(pred_token)
+        if pred_token == trg_vocab["<eos>"]:
+            break
+    return [trg_vocab.get_itos()[i] for i in trg_indexes]
+
+if __name__ == "__main__":
+    cfg = Config()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    dataset = load_from_disk("data/raw/")
+    src_tokenizer, src_vocab = build_vocab(dataset, cfg.source_lang)
+    trg_tokenizer, trg_vocab = build_vocab(dataset, cfg.target_lang)
+
+    enc = Encoder(len(src_vocab), cfg.emb_dim, cfg.hid_dim, cfg.n_layers)
+    dec = Decoder(len(trg_vocab), cfg.emb_dim, cfg.hid_dim, cfg.n_layers)
+    model = Seq2Seq(enc, dec, device).to(device)
+    model.load_state_dict(torch.load(cfg.model_save_path, map_location=device))
+
+    print(translate_sentence("I love cats", model, src_tokenizer, src_vocab, trg_vocab, device))

requirements.txt

@@ -0,0 +1,6 @@
+torch>=1.9.0
+torchtext>=0.10.0
+datasets>=2.14.0
+numpy>=1.21.0
+tqdm>=4.62.0
+streamlit>=1.22.0

resume_training.py

@@ -0,0 +1,14 @@
+import torch
+from train import main
+import os
+
+if __name__ == "__main__":
+    print("🔄 Resuming training from checkpoint...")
+    
+    # Check if checkpoint exists
+    if not os.path.exists("models/checkpoint.pth"):
+        print("❌ No checkpoint found. Starting fresh training...")
+    else:
+        print("✅ Checkpoint found. Resuming...")
+    
+    main()

train.py

@@ -0,0 +1,146 @@
+import torch
+import torch.optim as optim
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from datasets import load_from_disk
+import os
+from config import Config
+from utils.tokenizer import build_vocab
+from utils.preprocessing import collate_fn
+from models.seq2seq import Encoder, Decoder, Seq2Seq
+from tqdm import tqdm
+
+def save_checkpoint(epoch, model, optimizer, scaler, loss, path):
+    """Save training checkpoint"""
+    checkpoint = {
+        'epoch': epoch,
+        'model_state_dict': model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+        'scaler_state_dict': scaler.state_dict(),
+        'loss': loss,
+    }
+    torch.save(checkpoint, path)
+    print(f"✅ Checkpoint saved at epoch {epoch}")
+
+def load_checkpoint(model, optimizer, scaler, path, device):
+    """Load training checkpoint"""
+    if os.path.exists(path):
+        checkpoint = torch.load(path, map_location=device)
+        model.load_state_dict(checkpoint['model_state_dict'])
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        scaler.load_state_dict(checkpoint['scaler_state_dict'])
+        start_epoch = checkpoint['epoch'] + 1
+        best_loss = checkpoint['loss']
+        print(f"✅ Checkpoint loaded. Resuming from epoch {start_epoch}")
+        return start_epoch, best_loss
+    return 0, float('inf')  # Start from beginning if no checkpoint
+
+def train_one_epoch(model, dataloader, optimizer, criterion, device, scaler, epoch, cfg):
+    model.train()
+    total_loss = 0
+    optimizer.zero_grad()  # Zero gradients at start
+    
+    loop = tqdm(dataloader, desc=f"Epoch {epoch+1}", leave=False)
+
+    for batch_idx, (src, trg) in enumerate(loop):
+        src, trg = src.to(device), trg.to(device)
+
+        # Mixed precision training
+        with torch.cuda.amp.autocast(enabled=cfg.use_amp):
+            output = model(src, trg)
+            output_dim = output.shape[-1]
+            output = output[1:].reshape(-1, output_dim)
+            trg = trg[1:].reshape(-1)
+            loss = criterion(output, trg) / cfg.gradient_accumulation_steps  # Normalize loss
+
+        scaler.scale(loss).backward()
+
+        # Gradient accumulation
+        if (batch_idx + 1) % cfg.gradient_accumulation_steps == 0:
+            if cfg.use_gradient_clipping:
+                scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.max_grad_norm)
+            
+            scaler.step(optimizer)
+            scaler.update()
+            optimizer.zero_grad()
+
+        total_loss += loss.item() * cfg.gradient_accumulation_steps
+        loop.set_postfix(loss=loss.item() * cfg.gradient_accumulation_steps)
+
+    return total_loss / len(dataloader)
+
+
+
+def main():
+    cfg = Config()
+    device = cfg.device
+    print(f"Using device: {device}")
+
+    # Create directories if they don't exist
+    os.makedirs("models", exist_ok=True)
+    os.makedirs("models/tokenizers", exist_ok=True)
+
+    # Load dataset (full dataset)
+    dataset = load_from_disk("data/raw/")
+
+    # Build vocab using full dataset
+    src_tokenizer, src_vocab = build_vocab(dataset, cfg.source_lang)
+    trg_tokenizer, trg_vocab = build_vocab(dataset, cfg.target_lang)
+
+    # Save tokenizers and vocab for future use
+    torch.save({
+        'src_tokenizer': src_tokenizer,
+        'src_vocab': src_vocab,
+        'trg_tokenizer': trg_tokenizer,
+        'trg_vocab': trg_vocab
+    }, cfg.tokenizer_save_path + "tokenizers.pth")
+
+    # DataLoader with train split
+    collate = lambda batch: collate_fn(
+        batch, src_tokenizer, trg_tokenizer, src_vocab, trg_vocab, cfg.max_length,
+        src_lang=cfg.source_lang, trg_lang=cfg.target_lang
+    )
+    dataloader = DataLoader(dataset["train"], batch_size=cfg.batch_size, collate_fn=collate, shuffle=True)
+
+    # Model
+    enc = Encoder(len(src_vocab), cfg.embedding_dim, cfg.hidden_dim, cfg.num_layers)
+    dec = Decoder(len(trg_vocab), cfg.embedding_dim, cfg.hidden_dim, cfg.num_layers)
+    model = Seq2Seq(enc, dec, device).to(device)
+
+    optimizer = optim.Adam(model.parameters(), lr=cfg.learning_rate)
+    criterion = nn.CrossEntropyLoss(ignore_index=src_vocab["<pad>"])
+    scaler = torch.cuda.amp.GradScaler()
+
+    # Try to load checkpoint
+    start_epoch, best_loss = load_checkpoint(model, optimizer, scaler, cfg.checkpoint_path, device)
+
+    for epoch in range(start_epoch, cfg.num_epochs):
+        print(f"\nEpoch {epoch+1}/{cfg.num_epochs}")
+        
+        try:
+            loss = train_one_epoch(model, dataloader, optimizer, criterion, device, scaler, epoch, cfg)
+            print(f"Epoch {epoch+1}/{cfg.num_epochs} | Loss: {loss:.3f}")
+
+            # Save checkpoint after each epoch
+            save_checkpoint(epoch, model, optimizer, scaler, loss, cfg.checkpoint_path)
+
+            # Save best model
+            if loss < best_loss:
+                best_loss = loss
+                torch.save(model.state_dict(), cfg.best_model_path)
+                print(f"🎉 New best model saved with loss: {loss:.3f}")
+
+        except RuntimeError as e:
+            if "CUDA out of memory" in str(e):
+                print("⚠️ GPU out of memory. Saving checkpoint and exiting...")
+                save_checkpoint(epoch, model, optimizer, scaler, loss, cfg.checkpoint_path)
+                print("✅ Checkpoint saved. You can resume training later.")
+                break
+            else:
+                raise e
+
+    print("✅ Training completed!")
+
+if __name__ == "__main__":
+    main()