Spaces:
Paused
Paused
| """ | |
| Implements: 03_AI/00_AI_ARCHITECTURE.md (Isolation Forest) | |
| """ | |
| import pandas as pd | |
| from sklearn.ensemble import IsolationForest | |
| import pickle | |
| import json | |
| import logging | |
| from pathlib import Path | |
| from datetime import datetime | |
| logger = logging.getLogger(__name__) | |
| class IsolationForestEngine: | |
| """ | |
| Unsupervised ML model for unknown anomaly detection using engineered features. | |
| """ | |
| def __init__(self, model_dir: str): | |
| self.model_dir = Path(model_dir) | |
| self.model = None | |
| self.version = "1.0.0" | |
| self.features_used = [ | |
| "hour_of_day", "day_of_week", "is_weekend", "is_working_hour", | |
| "is_failure", "country_encoded", "is_mfa", "time_since_last_login", | |
| "rolling_failures_24h" | |
| ] | |
| def train(self, features: pd.DataFrame): | |
| logger.info(f"Training Isolation Forest model with {len(features)} samples...") | |
| X = features[self.features_used].fillna(0) | |
| self.model = IsolationForest(n_estimators=100, contamination=0.05, random_state=42) | |
| self.model.fit(X) | |
| self.persist() | |
| def persist(self): | |
| import hashlib | |
| self.model_dir.mkdir(parents=True, exist_ok=True) | |
| version_id = datetime.now().strftime("%Y%m%d_%H%M%S") | |
| model_path = self.model_dir / f"trained_isolation_forest_{version_id}.pkl" | |
| meta_path = self.model_dir / f"model_metadata_{version_id}.json" | |
| with open(model_path, "wb") as f: | |
| pickle.dump(self.model, f) | |
| with open(model_path, "rb") as f: | |
| checksum = hashlib.sha256(f.read()).hexdigest() | |
| metadata = { | |
| "model_version": version_id, | |
| "training_dataset_version": "1.0", | |
| "feature_schema_version": "1.0", | |
| "training_date": datetime.now().isoformat(), | |
| "feature_schema": self.features_used, | |
| "hyperparameters": {"n_estimators": self.model.n_estimators, "contamination": self.model.contamination}, | |
| "performance_metrics": {"estimated_anomaly_rate": self.model.contamination}, | |
| "model_checksum": checksum, | |
| "model_file": model_path.name | |
| } | |
| with open(meta_path, "w") as f: | |
| json.dump(metadata, f) | |
| self.version = version_id | |
| logger.info(f"Isolation Forest model {version_id} persisted to {model_path}") | |
| def load(self) -> bool: | |
| meta_files = sorted(self.model_dir.glob("model_metadata_*.json")) | |
| if not meta_files: | |
| return False | |
| latest_meta = meta_files[-1] | |
| with open(latest_meta, "r") as f: | |
| metadata = json.load(f) | |
| model_path = self.model_dir / metadata["model_file"] | |
| if model_path.exists(): | |
| with open(model_path, "rb") as f: | |
| self.model = pickle.load(f) | |
| self.version = metadata["model_version"] | |
| logger.info(f"Pre-trained Isolation Forest model {self.version} loaded successfully.") | |
| return True | |
| return False | |
| def predict(self, features: pd.DataFrame) -> pd.DataFrame: | |
| if self.model is None: | |
| if not self.load(): | |
| logger.warning("No saved model found. Initiating dynamic training...") | |
| self.train(features) | |
| logger.info("Executing Isolation Forest inference...") | |
| X = features[self.features_used].fillna(0) | |
| preds = self.model.predict(X) | |
| scores = self.model.decision_function(X) | |
| normalized_scores = (-scores + 0.5) * 100 | |
| normalized_scores = normalized_scores.clip(min=0, max=100) | |
| return pd.DataFrame({ | |
| "event_id": features["event_id"], | |
| "if_score": normalized_scores, | |
| "if_prediction": (preds == -1).astype(int) | |
| }) | |