""" MalwareGuard — URL Model Retraining Pipeline (CPU Optimized) Performs memory-safe chunked feature extraction, scaling, and trains the Stacking Ensemble classifier on the 2.5 Million URL Mega-Corpus. """ import os import sys import time import joblib import warnings import numpy as np import pandas as pd from sklearn.preprocessing import StandardScaler from sklearn.model_selection import train_test_split from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score # Ensure correct workspace path is in python path sys.path.append(os.path.dirname(os.path.abspath(__file__))) from feature_engineering import URLFeatureExtractor from models import build_xgb, build_lgbm, build_tfidf_lr_pipeline, ManualStackingEnsemble # Configurations DATA_DIR = "data" CSV_PATH = os.path.join(DATA_DIR, "url_megacorpus_2.5m.csv") MODELS_DIR = "urlmodel_info" SCALER_PATH = os.path.join(MODELS_DIR, "scaler.joblib") XGB_PATH = os.path.join(MODELS_DIR, "xgb_binary.joblib") LGBM_PATH = os.path.join(MODELS_DIR, "lgbm_binary.joblib") TFIDF_PATH = os.path.join(MODELS_DIR, "tfidf_lr_binary.joblib") META_PATH = os.path.join(MODELS_DIR, "hypermodel_meta.joblib") warnings.simplefilter("ignore") def check_dataset_exists(): """Verify that compiled CSV is ready.""" if not os.path.exists(CSV_PATH): print(f"[ERROR] Compiled dataset not found at {CSV_PATH}.") print("Please ensure the compilation script has finished running successfully first.") sys.exit(1) print(f"[INFO] Found compiled dataset at {CSV_PATH}") def featurize_in_chunks(csv_path, chunk_size=200000, n_jobs=-1): """Memory-safe featurization in chunks to prevent memory/RAM crashes.""" print("[INFO] Starting chunked feature extraction...") extractor = URLFeatureExtractor(n_jobs=n_jobs) # Pre-determine lines for accurate reporting total_rows = sum(1 for _ in open(csv_path, "r", encoding="utf-8")) - 1 print(f"Total rows to process: {total_rows:,}") chunks = [] chunk_idx = 0 start_time = time.time() for chunk in pd.read_csv(csv_path, chunksize=chunk_size): chunk_start = time.time() print(f"\n--- Featurizing Chunk {chunk_idx + 1} ({len(chunk):,} rows) ---") # Extract features feat_df = extractor.transform(chunk["url"], verbose=True) # Add labels feat_df.insert(0, "url", chunk["url"].values) feat_df["label"] = chunk["label"].values.astype(np.int8) chunks.append(feat_df) elapsed = time.time() - chunk_start print(f"Chunk {chunk_idx + 1} completed in {elapsed:.2f} seconds") chunk_idx += 1 print("\n[INFO] Concatenating featurized chunks...") full_feat_df = pd.concat(chunks, ignore_index=True) total_elapsed = time.time() - start_time print(f"[SUCCESS] Featurization complete! Total time: {total_elapsed/60:.2f} minutes") print(f"Feature matrix shape: {full_feat_df.shape[0]:,} rows × {full_feat_df.shape[1]} columns") return full_feat_df def train_pipeline(): """Execute the full model training and serialization pipeline.""" check_dataset_exists() # ── 1. Feature Extraction ────────────────────────────────────────────── feat_df = featurize_in_chunks(CSV_PATH) # Extract arrays X_url = feat_df["url"].values.astype(str) y = feat_df["label"].values.astype(np.int8) # Drop non-feature columns for scaling feature_cols = [c for c in feat_df.columns if c not in ("url", "label")] X_raw_feat = feat_df[feature_cols].values.astype(np.float32) # Clear df from memory to free up RAM del feat_df import gc gc.collect() # ── 2. Train / Test Split ────────────────────────────────────────────── print("\n[INFO] Splitting into Train and Test sets (80/20)...") X_feat_tr, X_feat_te, X_url_tr, X_url_te, y_tr, y_te = train_test_split( X_raw_feat, X_url, y, test_size=0.2, random_state=42, stratify=y ) # ── 3. Standardize Features ──────────────────────────────────────────── print("[INFO] Scaling features with StandardScaler...") scaler = StandardScaler() X_feat_tr_scaled = scaler.fit_transform(X_feat_tr).astype(np.float32) X_feat_te_scaled = scaler.transform(X_feat_te).astype(np.float32) # Save the scaler joblib.dump(scaler, SCALER_PATH) print(f"[SUCCESS] Scaler saved to {SCALER_PATH} ({scaler.n_features_in_} features)") # ── 4. Setup Stacking Ensemble ───────────────────────────────────────── print("\n[INFO] Configuring CPU-optimized base models...") # Speed optimized XGBoost (using hist tree method on CPU) xgb_base = build_xgb( num_classes=2, n_estimators=100, # Optimized for training speed max_depth=6, # Optimized down from 7 learning_rate=0.1, # Slightly higher learning rate for faster convergence tree_method="hist", device="cpu", n_jobs=-1 ) # Speed optimized LightGBM lgbm_base = build_lgbm( num_classes=2, n_estimators=100, # Optimized for speed max_depth=6, num_leaves=31, # Optimized learning_rate=0.1, device_type="cpu", n_jobs=-1 ) # Meta learner XGBoost meta_learner = build_xgb( num_classes=2, n_estimators=100, max_depth=3, learning_rate=0.1, tree_method="hist", device="cpu", n_jobs=-1 ) level0_models = { "xgb": xgb_base, "lgbm": lgbm_base } # Manual Stacking Ensemble with 2-fold cross validation for speed stacking_ensemble = ManualStackingEnsemble( level0=level0_models, meta=meta_learner, cv=2, num_classes=2 ) # ── 5. Stacking Ensemble Training ────────────────────────────────────── print("\n" + "="*50) print("TRAINING STACKING ENSEMBLE (3-Fold CV + Full Fit)...") print("="*50) start_train = time.time() stacking_ensemble.fit(X_feat_tr_scaled, X_url_tr, y_tr) train_duration = time.time() - start_train print(f"[SUCCESS] Training completed in {train_duration/60:.2f} minutes!") # ── 6. Serialize and Save Stacking Base and Meta Models ───────────────── print("\n[INFO] Serializing and saving all fitted base models...") joblib.dump(stacking_ensemble.fitted_l0_["xgb"], XGB_PATH) print(f" Saved XGBoost Base Model to {XGB_PATH}") joblib.dump(stacking_ensemble.fitted_l0_["lgbm"], LGBM_PATH) print(f" Saved LightGBM Base Model to {LGBM_PATH}") # Save the fitted Meta-Learner (hypermodel_meta.joblib) joblib.dump(stacking_ensemble.meta, META_PATH) print(f"[SUCCESS] Saved Stacking Meta-Learner to {META_PATH}") # ── 7. Evaluation ────────────────────────────────────────────────────── print("\n" + "="*50) print("EVALUATING ON TEST SET (20% Split)...") print("="*50) # Predict probabilities and classes p_meta = stacking_ensemble.predict_proba(X_feat_te_scaled, X_url_te) y_pred_score = p_meta[:, 1] y_pred = (y_pred_score > 0.5).astype(int) # Calculate metrics acc = accuracy_score(y_te, y_pred) prec = precision_score(y_te, y_pred) rec = recall_score(y_te, y_pred) f1 = f1_score(y_te, y_pred) auc = roc_auc_score(y_te, y_pred_score) print(f"Accuracy: {acc:.5f}") print(f"Precision: {prec:.5f}") print(f"Recall: {rec:.5f}") print(f"F1-Score: {f1:.5f}") print(f"AUC Score: {auc:.5f}") # Print Individual Model Scores on Test Set print("\nLevel-0 Estimator Performance on Test Set:") for name, base_model in stacking_ensemble.fitted_l0_.items(): X_input = X_feat_te_scaled p_base = base_model.predict_proba(X_input) y_pred_base = (p_base[:, 1] > 0.5).astype(int) base_acc = accuracy_score(y_te, y_pred_base) base_f1 = f1_score(y_te, y_pred_base) print(f" - {name:<10} | Accuracy: {base_acc:.5f} | F1: {base_f1:.5f}") print("\n[SUCCESS] ALL MODELS SUCCESSFULLY TRAINED AND DEPLOYED IN THE LOCAL WORKSPACE!") if __name__ == "__main__": train_pipeline()