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

# Define model names
en_ur_model_name = "Helsinki-NLP/opus-mt-en-ur"
ur_en_model_name = "Helsinki-NLP/opus-mt-ur-en"

# Load tokenizers and models
en_ur_tokenizer = MarianTokenizer.from_pretrained(en_ur_model_name)
en_ur_model = MarianMTModel.from_pretrained(en_ur_model_name)

ur_en_tokenizer = MarianTokenizer.from_pretrained(ur_en_model_name)
ur_en_model = MarianMTModel.from_pretrained(ur_en_model_name)

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

    tokenizer, model = (en_ur_tokenizer, en_ur_model) if direction == "English to Urdu" else (ur_en_tokenizer, ur_en_model)

    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        translated = model.generate(**inputs)
    output = tokenizer.decode(translated[0], skip_special_tokens=True)
    return output

# Gradio UI
iface = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(lines=4, label="Enter Text"),
        gr.Radio(["English to Urdu", "Urdu to English"], label="Translation Direction")
    ],
    outputs=gr.Textbox(label="Translated Text"),
    title="English ↔ Urdu Translator",
    description="Translate text between English and Urdu using Hugging Face MarianMT models."
)

iface.launch()