| """ |
| Script: train_deep.py |
| |
| Trains the Two-Tower neural recommender and pre-computes item embeddings. |
| |
| Outputs: |
| models/two_tower.pt (best checkpoint) |
| models/item_embeddings.npy (pre-computed item tower outputs) |
| models/item_index.pkl (item_idx → item metadata lookup) |
| data/outputs/figures/training_curve.png |
| |
| Usage: |
| python scripts/train_deep.py |
| python scripts/train_deep.py --debug |
| python scripts/train_deep.py --device cuda |
| """ |
|
|
| 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 numpy as np |
| import pandas as pd |
| import torch |
| from torch.utils.data import DataLoader |
|
|
| from src.config import FIGURES_DIR, MODELS_DIR, ensure_dirs, get_config |
| from src.data.dataset import BPRDataset |
| from src.models.deep import TwoTowerModel, TwoTowerTrainer, build_item_feature_matrix |
| from src.utils.logging import get_logger |
| from src.utils.seed import set_seed |
|
|
| logger = get_logger(__name__) |
|
|
| |
| ITEM_META_COLS = [ |
| "item_avg_rating", "item_num_ratings", "item_popularity", |
| "item_positive_rate", "price_norm", "category_idx", "brand_idx", |
| ] |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description="Train Two-Tower neural recommender.") |
| parser.add_argument("--debug", action="store_true") |
| parser.add_argument("--device", type=str, default=None, |
| help="Device: 'cpu', 'cuda', 'mps'. Auto-detected if not set.") |
| return parser.parse_args() |
|
|
|
|
| def detect_device(override: str | None) -> str: |
| if override: |
| return override |
| if torch.cuda.is_available(): |
| return "cuda" |
| if torch.backends.mps.is_available(): |
| return "mps" |
| return "cpu" |
|
|
|
|
| def plot_training_curve(history: list[dict], output_path: Path) -> None: |
| epochs = [h["epoch"] for h in history] |
| train_loss = [h["train_loss"] for h in history] |
| val_loss = [h["val_loss"] for h in history] |
|
|
| fig, ax = plt.subplots(figsize=(8, 4)) |
| ax.plot(epochs, train_loss, label="Train loss", linewidth=2, color="#3b82f6") |
| ax.plot(epochs, val_loss, label="Val loss", linewidth=2, color="#f97316", linestyle="--") |
| ax.set_xlabel("Epoch") |
| ax.set_ylabel("BPR Loss") |
| ax.set_title("Two-Tower Training Curve (BPR)") |
| ax.legend() |
| plt.tight_layout() |
| output_path.parent.mkdir(parents=True, exist_ok=True) |
| fig.savefig(output_path, dpi=150) |
| plt.close(fig) |
| logger.info("Training curve saved → %s", output_path) |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| cfg = get_config(debug=args.debug) |
| ensure_dirs(cfg) |
| set_seed(cfg.data.random_seed) |
|
|
| device = detect_device(args.device) |
| logger.info("Using device: %s", device) |
|
|
| logger.info("Loading data…") |
| train_df = pd.read_parquet(cfg.data.train_path) |
| val_df = pd.read_parquet(cfg.data.val_path) |
| item_features = pd.read_parquet(cfg.data.item_features_path) |
|
|
| |
| if "label" not in train_df.columns: |
| train_df["label"] = (train_df["rating"] >= cfg.data.positive_rating_threshold).astype(int) |
| if "label" not in val_df.columns: |
| val_df["label"] = (val_df["rating"] >= cfg.data.positive_rating_threshold).astype(int) |
|
|
| n_users = int(train_df["user_idx"].max() + 1) |
| n_items = int(item_features["item_idx"].max() + 1) |
| logger.info("Dataset: %d users, %d items", n_users, n_items) |
|
|
| |
| item_feat_matrix = build_item_feature_matrix(item_features, ITEM_META_COLS) |
| n_meta = item_feat_matrix.shape[1] |
| logger.info("Item metadata features: %d", n_meta) |
|
|
| |
| |
| train_ds = BPRDataset(train_df, item_feat_matrix) |
| val_ds = BPRDataset(val_df, item_feat_matrix, n_items=n_items) |
| logger.info("BPR dataset — train pairs: %d, val pairs: %d", len(train_ds), len(val_ds)) |
|
|
| train_loader = DataLoader( |
| train_ds, |
| batch_size=cfg.deep.batch_size, |
| shuffle=True, |
| num_workers=0, |
| pin_memory=(device == "cuda"), |
| ) |
| val_loader = DataLoader( |
| val_ds, |
| batch_size=cfg.deep.batch_size, |
| shuffle=False, |
| num_workers=0, |
| ) |
|
|
| logger.info("Building Two-Tower model…") |
| model = TwoTowerModel( |
| n_users=n_users, |
| n_items=n_items, |
| embedding_dim=cfg.deep.embedding_dim, |
| output_dim=cfg.deep.embedding_dim, |
| hidden_dims=cfg.deep.hidden_dims, |
| dropout=cfg.deep.dropout, |
| n_meta_features=n_meta, |
| ) |
| total_params = sum(p.numel() for p in model.parameters()) |
| logger.info("Model parameters: %d", total_params) |
|
|
| trainer = TwoTowerTrainer( |
| model=model, |
| device=device, |
| lr=cfg.deep.learning_rate, |
| weight_decay=cfg.deep.weight_decay, |
| patience=cfg.deep.patience, |
| model_path=cfg.deep.model_path, |
| ) |
|
|
| logger.info("Training for up to %d epochs…", cfg.deep.epochs) |
| history = trainer.fit(train_loader, val_loader, epochs=cfg.deep.epochs) |
|
|
| |
| plot_training_curve(history, FIGURES_DIR / "training_curve.png") |
|
|
| |
| logger.info("Pre-computing item embeddings…") |
| item_embeddings = trainer.build_item_embeddings( |
| n_items=n_items, |
| item_features=item_feat_matrix, |
| batch_size=2048, |
| ) |
| np.save(cfg.deep.item_embeddings_path, item_embeddings) |
| logger.info("Item embeddings saved → %s shape: %s", |
| cfg.deep.item_embeddings_path, item_embeddings.shape) |
|
|
| |
| item_index = item_features.set_index("item_idx")[["title", "category", "brand", "price"]].to_dict(orient="index") |
| with open(cfg.deep.item_index_path, "wb") as f: |
| pickle.dump(item_index, f) |
| logger.info("Item index saved → %s", cfg.deep.item_index_path) |
|
|
| logger.info("Two-Tower training complete.") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
