import gradio as gr
from transformers import pipeline

available_models = {
    "mt5.baseline": pipeline("text2text-generation", model="samzirbo/mt5.baseline"),
    "mt5.genered": pipeline("text2text-generation", model="samzirbo/mt5.gendered"),
    "mt5.balanced": pipeline("text2text-generation", model="samzirbo/mt5.balanced"),
    "mt5.gendered_balanced": pipeline("text2text-generation", model="samzirbo/mt5.gendered_balanced")
}

def translate_text(model_name, lang_dir, gender, input_text):

    model = available_models[model_name]

    src, tgt = lang_dir.split(" -> ")

    prompt = f"Translate {src} to {tgt} " + f"as a {gender} : " if gender and  "gendered" in model_name else f"Translate {src} to {tgt} : "
    inputs = prompt + input_text
    print(inputs)
    output_text = model(inputs, max_length=128)
    return output_text[0]['generated_text']


model_dropdown = gr.Dropdown(choices=list(available_models.keys()), label="Select Model", value="mt5.baseline")
lang_dropdown = gr.Dropdown(choices=["English -> Spanish", "Spanish -> English"], label="Language Direction", value="English -> Spanish")
gender_dropdown = gr.Dropdown(choices=["female", "male"], label="Select Gender", value=None)

iface = gr.Interface(fn=translate_text, 
                     inputs=[model_dropdown, lang_dropdown, gender_dropdown, "text"], 
                     outputs="text", 
                     title="Translation Interface",
                     description="Select a model, language direction, and input text to translate.")

iface.launch()