Update README.md

README.md

@@ -55,23 +55,55 @@ import torch
 tokenizer = AutoTokenizer.from_pretrained("your-username/model-name")
 model = AutoModelForSequenceClassification.from_pretrained("your-username/model-name")
 
-def predict_phishing(email_text):
+def predict_email(email_text):
     # Preprocess and tokenize
-    inputs = tokenizer(email_text, return_tensors="pt", truncation=True, max_length=512)
+    inputs = tokenizer(
+        email_text,
+        return_tensors="pt",
+        truncation=True,
+        max_length=512
+    )
     
     # Get prediction
     with torch.no_grad():
         outputs = model(**inputs)
         predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
     
+    # Get probabilities for each class
+    probs = predictions[0].tolist()
+    
+    # Create labels dictionary
+    labels = {
+        "legitimate_email": probs[0],
+        "phishing_url": probs[1],
+        "legitimate_url": probs[2],
+        "phishing_url_alt": probs[3]
+    }
+    
+    # Determine the most likely classification
+    max_label = max(labels.items(), key=lambda x: x[1])
+    
     return {
-        "is_phishing": bool(predictions[0][1] > 0.5),
-        "confidence": float(predictions[0][1])
+        "prediction": max_label[0],
+        "confidence": max_label[1],
+        "all_probabilities": labels
     }
+```
+
+**Example Usage**
 
+```python
 # Example usage
-email = "Your email text here..."
-result = predict_phishing(email)
-print(f"Is Phishing: {result['is_phishing']}")
+email = """
+Dear User,
+Your account security needs immediate attention. Please verify your credentials.
+Click here: http://suspicious-link.com
+"""
+
+result = predict_email(email)
+print(f"Prediction: {result['prediction']}")
 print(f"Confidence: {result['confidence']:.2%}")
+print("\nAll probabilities:")
+for label, prob in result['all_probabilities'].items():
+    print(f"{label}: {prob:.2%}")
 ```