import gradio as gr
from transformers import MBart50TokenizerFast, MBartForConditionalGeneration

# Load the fine-tuned model
model_name = "abdulwaheed1/urdu_to_english_translation_mbart"
tokenizer = MBart50TokenizerFast.from_pretrained(model_name, src_lang="ur_PK", tgt_lang="en_XX")
model = MBartForConditionalGeneration.from_pretrained(model_name)

def translate_urdu_to_english(text):
    # Tokenize the input text
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    # Generate translation
    outputs = model.generate(**inputs)
    # Decode the translated text
    translation = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
    return translation

# Create the Gradio interface
iface = gr.Interface(
    fn=translate_urdu_to_english,
    inputs=gr.Textbox(label="Input Urdu Text"),
    outputs=gr.Textbox(label="Translated English Text"),
    title="Urdu to English Translation",
    description="Enter Urdu text to get the English translation."
)

# Launch the app
iface.launch()