return confidence score

app.py

@@ -10,23 +10,25 @@ APP_NAME = "Mona Lisa Detection"
 
 logger.debug("loading processor and model.")
 processor = AutoFeatureExtractor.from_pretrained(
-    "drift-ai/autotrain-mona-lisa-detection-38345101350",
-    use_auth_token=True
+    "drift-ai/autotrain-mona-lisa-detection-38345101350", use_auth_token=True
 )
 model = AutoModelForImageClassification.from_pretrained(
-    "drift-ai/autotrain-mona-lisa-detection-38345101350",
-    use_auth_token=True
+    "drift-ai/autotrain-mona-lisa-detection-38345101350", use_auth_token=True
 )
 logger.debug("loading processor and model succeeded.")
 
 
 def process_image(image, model=model, processor=processor):
     logger.info("Making a prediction ...")
+
     inputs = processor(images=image, return_tensors="pt")
     outputs = model(**inputs)
     logits = outputs.logits
     predicted_class_idx = logits.argmax(-1).item()
-    result = model.config.id2label[predicted_class_idx]
+
+    label = {1: "Not Mona Lisa", 0: "Mona Lisa"}
+    predictions = logits.softmax(dim=-1).tolist()
+    result = {label[predicted_class_idx]: predictions[0][predicted_class_idx]}
     print("Predicted class:", result)
     logger.info("Prediction finished.")
     return result