from transformers import MarianMTModel, MarianTokenizer
import gradio as gr
import torch

# Load model and tokenizer
MODEL_NAME = "Helsinki-NLP/opus-mt-en-ur"

tokenizer = MarianTokenizer.from_pretrained(MODEL_NAME)
model = MarianMTModel.from_pretrained(MODEL_NAME)

# Optional: use GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

# Translation function
def translate(text):
    if not text.strip():
        return "Please enter some text."

    inputs = tokenizer(
        text,
        return_tensors="pt",
        padding=True,
        truncation=True
    ).to(device)

    translated = model.generate(**inputs)
    output = tokenizer.decode(translated[0], skip_special_tokens=True)

    return output

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 🌐 English to Urdu Translator")
    gr.Markdown("Translate English text into Urdu using AI")

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

    translate_btn = gr.Button("Translate")

    output_text = gr.Textbox(
        label="Urdu Translation",
        lines=5
    )

    # Button click event
    translate_btn.click(
        fn=translate,
        inputs=input_text,
        outputs=output_text
    )

# Launch app
if __name__ == "__main__":
    demo.launch()