import gradio as gr
from transformers import MBartForConditionalGeneration, MBart50Tokenizer

# Load model and tokenizer
model_name = "facebook/mbart-large-50"
model = MBartForConditionalGeneration.from_pretrained(model_name)
tokenizer = MBart50Tokenizer.from_pretrained(model_name)

def translate(text, direction):
    if direction == "Arab → Indonesia":
        src_lang = "ar_AR"
        tgt_lang = "id_ID"
    else:
        src_lang = "id_ID"
        tgt_lang = "ar_AR"
    
    tokenizer.src_lang = src_lang
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    generated_tokens = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang])
    translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
    return translation

# Define UI
demo = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(label="Masukkan teks"),
        gr.Radio(["Arab → Indonesia", "Indonesia → Arab"], label="Arah Terjemahan")
    ],
    outputs=gr.Textbox(label="Hasil Terjemahan"),
    title="Penerjemah Arab - Indonesia dengan mBART50",
    description="Masukkan teks dalam bahasa Arab atau Indonesia dan pilih arah terjemahan."
)

demo.launch()