Organized folders, and added a BERT-Mini model and explainer to use for email classifications

app/main.py

@@ -1,7 +1,7 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 from typing import List, Tuple, Optional 
-from .ml_logic import get_prediction_and_explanation # helper function from ml_logic
+from .ml import get_model_prediction, check_model_status
 
 app = FastAPI(title="AI-Powered Phishing Email Detection System")
 
@@ -10,6 +10,7 @@ class EmailInput(BaseModel):
     subject: Optional[str] = ""
     sender: Optional[str] = ""
     body: str
+    model_choice: Optional[str] = "nb" # Default to Naive Bayes
 
 # Define output data model
 class PredictionResponse(BaseModel):
@@ -24,16 +25,27 @@ class PredictionResponse(BaseModel):
 async def root():
     return {"message": "AI-Powered Phishing Email Detection API. POST to /predict with 'subject', 'sender', 'body'."}
 
+@app.get("/status")
+async def model_status():
+    return check_model_status() 
+
+
 @app.post("/predict", response_model=PredictionResponse)
 async def predict_email(email_input: EmailInput):
+ 
+    if email_input.model_choice not in ["nb", "bert-mini"]:
+        return PredictionResponse(prediction="Error", label=-1, confidence=0.0, explanation=[],
+                                  error="Invalid model_choice. Please use 'nb' or 'bert-mini'.")
     try:
