# Install necessary libraries (run this once)!pip install transformers gradio

# Now import them
from transformers import MarianMTModel, MarianTokenizer
import gradio as gr

# Define the models
models = {
    "English to Urdu": {
        "model_name": "Helsinki-NLP/opus-mt-en-ur"
    },
    "Urdu to English": {
        "model_name": "Helsinki-NLP/opus-mt-ur-en"
    }
}

# Load models and tokenizers
loaded_models = {}
for direction, info in models.items():
    tokenizer = MarianTokenizer.from_pretrained(info["model_name"])
    model = MarianMTModel.from_pretrained(info["model_name"])
    loaded_models[direction] = (tokenizer, model)

# Define the translation function
def translate_text(text, direction):
    tokenizer, model = loaded_models[direction]
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    translated = model.generate(**inputs, max_length=512)
    output = tokenizer.decode(translated[0], skip_special_tokens=True)
    return output

# Create Gradio Interface
iface = gr.Interface(
    fn=translate_text,
    inputs=[
        gr.Textbox(label="Enter text here", placeholder="Type your English or Urdu text..."),
        gr.Radio(["English to Urdu", "Urdu to English"], label="Select translation direction")
    ],
    outputs=gr.Textbox(label="Translated Text"),
    title="🌍 English ↔ Urdu Translator",
    description="Translate text between English and Urdu using Hugging Face pretrained models.",
    theme="default"
)

# Launch the app
iface.launch()