import gradio as gr
from transformers import MarianMTModel, MarianTokenizer

# Load models for both directions
en_to_ur_model_name = "Helsinki-NLP/opus-mt-en-ur"
ur_to_en_model_name = "Helsinki-NLP/opus-mt-ur-en"

en_to_ur_tokenizer = MarianTokenizer.from_pretrained(en_to_ur_model_name)
en_to_ur_model = MarianMTModel.from_pretrained(en_to_ur_model_name)

ur_to_en_tokenizer = MarianTokenizer.from_pretrained(ur_to_en_model_name)
ur_to_en_model = MarianMTModel.from_pretrained(ur_to_en_model_name)

# Translation function
def translate(text, direction):
    if direction == "English to Urdu":
        tokenizer = en_to_ur_tokenizer
        model = en_to_ur_model
    else:
        tokenizer = ur_to_en_tokenizer
        model = ur_to_en_model

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

# Gradio interface
demo = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(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 between English and Urdu using Hugging Face Transformers."
)

demo.launch()