app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline

# Carregar o modelo e o tokenizer
model_name = "vic35get/nhtsa_complaints_classifier"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
pipeline_clf = pipeline("text-classification", tokenizer=tokenizer, model=model)

# Função para inferência
def predict(text: str):
    classification = pipeline_clf(text)[0]
    return classification.get('label')
    
# Interface Gradio
iface = gr.Interface(fn=predict, inputs="text", outputs="text")

# Rodar a interface
iface.launch()