Spaces:
Sleeping
Sleeping
File size: 4,159 Bytes
e40ee72 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | """
Inference service for email threat classification:
Loads serialized ML pipelines, runs feature transformations, generates probabilities,
and assembles security intelligence reports.
"""
import os
import joblib
import pandas as pd
import numpy as np
from typing import Dict, Any
from src.data_loader import CLASS_NAMES, ID_TO_CLASS
from src.feature_engineering import extract_threat_indicators
from src.security_analyzer import assess_threat_level
class EmailThreatPredictor:
"""
High-performance production predictor for email threat classification.
"""
def __init__(self, model_path: str = None):
if model_path is None:
# Default location
default_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "models", "artifacts"))
model_path = os.path.join(default_dir, "best_threat_model.joblib")
self.model_path = model_path
self.pipeline = None
self._load_model()
def _load_model(self):
"""Loads serialized pipeline with safety checks."""
if not os.path.exists(self.model_path):
raise FileNotFoundError(
f"Model artifact not found at: {self.model_path}. "
"Please run train.py to build and save the model pipeline."
)
self.pipeline = joblib.load(self.model_path)
def predict(self, subject: str, body: str) -> Dict[str, Any]:
"""
Executes end-to-end inference on email subject and body.
Handles edge cases such as empty strings, whitespace, or missing inputs.
"""
clean_subj = str(subject or "").strip()
clean_body = str(body or "").strip()
# Handle purely empty input edge case
if not clean_subj and not clean_body:
return {
"predicted_class": "Legitimate",
"predicted_label_id": 0,
"probabilities": {"Legitimate": 0.90, "Spam": 0.05, "Suspicious": 0.05},
"confidence": 0.90,
"severity": "LOW RISK",
"badge_color": "#38a169",
"risk_score_out_of_10": 0.5,
"summary": "Empty message contains no malicious tokens or phishing indicators.",
"indicators": extract_threat_indicators("", ""),
"recommendations": ["No content detected. Enter email text to analyze."],
"warning": "Disclaimer: Research prototype designed for defense-in-depth assistance. Machine learning models cannot guarantee 100% email safety."
}
input_df = pd.DataFrame([{"subject": clean_subj, "body": clean_body}])
# Generate predictions
try:
probas = self.pipeline.predict_proba(input_df)[0]
pred_id = int(np.argmax(probas))
except Exception:
# Fallback if predict_proba is unavailable
pred_id = int(self.pipeline.predict(input_df)[0])
probas = [0.0, 0.0, 0.0]
probas[pred_id] = 1.0
predicted_class = ID_TO_CLASS[pred_id]
prob_dict = {
CLASS_NAMES[i]: round(float(probas[i]), 4)
for i in range(len(CLASS_NAMES))
}
# Extract indicators & assess security posture
indicators = extract_threat_indicators(clean_subj, clean_body)
assessment = assess_threat_level(predicted_class, prob_dict, indicators)
return {
"predicted_class": predicted_class,
"predicted_label_id": pred_id,
"probabilities": prob_dict,
"confidence": round(float(np.max(probas)), 4),
"severity": assessment["severity"],
"badge_color": assessment["badge_color"],
"risk_score_out_of_10": assessment["risk_score_out_of_10"],
"summary": assessment["summary"],
"indicators": indicators,
"recommendations": assessment["recommendations"],
"warning": "Disclaimer: Research prototype designed for defense-in-depth assistance. Machine learning models cannot guarantee 100% email safety. Always adhere to organizational security protocols."
}
|