import gradio as gr
from transformers import MarianMTModel, MarianTokenizer

# Load English → Urdu model
en_ur_model_name = 'Helsinki-NLP/opus-mt-en-ur'
en_ur_tokenizer = MarianTokenizer.from_pretrained(en_ur_model_name)
en_ur_model = MarianMTModel.from_pretrained(en_ur_model_name)

# Load Urdu → English model
ur_en_model_name = 'Helsinki-NLP/opus-mt-ur-en'
ur_en_tokenizer = MarianTokenizer.from_pretrained(ur_en_model_name)
ur_en_model = MarianMTModel.from_pretrained(ur_en_model_name)

# Define translation function
def translate_text(text, direction):
    if direction == "English to Urdu":
        tokenizer, model = en_ur_tokenizer, en_ur_model
    else:
        tokenizer, model = ur_en_tokenizer, ur_en_model

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

# Gradio Interface
iface = gr.Interface(
    fn=translate_text,
    inputs=[
        gr.Textbox(label="Enter Text", placeholder="Type text here..."),
        gr.Radio(["English to Urdu", "Urdu to English"], label="Translation Direction")
    ],
    outputs=gr.Textbox(label="Translated Text"),
    title="English ↔ Urdu Translator Chatbot",
    description="Translate between English and Urdu using pre-trained models from Hugging Face."
)

# Launch the interface
iface.launch()