Update app.py

app.py

@@ -13,20 +13,29 @@ tokenizer = AutoTokenizer.from_pretrained(checkpoint)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
+prob_thresh = 0.3
+
 # Function to perform inference
 def predict(text):
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
     outputs = model(**inputs)
     logits = outputs.logits
     probabilities = torch.sigmoid(logits)  # Use sigmoid for multi-label classification
+    print(probabilities)
 
     # Get predicted labels based on a threshold (e.g., 0.5)
-    predicted_labels = (probabilities > 0.5).nonzero()[:, 1].tolist()
+    predicted_labels = (probabilities > prob_thresh).nonzero()[:, 1].tolist()
+    # positions = (probabilities > 0.5).nonzero(as_tuple=False)
+    prob_values = probabilities[probabilities > prob_thresh].tolist()
+    print(predicted_labels)
+    print(prob_values)
 
     # Map label IDs back to label names
     predicted_labels_names = [model.config.id2label[label_id] for label_id in predicted_labels]
-
-    return predicted_labels_names
+    labels_dict = {model.config.id2label[label_id]: prob for label_id, prob in zip(predicted_labels, prob_values)}
+    print(labels_dict)
+    # labels_dict = {label: 1/len(predicted_labels_names) for label in predicted_labels_names}
+    return labels_dict
 
 
 # Create the Gradio interface