urlmodel_info/train_pipeline.py

"""
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()