import gradio as gr
import torch
from transformers import MarianMTModel, MarianTokenizer

# Ensure required libraries are installed
import os
os.system("pip install sentencepiece")

# Check if GPU is available and use it
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load models and tokenizers once (globally)
model_en_to_ur = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-en-ur").to(device)
tokenizer_en_to_ur = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ur")

model_ur_to_en = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-ur-en").to(device)
tokenizer_ur_to_en = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ur-en")

# Apply torch.compile() for optimization (if using PyTorch 2.0+)
if torch.__version__ >= "2.0":
    model_en_to_ur = torch.compile(model_en_to_ur)
    model_ur_to_en = torch.compile(model_ur_to_en)

# Function to translate text
def translate(text, direction):
    if not text.strip():
        return "Please enter some text to translate."

    if direction == "English to Urdu":
        tokenizer, model = tokenizer_en_to_ur, model_en_to_ur
    else:
        tokenizer, model = tokenizer_ur_to_en, model_ur_to_en

    # Tokenize input text (optimized padding)
    inputs = tokenizer(text, return_tensors="pt", padding="longest", truncation=True).to(device)

    # Generate translation
    with torch.no_grad():
        translated = model.generate(**inputs, max_length=512)

    # Decode output text
    output = tokenizer.decode(translated[0], skip_special_tokens=True)
    return output

# Gradio interface
interface = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(label="Enter Text"),
        gr.Radio(["English to Urdu", "Urdu to English"], label="Translation Direction", value="English to Urdu"),
    ],
    outputs=gr.Textbox(label="Translated Text"),
    title="⚡ Fast English ↔ Urdu Translator",
    description="Translate text between English and Urdu quickly using a neural machine translation model with GPU acceleration.",
)

# Launch the Gradio app
if __name__ == "__main__":
    interface.launch()