train_model.py

import numpy as np
import xgboost as xgb
from sklearn.metrics import accuracy_score, classification_report
import pickle
import os

data_dir = r"C:\Users\nlaks\OneDrive\Desktop\binary-risk-profiler\data\ember2018"

print("Loading dataset...")
X_train = np.load(os.path.join(data_dir, "X_train.npy"))
y_train = np.load(os.path.join(data_dir, "y_train.npy"))
X_test = np.load(os.path.join(data_dir, "X_test.npy"))
y_test = np.load(os.path.join(data_dir, "y_test.npy"))

print(f"Train: {X_train.shape}, Test: {X_test.shape}")

print("\nTraining XGBoost...")
model = xgb.XGBClassifier(
    n_estimators=300,
    max_depth=8,
    learning_rate=0.05,
    subsample=0.8,
    colsample_bytree=0.8,
    eval_metric='logloss',
    n_jobs=-1,
    random_state=42
)

model.fit(X_train, y_train)

print("\nEvaluating...")
y_pred = model.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
print(f"\nAccuracy: {accuracy:.4f}")
print(classification_report(y_test, y_pred, target_names=["Benign", "Malicious"]))

os.makedirs("model", exist_ok=True)
with open("model/xgb_model.pkl", "wb") as f:
    pickle.dump(model, f)

print("Model saved to model/xgb_model.pkl")