import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

model_name = "santyzenith/emergency_actors_classification"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

def generate_text(input_text):
    input_ids = tokenizer.encode(input_text, return_tensors="pt")
    output_ids = model.generate(input_ids, max_length=512, eos_token_id=2)
    output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return output_text

in_textbox = gr.Textbox(
    placeholder="Ingresa el texto"
)

out_textbox = gr.Textbox()

interface = gr.Interface(
    fn=generate_text,
    inputs=in_textbox,
    outputs=out_textbox,
    title="Identificación de actores en llamadas de Emergencia",
    description="""Esta aplicación utiliza un modelo ajustado en español para clasificar llamadas de emergencia del ECU 911: \n
    '<O>' y '</O>' indican el comienzo y fin de una oración del OPERADOR.\n
    '<A>' y '</A>' indican el comiendo y fin de una oración del ALERTANTE.\n
    """
)

interface.launch()