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

# Define model paths
MODEL_PATHS = {
    "Terjman-Nano-v2": "BounharAbdelaziz/Terjman-Nano-v2.0",
    "Terjman-Large-v2": "BounharAbdelaziz/Terjman-Large-v2.0",
    "Terjman-Ultra-v2": "BounharAbdelaziz/Terjman-Ultra-v2.0",
    "Terjman-Supreme-v2": "BounharAbdelaziz/Terjman-Supreme-v2.0"
}

# Load environment token
TOKEN = os.environ['TOKEN']

# Translation function for Nano and Large models
def translate_nano_large(text, model_path):
    translator = pipeline("translation", model=model_path, token=TOKEN)
    translated = translator(
        text, 
        max_length=512,
        num_beams=4,
        no_repeat_ngram_size=3,
        early_stopping=True,
        do_sample=False,
        pad_token_id=translator.tokenizer.pad_token_id,
        bos_token_id=translator.tokenizer.bos_token_id,
        eos_token_id=translator.tokenizer.eos_token_id
    )
    return translated[0]["translation_text"]

# Translation function for Ultra and Supreme models
def translate_ultra_supreme(text, model_path):
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path, token=TOKEN)
    tokenizer = AutoTokenizer.from_pretrained(model_path, src_lang="eng_Latn", tgt_lang="ary_Arab", token=TOKEN)
    translator = pipeline(
        "translation",
        model=model,
        tokenizer=tokenizer,
        max_length=512,
        src_lang="eng_Latn",
        tgt_lang="ary_Arab"
    )
    translation = translator(text)[0]['translation_text']
    return translation

# Main translation function
def translate_text(text, model_choice):
    model_path = MODEL_PATHS[model_choice]
    if model_choice in ["Terjman-Nano-v2", "Terjman-Large-v2"]:
        return translate_nano_large(text, model_path)
    elif model_choice in ["Terjman-Ultra-v2", "Terjman-Supreme-v2"]:
        return translate_ultra_supreme(text, model_path)
    else:
        return "Invalid model selection."

# Gradio app
def gradio_app():
    with gr.Blocks() as app:
        gr.Markdown("# 🇲🇦 Terjman-v2")
        gr.Markdown("Choose a model and enter the English text you want to translate to Moroccan Darija.")

        model_choice = gr.Dropdown(
            label="Select Model",
            choices=["Terjman-Nano-v2", "Terjman-Large-v2", "Terjman-Ultra-v2", "Terjman-Supreme-v2"],
            value="Terjman-Ultra-v2"
        )
        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, model_choice],
            outputs=output_text
        )

    return app

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