Update app.py

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 import torch
 from transformers import pipeline
 
-# 1. Load the phishing content detector pipeline
 device = 0 if torch.cuda.is_available() else -1
 detector = pipeline(
     "text-classification",
@@ -11,23 +10,17 @@ detector = pipeline(
     top_k=1,
 )
 
-# 2. Wrap prediction in a function for Gradio
 def classify_text(text: str):
-    """
-    Returns (label, score) for the given text.
-    Handles both flat and nested list outputs from the pipeline.
-    """
     if not text or not text.strip():
         return "⚠️ Please enter some text", 0.0
 
     preds = detector(text)
-    # If preds is [[{...}]], unwrap one level
     first = preds[0][0] if isinstance(preds[0], list) else preds[0]
-    label = first["label"]
+    raw = first["label"]
+    label = "Phishing" if raw == "LABEL_1" else "Legitimate"
     score = float(first["score"])
     return label, round(score, 4)
 
-# 3. Realistic example test cases
 examples = [
     # Phishing
     ["Congratulations! You've won a $1,000 gift card. Click here to claim: http://bit.ly/free-gift"],
@@ -48,7 +41,7 @@ with gr.Blocks(theme="default") as demo:
     gr.Markdown(
         """
         # 🚨 Phishing Content Detector  
-        Paste any email or message snippet below and this model will predict whether it's **PHISHING** or **LEGIT**.
+        Paste any email or message snippet below and this model will predict whether it's **Phishing** or **Legitimate**.
         """
     )
 
@@ -76,7 +69,7 @@ with gr.Blocks(theme="default") as demo:
 
     gr.Markdown(
         """
-        **Model:** Version 1. 
+        **Model:** Version 1.0
         """
     )