Spaces:
Sleeping
Sleeping
| """ | |
| 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." | |
| } | |