import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Set up OpenAI API key
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)

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