Update app.py

app.py

@@ -32,16 +32,23 @@ def respond(message, history, system_message=None, max_tokens=None, temperature=
 
     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]
 
-    # 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()}"
+        # Número de classes
+        num_labels = model.config.num_labels
+
+        if num_labels == 1:
+            prob = torch.sigmoid(logits).item()
+            output = f"Predicted logit: {logits.item():.4f}\nProbability (sigmoid): {prob:.4f}"
+        else:
+            probs = F.softmax(logits, dim=-1).squeeze().tolist()
+            pred_index = torch.argmax(logits, dim=-1).item()
+            pred_label = labels[pred_index]
+
+            output = f"Predicted label: {pred_label}\nProbabilities:\n"
+            for i, p in enumerate(probs):
+                output += f"{labels[i]}: {p:.4f}\n"
+            output += f"\nLogits: {logits.squeeze().tolist()}"
+
     return output
 
 # -------------------------------