import gradio as gr
from transformers import MarianMTModel, MarianTokenizer

model_en_ur_name = "Helsinki-NLP/opus-mt-en-ur"
model_ur_en_name = "Helsinki-NLP/opus-mt-ur-en"

tokenizer_en_ur = MarianTokenizer.from_pretrained(model_en_ur_name)
model_en_ur = MarianMTModel.from_pretrained(model_en_ur_name)

tokenizer_ur_en = MarianTokenizer.from_pretrained(model_ur_en_name)
model_ur_en = MarianMTModel.from_pretrained(model_ur_en_name)

def translate_text(text, direction):
    if not text.strip():
        return "⚠️ Please enter some text."

    if direction == "English → Urdu":
        inputs = tokenizer_en_ur(text, return_tensors="pt", padding=True, truncation=True)
        translated = model_en_ur.generate(**inputs)
        return tokenizer_en_ur.decode(translated[0], skip_special_tokens=True)

    else:
        inputs = tokenizer_ur_en(text, return_tensors="pt", padding=True, truncation=True)
        translated = model_ur_en.generate(**inputs)
        return tokenizer_ur_en.decode(translated[0], skip_special_tokens=True)

with gr.Blocks() as app:
    gr.Markdown("# 🌐 English ↔ Urdu Translator")
    gr.Markdown("Built with Hugging Face + Gradio")

    text_input = gr.Textbox(lines=4, label="Input Text")
    direction = gr.Radio(["English → Urdu", "Urdu → English"], value="English → Urdu")
    output = gr.Textbox(label="Output")

    btn = gr.Button("Translate")
    btn.click(translate_text, [text_input, direction], output)

app.launch()