output prob distribution

handler.py

@@ -18,6 +18,7 @@ class EndpointHandler:
         if not input_text:
             return {"error": "No input provided."}
 
+        # Tokenization
         inputs = self.tokenizer(
             input_text,
             return_tensors="pt",
@@ -27,14 +28,24 @@ class EndpointHandler:
         )
         inputs = {k: v.to(self.device) for k, v in inputs.items()}
 
+        # Forward pass
         with torch.no_grad():
             outputs = self.model(**inputs)
-            probs = torch.softmax(outputs.logits, dim=-1)
-            top_class_id = torch.argmax(probs, dim=-1).item()
-            top_class_label = self.id2label.get(top_class_id) or self.id2label.get(str(top_class_id))
-            top_class_prob = probs[0, top_class_id].item()
+            probs = torch.softmax(outputs.logits, dim=-1)[0]  # shape: (num_classes,)
+
+        # Get top class
+        top_class_id = torch.argmax(probs).item()
+        top_class_label = self.id2label.get(top_class_id) or self.id2label.get(str(top_class_id))
+        top_class_prob = probs[top_class_id].item()
+
+        # Convert full distribution to label->probability dict
+        prob_distribution = {
+            self.id2label.get(i) or self.id2label.get(str(i)): round(p.item(), 4)
+            for i, p in enumerate(probs)
+        }
 
         return {
             "label": top_class_label,
-            "confidence": round(top_class_prob, 4)
+            "probabilities": prob_distribution
         }
+