Spaces:
Sleeping
Sleeping
| """ | |
| Model architecture construction, pipeline assembly, and training orchestration. | |
| Integrates text vectorization, character n-grams, and cyber threat indicators. | |
| """ | |
| import time | |
| import joblib | |
| from typing import Dict, Any, Tuple | |
| from sklearn.pipeline import Pipeline, FeatureUnion | |
| from sklearn.feature_extraction.text import TfidfVectorizer | |
| from sklearn.preprocessing import StandardScaler | |
| from sklearn.dummy import DummyClassifier | |
| from sklearn.linear_model import LogisticRegression | |
| from xgboost import XGBClassifier | |
| from src.feature_engineering import ThreatFeatureExtractor, EmailTextCombiner | |
| def build_pipeline(model_type: str = "logistic_regression", params: Dict[str, Any] = None) -> Pipeline: | |
| """ | |
| Constructs a complete end-to-end scikit-learn pipeline combining: | |
| 1. Word-level TF-IDF (1, 2 n-grams) | |
| 2. Character-level TF-IDF (3, 5 n-grams for obfuscation detection) | |
| 3. Structural & behavioral threat features (urgency, credential, url counts) | |
| 4. Selected classification algorithm | |
| """ | |
| if params is None: | |
| params = {} | |
| word_max_features = params.get("word_max_features", 4000) | |
| char_max_features = params.get("char_max_features", 2000) | |
| # Word-level TF-IDF pipeline | |
| word_tfidf_pipeline = Pipeline([ | |
| ("combiner", EmailTextCombiner(subject_weight=2)), | |
| ("word_tfidf", TfidfVectorizer( | |
| ngram_range=(1, 2), | |
| max_features=word_max_features, | |
| sublinear_tf=True, | |
| strip_accents="unicode" | |
| )) | |
| ]) | |
| # Character-level TF-IDF pipeline for lexical obfuscation/typosquatting | |
| char_tfidf_pipeline = Pipeline([ | |
| ("combiner", EmailTextCombiner(subject_weight=1)), | |
| ("char_tfidf", TfidfVectorizer( | |
| analyzer="char", | |
| ngram_range=(3, 5), | |
| max_features=char_max_features, | |
| sublinear_tf=True | |
| )) | |
| ]) | |
| # Cyber domain behavioral feature extraction | |
| threat_features_pipeline = Pipeline([ | |
| ("extractor", ThreatFeatureExtractor()), | |
| ("scaler", StandardScaler()) | |
| ]) | |
| # Unified feature representation | |
| feature_union = FeatureUnion([ | |
| ("word_features", word_tfidf_pipeline), | |
| ("char_features", char_tfidf_pipeline), | |
| ("threat_indicators", threat_features_pipeline) | |
| ]) | |
| # Instantiate chosen classifier | |
| if model_type == "dummy": | |
| clf = DummyClassifier( | |
| strategy=params.get("strategy", "stratified"), | |
| random_state=params.get("random_state", 42) | |
| ) | |
| elif model_type == "logistic_regression": | |
| clf = LogisticRegression( | |
| C=params.get("C", 1.0), | |
| max_iter=params.get("max_iter", 1000), | |
| class_weight=params.get("class_weight", "balanced"), | |
| solver="lbfgs", | |
| random_state=params.get("random_state", 42) | |
| ) | |
| elif model_type == "logistic_regression_security_tuned": | |
| # Tuned with higher penalty weighting on Suspicious (Class 2) for zero-trust defense | |
| custom_weights = params.get("class_weight", {0: 1.0, 1: 1.0, 2: 1.4}) | |
| clf = LogisticRegression( | |
| C=params.get("C", 2.5), | |
| max_iter=params.get("max_iter", 1500), | |
| class_weight=custom_weights, | |
| solver="lbfgs", | |
| random_state=params.get("random_state", 42) | |
| ) | |
| elif model_type == "xgboost": | |
| clf = XGBClassifier( | |
| n_estimators=params.get("n_estimators", 100), | |
| max_depth=params.get("max_depth", 5), | |
| learning_rate=params.get("learning_rate", 0.1), | |
| subsample=params.get("subsample", 0.85), | |
| colsample_bytree=params.get("colsample_bytree", 0.85), | |
| eval_metric="mlogloss", | |
| random_state=params.get("random_state", 42), | |
| n_jobs=4 | |
| ) | |
| else: | |
| raise ValueError(f"Unsupported model_type: '{model_type}'") | |
| pipeline = Pipeline([ | |
| ("features", feature_union), | |
| ("classifier", clf) | |
| ]) | |
| return pipeline | |
| def train_model( | |
| model_type: str, | |
| train_df, | |
| params: Dict[str, Any] = None | |
| ) -> Tuple[Pipeline, float]: | |
| """ | |
| Fits the complete pipeline and measures training execution latency. | |
| """ | |
| pipeline = build_pipeline(model_type, params) | |
| start_time = time.time() | |
| pipeline.fit(train_df, train_df["label_id"].values) | |
| duration = time.time() - start_time | |
| return pipeline, duration | |
| def save_pipeline(pipeline: Pipeline, output_path: str) -> str: | |
| """ | |
| Serializes complete fitted pipeline to file using joblib. | |
| """ | |
| import os | |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) | |
| joblib.dump(pipeline, output_path) | |
| return output_path | |
| def load_pipeline(pipeline_path: str) -> Pipeline: | |
| """ | |
| Loads saved pipeline artifact. | |
| """ | |
| return joblib.load(pipeline_path) | |