Update app.py

app.py

@@ -82,8 +82,8 @@ def classify_email_single_model(text, model_name):
     return response
 
     
-# Classification function (all models together)
-def classify_email(text):
+# Classification function (all models together, probabilities for each model)
+def classify_email_detailed(text):
     votes = {"Safe": 0, "Spam": 0, "Phishing": 0}
     probabilities = {}
     
@@ -111,11 +111,43 @@ def classify_email(text):
         response += "\n"
         
     return response
+
+# Classification function (all models together, just the votes)
+def classify_email_simple(text):
+    votes = {"Safe": 0, "Spam": 0, "Phishing": 0}
+    
+    for model_name in model_names:
+        response = classify_email_single_model(text, model_name)
+        vote = response['prediction']
+        votes[vote] += 1
+        
+    response = ""
+    i = 1
+    for label, vote_count in votes.items():
+        vote_or_votes = "vote" if vote_count == 1 else "votes"
+        if i != 3:
+            response += f"{label}: {vote_count} {vote_or_votes}, "
+        else:
+            response += f"{label}: {vote_count} {vote_or_votes}\n"
+        i += 1
+    response += "\n"
+        
+    return response
+
+
+def classify_email(text, mode):
+    if mode == "Tikai balsis":
+        return classify_email_simple(text)
+    else:
+        return classify_email_detailed(text)
+
     
 # Gradio UI
 demo = gr.Interface(
     fn=classify_email,
-    inputs=gr.Textbox(lines=10, placeholder="Ievietojiet savu e-pastu šeit..."),
+    inputs=[gr.Textbox(lines=10, placeholder="Ievietojiet savu e-pastu šeit..."),
+            gr.Radio(choices=["Tikai balsis", "Balsis un varbūtības"], label='Klasifikācijas veids')
+    ],
     outputs="text",
     title="E-pastu klasifikators (vairāku modeļu balsošana)",
     description="Autori: Kristaps Tretjuks un Aleksejs Gorlovičs"