Spaces:
Paused
Paused
File size: 2,484 Bytes
5a3b9db | 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 | import pandas as pd
import time
import logging
from datetime import datetime
from pathlib import Path
from src.ai.risk.fusion_engine import RiskFusionEngine
from src.ai.classification.xgboost_classifier import XGBoostClassifier
logger = logging.getLogger(__name__)
class Phase4Pipeline:
"""
Module 4: Classification Pipeline.
Orchestrates Enterprise Risk Fusion and Attack Classification.
Output conforms to Documented Output Contract.
"""
def __init__(self, config: dict, model_dir: str, output_dir: str):
self.config = config
self.model_dir = Path(model_dir)
self.output_dir = Path(output_dir)
self.risk_engine = RiskFusionEngine(config)
self.classifier = XGBoostClassifier(str(model_dir), config)
def execute(self, anomaly_scores_path: str, features_path: str) -> pd.DataFrame:
start_time = time.time()
logger.info("Starting Phase 4 Classification Pipeline...")
# 1. Load Data
anomalies_df = pd.read_parquet(anomaly_scores_path)
features_df = pd.read_parquet(features_path)
# 2. Execute Enterprise Risk Fusion
risk_df = self.risk_engine.calculate_risk(anomalies_df)
# 3. Execute Attack Classification
if not self.classifier.load():
logger.warning("No pre-trained classification model found. Training dummy model for execution...")
self.classifier.train(features_df, features_df["attack_label"])
class_df = self.classifier.predict(features_df)
# 4. Merge Results & Ensure Output Contract
result_df = risk_df.merge(class_df, on="event_id")
# Add Traceability Metadata
det_engine_version = anomalies_df.get("model_version", pd.Series(["1.0"] * len(anomalies_df)))
result_df["detection_engine_version"] = det_engine_version
result_df["feature_schema_version"] = self.config.get("feature_schema_version", "1.0.0")
result_df["processing_timestamp"] = datetime.now().isoformat()
result_df["processing_duration_sec"] = time.time() - start_time
# Save output
self.output_dir.mkdir(parents=True, exist_ok=True)
out_path = self.output_dir / "risk_scores.parquet"
result_df.to_parquet(out_path, index=False)
logger.info(f"Phase 4 Pipeline completed successfully. Output persisted to {out_path}")
return result_df
|