# Install necessary libraries first
!pip install gradio transformers sentencepiece

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

# Load models
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)

# Define translation functions
def translate_en_to_ur(text):
    inputs = en_to_ur_tokenizer(text, return_tensors="pt", padding=True)
    translated = en_to_ur_model.generate(**inputs)
    result = en_to_ur_tokenizer.decode(translated[0], skip_special_tokens=True)
    return result

def translate_ur_to_en(text):
    inputs = ur_to_en_tokenizer(text, return_tensors="pt", padding=True)
    translated = ur_to_en_model.generate(**inputs)
    result = ur_to_en_tokenizer.decode(translated[0], skip_special_tokens=True)
    return result

# Create Gradio Interface
with gr.Blocks() as app:
    gr.Markdown("## 📝 English ↔ Urdu Translator (Free, Open Source)")

    with gr.Row():
        input_text = gr.Textbox(lines=4, placeholder="Enter your text here...")

    with gr.Row():
        en_to_ur_button = gr.Button("Translate English → Urdu")
        ur_to_en_button = gr.Button("Translate Urdu → English")

    output_text = gr.Textbox(lines=4, label="Translated Text")

    en_to_ur_button.click(fn=translate_en_to_ur, inputs=input_text, outputs=output_text)
    ur_to_en_button.click(fn=translate_ur_to_en, inputs=input_text, outputs=output_text)

# Launch app
app.launch()