Update app.py

app.py

@@ -67,6 +67,7 @@ def preprocess(text):
     text = ' '.join([lemmatizer.lemmatize(word) for word in text.split() if word not in stop_words])  # lemmatizācija
     return text
 
+
 # Classification function (single model)
 def classify_email_single_model(text, model_name):
     text = preprocess(text)
@@ -74,15 +75,22 @@ def classify_email_single_model(text, model_name):
     with torch.no_grad():
         outputs = models[model_name](**inputs)
         prediction = torch.argmax(outputs.logits, dim=1).item()
-    return labels[prediction]
+    probs = F.softmax(logits, dim=1)
+    probs_percent = probs.cpu().numpy() * 100
+    response = {"prediction": labels[prediction], "probabilities": probs_percent}
+    return response
 
+    
 # Classification function (all models together)
 def classify_email(text):
     votes = {"Safe": 0, "Spam": 0, "Phishing": 0}
+    probabilities = {}
     
     for model_name in model_names:
-        vote = classify_email_single_model(text, model_name)
+        response = classify_email_single_model(text, model_name)
+        vote = response['prediction']
         votes[vote] += 1
+        probabilities[model_name] = response['probabilities']
         
         
     response = ""
@@ -92,8 +100,14 @@ def classify_email(text):
         if i != 3:
             response += f"{label}: {vote_count} {vote_or_votes}, "
         else:
-            response += f"{label}: {vote_count} {vote_or_votes}"
+            response += f"{label}: {vote_count} {vote_or_votes}\n"
         i += 1
+
+    for model_name in model_names:
+        response += f"{model_name}: "
+        for j, prob in enumerate(probabilities[model_name]):
+            response += f"{num_to_label[j]}: {prob:.2f}%"
+        response += "\n"
         
     return response