Spaces:
Sleeping
Sleeping
File size: 4,854 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 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | """
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)
|