update main py

now prints out all classes confidence

main.py

@@ -84,29 +84,15 @@ def get_preds(user_text, model=model, tokenizer=tokenizer, label_names=label_nam
 
     # Process the outputs as needed
     # For example, if you have a classification task:
-    predictions = torch.nn.functional.softmax(outputs.logits, dim=1)
+    predictions = torch.nn.functional.softmax(outputs.logits, dim=1).tolist()
+    predictions = [item for sublist in predictions for item in sublist]
 
-    # Get the predicted label indices for each prediction
-    predicted_label_indices = predictions.argmax(dim=1).cpu().numpy()
+    print(predictions)
+    
+    dicti = dict(zip(label_names, predictions))
+    print(dicti)
 
-    # Get the confidence scores (probabilities) for each prediction
-    confidence_scores = predictions.max(dim=1).values.cpu().numpy()
-
-    # Combine the category names, predicted label indices, and confidence scores into a list of tuples
-    predictions_with_confidence = list(zip(predicted_label_indices, confidence_scores))
-
-    # Sort the predictions with confidence scores in descending order
-    predictions_with_confidence_sorted = sorted(predictions_with_confidence, key=lambda x: x[1], reverse=True)
-
-    # Print all predictions with category names and confidence scores in descending order
-    category_list, conf_sc = [], []
-    for label_index, confidence_score in predictions_with_confidence_sorted:
-        category_name = model.config.id2label[str(label_index)]
-        print(f"Category: {category_name}, Confidence: {confidence_score:.2f}")
-        category_list.append(category_name)
-        conf_sc.append(confidence_score)
-
-    return category_list
+    return dicti
 
 @app.get("/", tags=["Home"])
 def api_home(user_text: str):
@@ -116,6 +102,8 @@ def api_home(user_text: str):
 def inference(user_text: str):
     return get_preds(user_text=user_text)
 
+# get_preds(user_text=user_text)
+
 
 if __name__=='__main__':
-    uvicorn.run('main:app', reload=True)
+    uvicorn.run('main:app', reload=True)