train.py

"""
Model Training Script — run this once before starting the server.

Pattern: Template Method (Behavioral)
Each pipeline class extends MLPipeline, implementing only build_model() and save().
The training skeleton (load → preprocess → features → train → evaluate → save)
is inherited from the base class.

Usage:
    python train.py
    python train.py --models rf xgb
"""
import argparse
import sys
import os

# Ensure project root is on the path
sys.path.insert(0, os.path.dirname(__file__))

from sklearn.ensemble import RandomForestClassifier
import xgboost as xgb

from backend.infrastructure.ml.ml_pipeline import MLPipeline


# =============================================================================
# Concrete Pipeline 1: Random Forest
# =============================================================================
class RandomForestPipeline(MLPipeline):
    """Template Method — ConcreteClass for Random Forest training."""

    def build_model(self):
        return RandomForestClassifier(
            n_estimators=100,
            max_depth=20,
            n_jobs=-1,
            random_state=42,
            class_weight="balanced",
        )

    def save(self):
        self._save_model("rf_model.pkl")


# =============================================================================
# Concrete Pipeline 2: XGBoost
# =============================================================================
class XGBoostPipeline(MLPipeline):
    """Template Method — ConcreteClass for XGBoost training."""

    def build_model(self):
        return xgb.XGBClassifier(
            n_estimators=200,
            max_depth=8,
            learning_rate=0.1,
            use_label_encoder=False,
            eval_metric="mlogloss",
            random_state=42,
            n_jobs=-1,
        )

    def save(self):
        self._save_model("xgb_model.pkl")


# =============================================================================
# Main
# =============================================================================
PIPELINE_MAP = {
    "rf":  RandomForestPipeline,
    "xgb": XGBoostPipeline,
}

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train ransomware classification models.")
    parser.add_argument(
        "--models", nargs="+", choices=list(PIPELINE_MAP.keys()),
        default=list(PIPELINE_MAP.keys()),
        help="Which models to train (default: all)",
    )
    args = parser.parse_args()

    for key in args.models:
        pipeline = PIPELINE_MAP[key]()
        pipeline.run()

    print("\n Training complete. Models saved to models/")
    print("   Run: uvicorn backend.main:app --reload")