-        result = get_prediction_and_explanation(
-            email_input.subject or "",
-            email_input.sender or "", 
-            email_input.body
+        result = get_model_prediction(
+            subject=email_input.subject or "", 
+            sender=email_input.sender or "",
+            body=email_input.body,
+            model_choice=email_input.model_choice
         )
-        if "error" in result and result["error"]:
-             return PredictionResponse(prediction="Error", label=-1, confidence=0.0, explanation=[], error=result["error"])
         return PredictionResponse(**result)
+
     except Exception as e:
-        return PredictionResponse(prediction="Error", label=-1, confidence=0.0, explanation=[], error=f"API error: {str(e)}")
+        # Fallback for truly unexpected errors in the endpoint itself
+        return PredictionResponse(prediction="Error", label=-1, confidence=0.0, explanation=[],
+                                  error=f"Critical API endpoint error: {str(e)}")

app/ml/__init__.py

@@ -0,0 +1,36 @@
+# app/ml/__init__.py
+from .nb_model import get_prediction_and_explanation_nb, nb_model, nb_preprocessor
+from .bert_mini_model import get_prediction_and_explanation_bert_mini, bert_mini_model, bert_mini_tokenizer
+from typing import Dict
+
+def get_model_prediction(subject: str, sender: str, body: str, model_choice: str = "nb") -> Dict:
+    """
+    # Dispatcher function to get predictions from the chosen model.
+    """
+    if model_choice == "bert-mini":
+        if bert_mini_model is None or bert_mini_tokenizer is None: 
+             return {"error": "BERT-Mini Model/Tokenizer is not available. Check server logs.",
+                    "prediction": "Error", "label": -1, "confidence": 0.0, "explanation": []}
+        return get_prediction_and_explanation_bert_mini(subject, sender, body)
+    elif model_choice == "nb":
+        if nb_model is None or nb_preprocessor is None: # Check if NB loaded successfully
+            return {"error": "Multinomial Naive Bayes Model/Preprocessor is not available. Check server logs.",
+                    "prediction": "Error", "label": -1, "confidence": 0.0, "explanation": []}
+        return get_prediction_and_explanation_nb(subject, sender, body)
+    else:
+        return {"error": f"Invalid model_choice: '{model_choice}'. Choose 'nb' or 'bert-mini'.",
+                "prediction": "Error", "label": -1, "confidence": 0.0, "explanation": []}
+
+# You can also add a health check function here if needed
+def check_model_status():
+    status = {
+        "naive_bayes": {
+            "model_loaded": nb_model is not None,
+            "preprocessor_loaded": nb_preprocessor is not None
+        },
+        "bert-mini": {
+            "model_loaded": bert_mini_model is not None,
+            "tokenizer_loaded": bert_mini_tokenizer is not None
+        }
+    }
+    return status

app/ml/bert_mini_model.py

@@ -0,0 +1,121 @@
+# app/ml/bert_mini_model.py
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+import numpy as np
+from .common import simple_text_clean, CLASS_NAMES
+import traceback
+from transformers_interpret import SequenceClassificationExplainer 
+
+# # Load BERT-mini model and tokenizer from Hugging Face Hub 
+BERT_MODEL_ID = "lleratodev/720-bert-mini-phishing" 
+bert_mini_tokenizer = None
+bert_mini_model = None
+device = None
+cls_explainer_bert_mini = None
+
+try:
+    bert_mini_tokenizer = AutoTokenizer.from_pretrained(BERT_MODEL_ID)
+    bert_mini_model = AutoModelForSequenceClassification.from_pretrained(BERT_MODEL_ID)
+    
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    bert_mini_model.to(device)
+    bert_mini_model.eval() 
+    cls_explainer_bert_mini = SequenceClassificationExplainer(bert_mini_model, bert_mini_tokenizer)
+    
+    print(f"BERT-mini model ('{BERT_MODEL_ID}'), tokenizer, and Transformers-Interpret Explainer loaded successfully.")
+    print(f"BERT-mini model running on device: {device}")
+
+except Exception as e:
+    print(f"FATAL ERROR (BERT): Could not load model/tokenizer '{BERT_MODEL_ID}' or initialize Transformers-Interpret Explainer: {e}")
+    traceback.print_exc()
+
+# # Using BERT-Mini model to make email classifications 
+def bert_mini_predict_probability_for_lime(text_instances: list) -> np.ndarray:
+    if bert_mini_tokenizer is None or bert_mini_model is None:
+        # Return neutral probabilities if model isn't loaded (number of instances, number of classes)
+        return np.array([[1.0/len(CLASS_NAMES)] * len(CLASS_NAMES)] * len(text_instances))
+
+    all_probabilities = []
+    try:
+        for text_instance in text_instances:
+            inputs = bert_mini_tokenizer(text_instance, return_tensors="pt", truncation=True, padding="max_length", max_length=512)
+            inputs = {k: v.to(device) for k, v in inputs.items()}
+
+            with torch.no_grad():
+                outputs = bert_mini_model(**inputs)
+                logits = outputs.logits
+            
+            probabilities_tensor = torch.softmax(logits, dim=-1)
+        
+            probabilities_for_instance = probabilities_tensor.cpu().numpy().squeeze() 
+            all_probabilities.append(probabilities_for_instance)
+        
+        return np.array(all_probabilities)
+
+    except Exception as e:
+        print(f"Error in bert_mini_predict_probability_for_lime: {e}")
+        traceback.print_exc()
+        return np.array([[1.0/len(CLASS_NAMES)] * len(CLASS_NAMES)] * len(text_instances))
+
+
+def get_prediction_and_explanation_bert_mini(subject: str, sender: str, body: str) -> dict:
+    if bert_mini_tokenizer is None or bert_mini_model is None or cls_explainer_bert_mini is None:
+        return {"error": "BERT-Mini Model/Tokenizer/Explainer not loaded correctly. Check server logs.",
+                "prediction": "Error", "label": -1, "confidence": 0.0, "explanation": []}
+
+    cleaned_sender = simple_text_clean(sender)
+    cleaned_subject = simple_text_clean(subject)
+    cleaned_body = simple_text_clean(body)
+    
+    combined_text_for_prediction = f"{cleaned_sender} {cleaned_subject} {cleaned_body}"
+    text_for_lime_explanation = combined_text_for_prediction
+
+
+    try:
+        inputs = bert_mini_tokenizer(combined_text_for_prediction, return_tensors="pt", truncation=True, padding="max_length", max_length=512)
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = bert_mini_model(**inputs)
+            logits = outputs.logits
+
+        probabilities_tensor = torch.softmax(logits, dim=-1)
+        probabilities = probabilities_tensor.cpu().numpy()[0]
+
+        prediction_label_int = np.argmax(probabilities).item()
+        confidence_score = probabilities[prediction_label_int].item()
+        predicted_class_name = CLASS_NAMES[prediction_label_int]
+
+        explanation_data = []
+        try:
+            word_attributions = cls_explainer_bert_mini(
+                combined_text_for_prediction,
+                index = prediction_label_int 
+            )
+            
+            explanation_data = [(word, float(score)) for word, score in word_attributions]
+            
+            explanation_data.sort(key=lambda x: abs(x[1]), reverse=True)
+            explanation_data = explanation_data[:15] 
+
+        except Exception as e:
+            print(f"Transformers-Interpret explanation error: {e}")
+            traceback.print_exc()
+            explanation_data = [("Explanation error with Transformers-Interpret", 0.0)]
+        # --- End Explanation ---
+
+        return {
+            "prediction": predicted_class_name,
+            "label": int(prediction_label_int),
+            "confidence": float(confidence_score),
+            "explanation": explanation_data,
+            "error": None
+        }
+    except Exception as e:
+        print(f"--- ORIGINAL ERROR in predict_with_bert_mini ---")
+        print(f"Error type: {type(e)}")
+        print(f"Error message: {str(e)}")
+        print("Traceback:")
+        traceback.print_exc()
+        return {"error": f"BERT-Mini Prediction error: {str(e)}", "prediction": "Error",
+                "label": -1, "confidence": 0.0, "explanation": []}

app/ml/common.py

@@ -0,0 +1,14 @@
+import re
+
+# Text cleaning function, makes everything lowercase, removed non alpha-numeric characters and normalize white spaces
+def simple_text_clean(text: str) -> str:
+    if isinstance(text, str):
+        text = text.lower()
+        text = re.sub(r'[^a-z0-9\s]', '', text) # Keep spaces, remove other non-alphanumeric
+        text = re.sub(r'\s+', ' ', text).strip()
+    else:
+        text = '' 
+    return text
+
+# Class names for predictions 
+CLASS_NAMES = ['Legitimate', 'Phishing'] # 0: Legitimate, 1: Phishing

app/ml/nb_model.py

@@ -0,0 +1,131 @@
+import joblib
+import pandas as pd
+import numpy as np
+import os
+from lime.lime_text import LimeTextExplainer
+from .common import simple_text_clean, CLASS_NAMES 
+
+ASSETS_DIR = os.path.join(os.path.dirname(__file__), '..', 'assets') 
+PREPROCESSOR_FILENAME = "email_preprocessor_20250506_203148.joblib"
+MODEL_FILENAME = "phishing_nb_model_20250506_203148.joblib"
+PREPROCESSOR_PATH = os.path.join(ASSETS_DIR, PREPROCESSOR_FILENAME)
+MODEL_PATH = os.path.join(ASSETS_DIR, MODEL_FILENAME)
+
+nb_preprocessor = None
+nb_model = None
+lime_explainer_nb = None
+
+try:
+    nb_preprocessor = joblib.load(PREPROCESSOR_PATH)
+    nb_model = joblib.load(MODEL_PATH)
+    lime_explainer_nb = LimeTextExplainer(class_names=CLASS_NAMES)
+    print("Multinomial NB model, Preprocessor, and LIME Explainer loaded successfully.")
+except FileNotFoundError:
+    print(f"FATAL ERROR (Naive Bayes): Could not find model ('{MODEL_PATH}') or nb_preprocessor ('{PREPROCESSOR_PATH}').")
+    print("Ensure files are in 'app/assets/' and filenames are correct.")
+except Exception as e:
+    print(f"Error loading Multinomial NB model/preprocessor or initializing LIME: {e}")
+
+def model_predict_probability_for_lime(combined_texts):
+    if nb_preprocessor is None or nb_model is None:
+        return np.array([[0.5, 0.5]] * len(combined_texts))
+    
+    subjects = []
+    senders = []
+    bodies = []
+
+    for combined_text in combined_texts:
+        s_marker = "subject: "
+        d_marker = " sender: " 
+        b_marker = " body: "    
+
+        s_text, d_text, b_text = "", "", ""
+
+        if d_marker in combined_text:
+            s_text_part, rest = combined_text.split(d_marker, 1)
+            if s_marker in s_text_part:
+                s_text = s_text_part.replace(s_marker, "").strip()
+            
+            if b_marker in rest:
+                d_text_part, b_text_part = rest.split(b_marker, 1)
+                d_text = d_text_part.strip()
+                b_text = b_text_part.strip()
+            else: 
+                d_text = rest.strip()
+        else: 
+             if s_marker in combined_text and b_marker in combined_text :
+                  s_text_part, b_text_part = combined_text.split(b_marker, 1)
+                  s_text = s_text_part.replace(s_marker, "").strip()
+                  b_text = b_text_part.strip()
+             elif s_marker in combined_text: 
+                  s_text = combined_text.replace(s_marker,"").strip()
+             else: 
+                  b_text = combined_text.strip()
+
+
+        subjects.append(simple_text_clean(s_text))
+        senders.append(simple_text_clean(d_text))
+        bodies.append(simple_text_clean(b_text))
+
+    data_for_lime = pd.DataFrame({
+        'subject': subjects,
+        'sender': senders,
+        'body': bodies
+    })
+    
+    try:
+        vectorized_input = nb_preprocessor.transform(data_for_lime)
+        probabilities = nb_model.predict_proba(vectorized_input)
+        return probabilities
+    except Exception as e:
+        print(f"Error in model_predict_probability_for_lime function during transform/predict: {e}")
+        return np.array([[0.5, 0.5]] * len(combined_texts))
+
+def get_prediction_and_explanation_nb(subject: str, sender: str, body: str):
+    if nb_preprocessor is None or nb_model is None:
+        return {"error": "Model/Preprocessor not loaded. Check server logs.", "prediction": "Error", "label": -1, "confidence": 0.0, "explanation": []}
+    
+    cleaned_subject = simple_text_clean(subject)
+    cleaned_sender = simple_text_clean(sender)
+    cleaned_body = simple_text_clean(body)
+
+    input_df_for_model = pd.DataFrame({
+        'subject': [cleaned_subject],
+        'sender': [cleaned_sender],
+        'body': [cleaned_body]
+        })
+
+    try:
+        vectorized_input = nb_preprocessor.transform(input_df_for_model)
+        prediction_label_int = nb_model.predict(vectorized_input)[0]
+        probabilities = nb_model.predict_proba(vectorized_input)[0]
+        
+        predicted_class_name = CLASS_NAMES[prediction_label_int]
+        confidence_score = probabilities[prediction_label_int]
+    except Exception as e:
+        return {"error": f"Prediction error: {e}", "prediction": "Error", 
+                "label": -1, "confidence": 0.0, "explanation": []}
+
+    text_for_lime = f"{cleaned_subject} : {cleaned_sender} : {cleaned_body}"
+
+    explanation_data = []
+    try:
+        exp = lime_explainer_nb.explain_instance(
+            text_instance=text_for_lime, 
+            classifier_fn=model_predict_probability_for_lime, 
+            num_features=15, 
+            top_labels=1,  
+            labels=(prediction_label_int,)
+        )
+        explanation_data = exp.as_list(label=prediction_label_int) 
+        print(f"LIME Explanation (Top 3): {explanation_data[:3]}")
+    except Exception as e:
+        print(f"LIME explanation error: {e}")
+        explanation_data = [("LIME explanation error or N/A", 0.0)]
+
+    return {
+        "prediction": predicted_class_name,
+        "label": int(prediction_label_int),
+        "confidence": float(confidence_score),
+        "explanation": explanation_data
+    }

app/ml_logic.py

@@ -5,6 +5,9 @@ from lime.lime_text import LimeTextExplainer
 import numpy as np
 import os
 
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+
 # Configure and setup model and preprocessor files
 ASSETS_DIR = os.path.join(os.path.dirname(__file__), 'assets')
 PREPROCESSOR_FILENAME = "email_preprocessor_20250506_203148.joblib" 
@@ -27,6 +30,26 @@ except Exception as e:
     preprocessor = None
     model = None
 
+# --- Load BERT-mini model and tokenizer from Hugging Face Hub ---
+# Replace with your actual Hugging Face model ID
+BERT_MODEL_ID = "lleratodev/720-bert-mini-phishing" # e.g., "LeratoLetsepe/phishing-bert-mini"
+try:
+    bert_tokenizer = AutoTokenizer.from_pretrained(BERT_MODEL_ID)
+    bert_model = AutoModelForSequenceClassification.from_pretrained(BERT_MODEL_ID)
+    bert_model.eval() # Set model to evaluation mode
+    print(f"BERT-mini model ('{BERT_MODEL_ID}') and tokenizer loaded successfully from Hugging Face Hub.")
+    # Determine device for BERT model (CPU by default, can be adapted for GPU)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    bert_model.to(device)
+    print(f"BERT model moved to device: {device}")
+
+except Exception as e:
+    print(f"FATAL ERROR (BERT): Could not load model/tokenizer '{BERT_MODEL_ID}' from Hugging Face Hub: {e}")
+    print("Ensure the model ID is correct, you have an internet connection, and the model files are correctly set up on the Hub.")
+    bert_tokenizer = None
+    bert_model = None
+# --- End BERT Loading ---
+
 # Text cleaning function, makes everything lowercase, removed non alpha-numeric characters and normalize white spaces
 def simple_text_clean(text):
     if isinstance(text, str):
@@ -37,7 +60,7 @@ def simple_text_clean(text):
         text = ''
     return text
 
-# For explanability, LIME setup
+# For explanability, LIME setup - # LIME probability function for MultinomialNB model
 class_names = ['Legitimate', 'Phishing'] # 0: Legitimate, 1: Phishing
 explainer = LimeTextExplainer(class_names=class_names)
 

requirements.txt

@@ -1,10 +1,14 @@
 fastapi
 uvicorn[standard]
-scikit-learn
+scikit-learn==1.5.1
+numpy==1.26.4
 pandas
 joblib
 scipy
-numpy
 lime
 python-multipart
-dill 
+dill 
+
+transformers
+transformers-interpret
+torch