import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline

# Fixed model path
MODEL_PATH = "BounharAbdelaziz/Terjman-Ultra-v2.0"
MAX_LEN = 1024

# Translation function
def translate_text(text):

    # Load the model
    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH)

    # Load the tokenizer with specific tokenizer class
    tokenizer = AutoTokenizer.from_pretrained(
    MODEL_PATH,
    src_lang="eng_Latn",
    tgt_lang="ary_Arab"
    )

    # Create translation pipeline with language configuration
    translator = pipeline(
    "translation",
    model=model,
    tokenizer=tokenizer,
    max_length=MAX_LEN,
    src_lang="eng_Latn",
    tgt_lang="ary_Arab"
    )

    # run prediction
    translation = translator(text)[0]['translation_text']
    return translation

# Gradio app
def gradio_app():
    with gr.Blocks() as app:
        gr.Markdown("# 🇲🇦 Terjman-Ultra v2.0:")
        gr.Markdown("Enter the English text you want to translate to Moroccan Darija.")
        input_text = gr.Textbox(
            label="Input Text", 
            placeholder="Enter text to translate...",
            lines=3
        )
        output_text = gr.Textbox(
            label="Translated Text", 
            interactive=False,
            lines=3
        )
        translate_button = gr.Button("Translate")
        
        # Link input and output
        translate_button.click(
            fn=translate_text,
            inputs=input_text,
            outputs=output_text
        )
    return app

# Run the app
if __name__ == "__main__":
    app = gradio_app()
    app.launch()