Update app.py

app.py

@@ -1,40 +1,48 @@
 import gradio as gr
 import torch
 from transformers import DistilBertTokenizer, AutoModelForSequenceClassification
+import torch.nn.functional as F
 
 # -------------------------------
 # Configuração do modelo
 # -------------------------------
 MODEL_NAME = "hermanshid/distilbert-id-law"
 
-# Como o modelo não tem tokenizer, usamos um tokenizer de DistilBERT base compatível
+# Tokenizer compatível do DistilBERT base
 tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
 
-# Carrega o modelo customizado
+# Modelo customizado
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
-
-# Definir dispositivo (CPU suficiente para Space gratuito)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 model.eval()
 
+# -------------------------------
+# Labels dummy para visualização
+# -------------------------------
+num_labels = model.config.num_labels
+labels = [f"Class {i}" for i in range(num_labels)]
+
 # -------------------------------
 # Função de inferência
 # -------------------------------
 def respond(message, history, system_message=None, max_tokens=None, temperature=None, top_p=None, hf_token=None):
-    """
-    Recebe texto do usuário, retorna a classe prevista pelo DistilBERT.
-    Os parâmetros history, system_message, max_tokens etc são mantidos
-    apenas para compatibilidade com ChatInterface.
-    """
     inputs = tokenizer(message, return_tensors="pt", truncation=True, padding=True)
     inputs = {k: v.to(device) for k, v in inputs.items()}
 
     with torch.no_grad():
         logits = model(**inputs).logits
+        probs = F.softmax(logits, dim=-1).squeeze().tolist()  # probabilidades
+        pred_index = torch.argmax(logits, dim=-1).item()
+        pred_label = labels[pred_index]
 
-    pred_class = torch.argmax(logits, dim=-1).item()
-    return f"Predicted class: {pred_class}"
+    # Monta mensagem de saída detalhada
+    output = f"Predicted label: {pred_label}\n\n"
+    output += "Probabilities:\n"
+    for i, p in enumerate(probs):
+        output += f"{labels[i]}: {p:.4f}\n"
+    output += f"\nLogits: {logits.squeeze().tolist()}"
+    return output
 
 # -------------------------------
 # Interface estilo Chatbot
@@ -59,4 +67,4 @@ with gr.Blocks() as demo:
 # Inicializa a aplicação
 # -------------------------------
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()