import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM



# Model in Hugging Hub
tokenizer = AutoTokenizer.from_pretrained("Andresmfs/st5s-es-inclusivo")
model = AutoModelForSeq2SeqLM.from_pretrained("Andresmfs/st5s-es-inclusivo")

def make_neutral(phrase):
    # Define prompt for converting gendered text to neutral
    input_ids = tokenizer(phrase, return_tensors="pt").input_ids

    # Call the LLM to generate neutral text
    outputs = model.generate(input_ids)

    return tokenizer.decode(outputs[0], skip_special_tokens=True)

theme = gr.themes.Default(primary_hue="blue").set(
    loader_color="#FF0000",
    slider_color="#FF0000",
)

with gr.Blocks(theme=theme) as demo:
    
    iface = gr.Interface(
        fn=make_neutral,
        inputs="text",
        outputs="text",
        title="Gender-Neutralizer",
        description="Enter a gendered phrase and get it converted into neutral form."
    )
    iface.launch()