Spaces:
Sleeping
Sleeping
File size: 7,093 Bytes
38ae75d |
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 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 |
import pandas as pd
import numpy as np
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping, LearningRateMonitor
from pytorch_lightning.loggers import TensorBoardLogger
import torch
from utils import prepare_ground_truth, calculate_metrics
from data_prepare import prepare_data, SASRecDataset, SASRecDataModule
from models import recommend_popular_items_and_evaluate, recommend_item_item_and_evaluate, recommend_als_and_evaluate, SASRec
def train_and_eval_SASRec_model(train_set, validation_set, test_set, checkpoint_dir_path='checkpoints/',
checkpoint_path=None, n_epochs=10, mode='train',
batchsize=256, max_token_len=50, learning_rate=1e-3, hidden_dim=128,
num_heads=2, num_layers=2, dropout=0.2, weight_decay=1e-6):
"""
Train or evaluate a SASRec sequential recommendation model using PyTorch Lightning.
This function wraps the entire SASRec pipeline:
- Initializes the SASRecDataModule (handles dataset preprocessing and dataloaders).
- Builds the SASRec Transformer-based model.
- Configures training callbacks (checkpointing, early stopping, LR monitoring).
- Runs either training (`mode='train'`) or evaluation on the test set (`mode='test'`).
Args
----------
train_set : pd.DataFrame
Training interactions dataset .
validation_set : pd.DataFrame
Validation dataset with the same structure as `train_set`.
test_set : pd.DataFrame
Test dataset with the same structure as `train_set`.
checkpoint_dir_path : str, optional (default='checkpoints/')
Directory to save model checkpoints.
checkpoint_path : str or None, optional (default=None)
Path to a checkpoint file for resuming training or loading a pretrained model for testing.
n_epochs : int, optional (default=10)
Number of training epochs.
mode : {'train', 'test'}, optional (default='train')
- `'train'`: trains the model on the training/validation data.
- `'test'`: evaluates the model on the test set using a checkpoint.
batchsize : int, optional (default=256)
Batch size for training and evaluation.
max_token_len : int, optional (default=50)
Maximum sequence length per user (recent interactions kept).
learning_rate : float, optional (default=1e-3)
Learning rate for the AdamW optimizer.
hidden_dim : int, optional (default=128)
Dimensionality of item and positional embeddings.
num_heads : int, optional (default=2)
Number of attention heads in each Transformer encoder layer.
num_layers : int, optional (default=2)
Number of Transformer encoder layers.
dropout : float, optional (default=0.2)
Dropout probability applied in embeddings and Transformer layers.
weight_decay : float, optional (default=1e-6)
Weight decay regularization coefficient for AdamW.
"""
# --- 1. Initialize DataModule ---
print("Initializing DataModule...")
datamodule = SASRecDataModule(
train_df=train_set,
val_df=validation_set,
test_df=test_set,
batch_size=batchsize,
max_len=max_token_len
)
datamodule.setup()
# --- 2. Initialize Model ---
print("Initializing SASRec model...")
model = SASRec(
vocab_size=datamodule.vocab_size,
max_len=max_token_len,
hidden_dim=hidden_dim,
num_heads=num_heads,
num_layers=num_layers,
dropout=dropout,
learning_rate=learning_rate,
weight_decay=weight_decay
)
# --- 3. Configure Training Callbacks ---
checkpoint_callback = ModelCheckpoint(
dirpath=checkpoint_dir_path,
filename="sasrec-{epoch:02d}-{val_hitrate@10:.4f}",
save_top_k=1,
verbose=True,
monitor="val_hitrate@10",
mode="max"
)
early_stopping_callback = EarlyStopping(
monitor="val_hitrate@10", # stop if ranking metric stagnates
patience=5,
mode="max"
)
lr_monitor = LearningRateMonitor(logging_interval="step")
logger = TensorBoardLogger("lightning_logs", name="sasrec")
# --- 4. Initialize Trainer ---
print("Initializing PyTorch Lightning Trainer...")
trainer = pl.Trainer(
logger=logger,
callbacks=[checkpoint_callback, early_stopping_callback, lr_monitor],
max_epochs=n_epochs,
accelerator='auto',
devices=1,
gradient_clip_val=1.0, # helps with exploding gradients
)
if mode == 'train' :
# --- 5. Start Training ---
print(f"Starting training for up to {n_epochs} epochs...")
trainer.fit(model, datamodule,
ckpt_path=checkpoint_path
)
elif mode == 'test':
# --- 6. Test on best checkpoint ---
print("Evaluating on test set...")
trainer.test(model, datamodule,
ckpt_path=checkpoint_path
)
# --- Main Execution Block ---
if __name__ == "__main__":
# --- Configuration ---
BATCH_SIZE = 256
MAX_TOKEN_LEN = 50 # 50–100 is standard
LEARNING_RATE = 1e-3
HIDDEN_DIM = 128
NUM_HEADS = 2
NUM_LAYERS = 2
DROPOUT = 0.2
WEIGHT_DECAY = 1e-6
N_EPOCHS = 50
CHECKPOINT_SAVE_PATH = 'checkpoints/'
CHECKPOINT_LOAD_PATH = None # or specify a path to a checkpoint file
MODE = 'train' # 'train' or 'test'
train_set, validation_set, test_set = prepare_data(data_folder='data/')
if train_set is not None:
results = {}
full_train_set = pd.concat([train_set, validation_set])
# Evaluate classical models
print("\n>>> Running evaluations on the VALIDATION set <<<")
results['Popularity (Validation)'] = recommend_popular_items_and_evaluate(train_set, validation_set)
results['Item-Item CF (Validation)'] = recommend_item_item_and_evaluate(train_set, validation_set)
results['ALS (Validation)'] = recommend_als_and_evaluate(train_set, validation_set)
print("\n>>> Running final evaluations on the TEST set <<<")
results['Popularity (Test)'] = recommend_popular_items_and_evaluate(full_train_set, test_set)
results['Item-Item CF (Test)'] = recommend_item_item_and_evaluate(full_train_set, test_set)
results['ALS (Test)'] = recommend_als_and_evaluate(full_train_set, test_set)
print("\n--- Final Evaluation Results ---")
results_df = pd.DataFrame.from_dict(results, orient='index')
print(results_df)
print("--------------------------------")
# Train and evaluate SASRec model
print("\n>>> Training and evaluating SASRec model <<<")
train_and_eval_SASRec_model(train_set, validation_set, test_set, n_epochs=10, mode='train')
print("\n>>> Evaluating trained SASRec model on TEST set <<<")
train_and_eval_SASRec_model(train_set, validation_set, test_set, mode='test') |