import gradio as gr
from transformers import MarianMTModel, MarianTokenizer

# Load English to Arabic translation model and tokenizer
en_ar_model_name = "Helsinki-NLP/opus-mt-en-ar"
en_ar_model = MarianMTModel.from_pretrained(en_ar_model_name)
en_ar_tokenizer = MarianTokenizer.from_pretrained(en_ar_model_name)

# Load Arabic to English translation model and tokenizer
ar_en_model_name = "Helsinki-NLP/opus-mt-ar-en"
ar_en_model = MarianMTModel.from_pretrained(ar_en_model_name)
ar_en_tokenizer = MarianTokenizer.from_pretrained(ar_en_model_name)

def translate_en_to_ar(text):
    inputs = en_ar_tokenizer.encode(text, return_tensors="pt")
    translation = en_ar_model.generate(inputs, max_length=128)
    translated_text = en_ar_tokenizer.decode(translation[0], skip_special_tokens=True)
    return translated_text

def translate_ar_to_en(text):
    inputs = ar_en_tokenizer.encode(text, return_tensors="pt")
    translation = ar_en_model.generate(inputs, max_length=128)
    translated_text = ar_en_tokenizer.decode(translation[0], skip_special_tokens=True)
    return translated_text

# Create Gradio interfaces for both translation directions
en_ar_interface = gr.Interface(
    fn=translate_en_to_ar,
    inputs=gr.inputs.Textbox(),
    outputs=gr.outputs.Textbox(),
    title="English to Arabic Translation",
    description="Translate English text to Arabic."
)

ar_en_interface = gr.Interface(
    fn=translate_ar_to_en,
    inputs=gr.inputs.Textbox(),
    outputs=gr.outputs.Textbox(),
    title="Arabic to English Translation",
    description="Translate Arabic text to English."
)

# # Combine the interfaces in a single app
# app = gr.Interface(
#     fn=[en_ar_interface, ar_en_interface],
#     layout="horizontal",
#     title="Translation App",
#     description="Translate text between English and Arabic."
# )



if __name__ == "__main__":
    # app.launch()
    # ar_en_interface.launch()
    gr.Interface([en_ar_interface, ar_en_interface]).launch()