File size: 4,987 Bytes
fa98c11 | 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 | """
Script: train_classical.py
Trains the LightGBM reranker on the full feature set and saves the model.
Outputs:
models/lgbm_model.pkl
data/outputs/figures/feature_importance.png
Usage:
python scripts/train_classical.py
python scripts/train_classical.py --debug
python scripts/train_classical.py --feature-set id_metadata # ablation tier
"""
from __future__ import annotations
import argparse
import pickle
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import matplotlib.pyplot as plt
import pandas as pd
from src.config import FIGURES_DIR, MODELS_DIR, PROCESSED_DIR, ensure_dirs, get_config
from src.features.builder import FEATURE_SETS, build_feature_matrix
from src.models.classical import LightGBMReranker
from src.utils.logging import get_logger
logger = get_logger(__name__)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Train LightGBM reranker.")
parser.add_argument("--debug", action="store_true")
parser.add_argument(
"--feature-set",
type=str,
default="id_metadata_text_history",
choices=FEATURE_SETS,
help="Feature tier to use for training.",
)
parser.add_argument(
"--output-name",
type=str,
default="lgbm_model",
help="Stem of the output .pkl filename (in models/).",
)
return parser.parse_args()
def plot_feature_importance(model: LightGBMReranker, output_path: Path) -> None:
"""Save a horizontal bar chart of feature importances."""
imp = model.feature_importance()
top = imp.head(20)
fig, ax = plt.subplots(figsize=(9, 6))
ax.barh(top["feature"][::-1], top["importance"][::-1], color="#3b82f6")
ax.set_xlabel("Importance (Gain)")
ax.set_title("LightGBM Feature Importance (Top 20)")
plt.tight_layout()
output_path.parent.mkdir(parents=True, exist_ok=True)
fig.savefig(output_path, dpi=150)
plt.close(fig)
logger.info("Feature importance plot saved → %s", output_path)
def main() -> None:
args = parse_args()
cfg = get_config(debug=args.debug)
ensure_dirs(cfg)
logger.info("Loading pre-built feature matrices…")
# Try to load pre-built features; fall back to rebuilding
if cfg.data.train_features_path.exists():
train_feat = pd.read_parquet(cfg.data.train_features_path)
val_feat = pd.read_parquet(cfg.data.val_features_path)
else:
logger.warning("Pre-built features not found. Loading raw splits and rebuilding…")
train_df = pd.read_parquet(cfg.data.train_path)
val_df = pd.read_parquet(cfg.data.val_path)
user_features = pd.read_parquet(cfg.data.user_features_path)
item_features = pd.read_parquet(cfg.data.item_features_path)
artifacts_path = PROCESSED_DIR / "feature_artifacts.pkl"
with open(artifacts_path, "rb") as f:
artifacts = pickle.load(f)
user_profiles_path = PROCESSED_DIR / "user_profiles.parquet"
if user_profiles_path.exists():
user_profiles = pd.read_parquet(user_profiles_path)
else:
from src.features.text_features import build_user_text_profiles
user_profiles = build_user_text_profiles(train_df, item_features)
train_feat, _ = build_feature_matrix(
train_df, user_features, item_features, user_profiles,
artifacts["vectorizer"], feature_set=args.feature_set,
)
val_feat, _ = build_feature_matrix(
val_df, user_features, item_features, user_profiles,
artifacts["vectorizer"], feature_set=args.feature_set,
)
# Determine feature columns (everything except 'label' and index cols)
exclude = {"label", "user_idx", "item_idx", "user_id", "item_id", "rating", "timestamp"}
feature_cols = [c for c in train_feat.columns if c not in exclude]
logger.info("Training LightGBM on %d features, %d rows…", len(feature_cols), len(train_feat))
model = LightGBMReranker(
n_estimators=cfg.classical.n_estimators,
learning_rate=cfg.classical.learning_rate,
num_leaves=cfg.classical.num_leaves,
max_depth=cfg.classical.max_depth,
min_child_samples=cfg.classical.min_child_samples,
subsample=cfg.classical.subsample,
colsample_bytree=cfg.classical.colsample_bytree,
reg_alpha=cfg.classical.reg_alpha,
reg_lambda=cfg.classical.reg_lambda,
n_jobs=cfg.classical.n_jobs,
random_state=cfg.data.random_seed,
early_stopping_rounds=cfg.classical.early_stopping_rounds,
)
model.fit(train_feat, val_feat, feature_cols)
# Save model
model_path = MODELS_DIR / f"{args.output_name}.pkl"
model.save(model_path)
# Feature importance plot
plot_feature_importance(model, FIGURES_DIR / "feature_importance.png")
logger.info("LightGBM training complete.")
if __name__ == "__main__":
main()
|