import gradio as gr
from transformers import MarianMTModel, MarianTokenizer

model_names = {
    "English to Urdu": "Helsinki-NLP/opus-mt-en-ur",
    "Urdu to English": "Helsinki-NLP/opus-mt-ur-en"
}

models = {}

def load_model(direction):
    model_name = model_names[direction]
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    model = MarianMTModel.from_pretrained(model_name)
    return tokenizer, model

def translate(text, direction):
    try:
        if not text.strip():
            return "Please enter text to translate."
        tokenizer, model = load_model(direction)
        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
        output = model.generate(**inputs)
        return tokenizer.decode(output[0], skip_special_tokens=True)
    except Exception as e:
        return f"Error: {str(e)}"

iface = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(label="Enter Text", lines=4),
        gr.Radio(["English to Urdu", "Urdu to English"], label="Select Direction")
    ],
    outputs=gr.Textbox(label="Translated Text"),
    title="English ↔ Urdu Translator",
    description="Use this app to translate between English and Urdu using Hugging Face MarianMT models.",
    allow_flagging="never"
)

iface.launch()