import gradio as gr
from huggingface_hub import InferenceClient
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load both translation models from Hugging Face
# English to Moroccan Arabic (Darija)
tokenizer_eng_to_darija = AutoTokenizer.from_pretrained("Saidtaoussi/AraT5_Darija_to_MSA")
model_eng_to_darija = AutoModelForSeq2SeqLM.from_pretrained("Saidtaoussi/AraT5_Darija_to_MSA")

# Moroccan Arabic (Darija) to Modern Standard Arabic (MSA)
tokenizer_darija_to_msa = AutoTokenizer.from_pretrained("lachkarsalim/Helsinki-translation-English_Moroccan-Arabic")
model_darija_to_msa = AutoModelForSeq2SeqLM.from_pretrained("lachkarsalim/Helsinki-translation-English_Moroccan-Arabic")

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    translation_choice: str,
):
    # Ensure there's no empty input
    if not message.strip():
        return "Error: Please enter a valid text to translate."

    # Initialize the response variable
    response = ""

    # Translate based on the user's choice
    try:
        if translation_choice == "Moroccan Arabic to MSA":
            # Translate Moroccan Arabic (Darija) to Modern Standard Arabic
            inputs = tokenizer_darija_to_msa(message, return_tensors="pt", padding=True)
            outputs = model_darija_to_msa.generate(inputs["input_ids"], num_beams=5, max_length=max_tokens, early_stopping=True)
            response = tokenizer_darija_to_msa.decode(outputs[0], skip_special_tokens=True)

        elif translation_choice == "English to Moroccan Arabic":
            # Translate English to Moroccan Arabic (Darija)
            inputs = tokenizer_eng_to_darija(message, return_tensors="pt", padding=True)
            outputs = model_eng_to_darija.generate(inputs["input_ids"], num_beams=5, max_length=max_tokens, early_stopping=True)
            response = tokenizer_eng_to_darija.decode(outputs[0], skip_special_tokens=True)
    except Exception as e:
        response = f"Error occurred: {str(e)}"

    return response


# Gradio interface setup without pre-filled system message
demo = gr.Interface(
    fn=respond,
    inputs=[
        gr.Textbox(value="", label="Enter Your Text", placeholder="Type your sentence here..."),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Dropdown(
            label="Choose Translation Direction",
            choices=["English to Moroccan Arabic", "Moroccan Arabic to MSA"],
            value="English to Moroccan Arabic"
        ),
    ],
    outputs="text"
)

# Launch the interface
demo.launch()