import gradio as gr
from transformers import pipeline
import os

# Fixed model path
MODEL_PATH = "BounharAbdelaziz/Terjman-Large-v2.2-bs-16-lr-0.001-ep-2-wp-0.1-gacc-8-gnm-1.0-mx-512-v2.2" #"BounharAbdelaziz/Terjman-Large-v2.2"

TOKEN=os.environ['TOKEN']

# Load the translation pipeline with the fixed model
translator = pipeline("translation", model=MODEL_PATH, token=TOKEN)

# Translation function
def translate_text(text):
    # Perform translation
    translated = translator(
        text, 
        max_length=512,
        num_beams=4,  # Beam search for better quality
        no_repeat_ngram_size=3,  # Avoid repetition
        early_stopping=True,  # Stop when viable translations are found
        do_sample=False,  # Deterministic output
        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"]

# Gradio app
def gradio_app():
    with gr.Blocks() as app:
        gr.Markdown("# 🇲🇦 Terjman-Large 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...")
        output_text = gr.Textbox(label="Translated Text", interactive=False)

        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()