Update app.py

app.py

@@ -1,21 +1,19 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
+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):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
-    with torch.no_grad():
-        outputs = model(**inputs)
-    return torch.argmax(outputs.logits, dim=1).item()
-
+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()
+iface.launch()