{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.12.12"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":4120,"databundleVersionId":44256,"isSourceIdPinned":false}],"dockerImageVersionId":31260,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":26969.495825,"end_time":"2026-02-11T18:21:39.089215","environment_variables":{},"exception":true,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2026-02-11T10:52:09.593390","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# CTR Prediction using different ML models\n\n**Author:** Arya Pathak \n**Dataset:** Avazu CTR Prediction (10M sampled impressions) \n**Environment:** Kaggle Notebooks (GPU-accelerated) \n\n---\n\n## Overview\n\nThis notebook compares click-through rate prediction models, progressing from logistic regression through factorization machines to attention-based deep learning. The key finding is that **selective feature interaction modeling (AFM) outperforms more complex approaches**.\n\n**References:**\n\nData processing approach inspired by the following notebooks:\n- https://www.kaggle.com/code/avyaya/mlfinal\n- https://www.kaggle.com/code/sharanharsoor/ctr-analysis-of-different-ml-models\n\nExample of DeepFM implementation: https://medium.com/@ngiengkianyew/deconstructing-the-deepfm-22e64a803c33\n\n## Contents\n1. Setup\n2. Model Definitions\n3. Data Loading and Preprocessing\n4. Training Infrastructure\n5. Hyperparameter Tuning\n6. Experiments and Results\n\n---\n","metadata":{"papermill":{"duration":0.007927,"end_time":"2026-02-11T10:52:13.065546","exception":false,"start_time":"2026-02-11T10:52:13.057619","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"## 1. Setup","metadata":{"papermill":{"duration":0.00623,"end_time":"2026-02-11T10:52:13.078467","exception":false,"start_time":"2026-02-11T10:52:13.072237","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# required libraries\nimport numpy as np \nimport pandas as pd \nimport random\nimport seaborn as sns\nimport matplotlib.pyplot as plt\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.preprocessing import LabelEncoder\nfrom sklearn.metrics import roc_auc_score, log_loss, roc_curve, RocCurveDisplay\nimport torch \nimport torch.nn as nn \nimport torch.optim as optim \nfrom torch.utils.data import DataLoader, TensorDataset\nfrom dataclasses import dataclass, asdict\n\n# New requirements\nimport polars as pl\nimport math\nfrom rich.console import Console\nfrom rich.live import Live\nfrom rich.table import Table\nfrom rich.panel import Panel\nfrom rich.progress import Progress, BarColumn, TextColumn, TimeElapsedColumn, TimeRemainingColumn, SpinnerColumn\nfrom rich.layout import Layout\nfrom rich.align import Align\nfrom rich.text import Text\nimport time\nimport json\nfrom tqdm.notebook import tqdm\nfrom IPython.display import clear_output, display\nfrom sklearn.compose import ColumnTransformer\nimport polars.selectors as cs\nimport os\nimport itertools\nfrom dataclasses import dataclass\nfrom datetime import datetime\nfrom torch.utils.tensorboard import SummaryWriter\nfrom typing import Dict, Any, Optional, List, Tuple, Callable","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:52:13.092740Z","iopub.status.busy":"2026-02-11T10:52:13.092421Z","iopub.status.idle":"2026-02-11T10:52:43.562328Z","shell.execute_reply":"2026-02-11T10:52:43.561711Z"},"papermill":{"duration":30.479358,"end_time":"2026-02-11T10:52:43.564072","exception":false,"start_time":"2026-02-11T10:52:13.084714","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2. Model Definitions\n\nThis section defines the neural network architectures for CTR prediction.\n\n**Implementation reference:** [Deconstructing DeepFM](https://medium.com/@ngiengkianyew/deconstructing-the-deepfm-22e64a803c33)","metadata":{"papermill":{"duration":0.006466,"end_time":"2026-02-11T10:52:43.578203","exception":false,"start_time":"2026-02-11T10:52:43.571737","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"### 2.1 DeepFM","metadata":{"papermill":{"duration":0.006259,"end_time":"2026-02-11T10:52:43.590740","exception":false,"start_time":"2026-02-11T10:52:43.584481","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# DeepFM: Combines FM for pairwise interactions with DNN for higher-order interactions\n# For simplicity, we will only embed each column to 10 dimensions (each column will be represented by 10 dimensional vectors)\nclass DeepFM(nn.Module):\n def __init__(self,field_dims, embed_dims=10, mlp_dims = [400,400,400], drop_rate = 0.1):\n '''\n field_dims: a list containing how many unique amount in each columns\n drop_rate: can be finetuned\n '''\n super().__init__()\n self.num_fields = len(field_dims)\n self.embed_dims = embed_dims\n self.embeddings = nn.ModuleList([\n nn.Embedding(dim, embed_dims) for dim in field_dims\n ])\n self.linear_weights = nn.ModuleList([\n nn.Embedding(num_embeddings=dim, embedding_dim=1)\n for dim in field_dims\n ])\n self.bias = nn.Parameter(torch.zeros(1))\n mlp_layers = []\n input_dim = self.num_fields * embed_dims\n for mlp_dim in mlp_dims:\n mlp_layers.extend([\n nn.Linear(input_dim, mlp_dim),\n nn.BatchNorm1d(mlp_dim),\n nn.ReLU(),\n nn.Dropout(drop_rate)\n ])\n input_dim = mlp_dim\n\n mlp_layers.append(nn.Linear(mlp_dims[-1], 1))\n \n self.dnn = nn.Sequential(*mlp_layers)\n \n def forward(self,x):\n embeds = torch.stack([emb(x[:, i]) for i, emb in enumerate(self.embeddings)], dim=1)\n #Linear Part\n linear_terms = torch.cat([\n self.linear_weights[i](x[:, i]) # (batch, 1)\n for i in range(self.num_fields)\n ], dim=1)\n linear_output = torch.sum(linear_terms, dim=1)\n #FM Part\n square_of_sum = torch.sum(embeds, axis=1).pow(2)\n sum_of_squares = torch.sum(embeds.pow(2), axis=1)\n cross_term = square_of_sum - sum_of_squares\n fm_output = 0.5 * torch.sum(cross_term, dim=-1)\n #DNN Part\n embeds_flat = embeds.view(embeds.size(0), -1)\n dnn_output = self.dnn(embeds_flat).squeeze()\n #Total\n x = self.bias + linear_output + fm_output + dnn_output\n prediction = torch.sigmoid(x)\n return prediction","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:52:43.604908Z","iopub.status.busy":"2026-02-11T10:52:43.604400Z","iopub.status.idle":"2026-02-11T10:52:43.613543Z","shell.execute_reply":"2026-02-11T10:52:43.612975Z"},"papermill":{"duration":0.018027,"end_time":"2026-02-11T10:52:43.614952","exception":false,"start_time":"2026-02-11T10:52:43.596925","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 2.2 AFM (Attentional Factorization Machine)","metadata":{"papermill":{"duration":0.006119,"end_time":"2026-02-11T10:52:43.627503","exception":false,"start_time":"2026-02-11T10:52:43.621384","status":"completed"},"tags":[]}},{"cell_type":"code","source":"class AFM(nn.Module):\n def __init__(self, field_dims, embed_dims=10, attention_dims=20):\n super().__init__()\n self.num_fields = len(field_dims)\n self.embed_dims = embed_dims\n #embedding layers\n self.embeddings = nn.ModuleList([\n nn.Embedding(dim, embed_dims) for dim in field_dims\n ])\n #linear layers\n self.linear_weights = nn.ModuleList([\n nn.Embedding(num_embeddings=dim, embedding_dim=1)\n for dim in field_dims\n ])\n self.bias = nn.Parameter(torch.zeros(1))\n #attention layers\n self.attention_W = nn.Linear(embed_dims, attention_dims, bias=True)\n self.attention_h = nn.Linear(attention_dims, 1, bias=False)\n #projection to scalar\n self.projection = nn.Linear(embed_dims, 1, bias=False)\n\n def forward(self,x):\n embeds = torch.stack([emb(x[:, i]) for i, emb in enumerate(self.embeddings)], dim=1)\n # Linear Part\n linear_terms = torch.cat([\n self.linear_weights[i](x[:, i]) # (batch, 1)\n for i in range(self.num_fields)\n ], dim=1)\n linear_output = torch.sum(linear_terms, dim=1)\n \n # Interaction Layer\n interactions = []\n for i in range(self.num_fields):\n for j in range(i + 1, self.num_fields):\n # Element-wise product\n interaction = embeds[:, i] * embeds[:, j]\n interactions.append(interaction)\n \n interactions = torch.stack(interactions, dim=1)\n num_pairs = interactions.size(1)\n\n # Attention Layer\n attention_input = interactions \n attention_input_flat = attention_input.view(-1, self.embed_dims)\n attention_hidden = torch.relu(self.attention_W(attention_input_flat)) \n attention_scores = self.attention_h(attention_hidden) \n attention_scores = attention_scores.view(x.size(0), num_pairs, 1)\n attention_weights = torch.softmax(attention_scores, dim=1)\n weighted_interactions = attention_weights * interactions\n afm_output = torch.sum(weighted_interactions, dim=1)\n afm_output = self.projection(afm_output).squeeze(1)\n\n # Final Layer\n output = self.bias + linear_output + afm_output\n prediction = torch.sigmoid(output)\n \n return prediction","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:52:43.641091Z","iopub.status.busy":"2026-02-11T10:52:43.640843Z","iopub.status.idle":"2026-02-11T10:52:43.649594Z","shell.execute_reply":"2026-02-11T10:52:43.649013Z"},"papermill":{"duration":0.017239,"end_time":"2026-02-11T10:52:43.650891","exception":false,"start_time":"2026-02-11T10:52:43.633652","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 2.3 AFM-MLP (My Proposed Model)\n\nCombines AFM's attention-weighted pairwise interactions with an MLP branch for higher-order feature interactions.","metadata":{"papermill":{"duration":0.006436,"end_time":"2026-02-11T10:52:43.663601","exception":false,"start_time":"2026-02-11T10:52:43.657165","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# AFM-MLP: Combines attention-weighted FM with deep neural network\nclass DeepAFM(nn.Module):\n def __init__(self, field_dims, embed_dims=10, mlp_dims=(400,400,400),attention_dim=20, drop_rate = 0.1):\n super().__init__()\n self.num_fields = len(field_dims)\n self.embed_dims = embed_dims\n # Embedding layers\n self.embeddings = nn.ModuleList([\n nn.Embedding(dim, embed_dims) for dim in field_dims\n ])\n # Linear layers\n self.linear_weights = nn.ModuleList([\n nn.Embedding(num_embeddings=dim, embedding_dim=1)\n for dim in field_dims\n ])\n self.bias = nn.Parameter(torch.zeros(1))\n # Attention layers\n self.attention_W = nn.Linear(embed_dims, attention_dim, bias=True)\n self.attention_h = nn.Linear(attention_dim, 1, bias=False)\n # Projection to scalar\n self.projection = nn.Linear(embed_dims, 1, bias=False)\n # DNN/MLP part\n mlp_layers = []\n input_dim = self.num_fields * embed_dims\n for mlp_dim in mlp_dims:\n mlp_layers.extend([\n nn.Linear(input_dim, mlp_dim),\n nn.BatchNorm1d(mlp_dim),\n nn.ReLU(),\n nn.Dropout(drop_rate)\n ])\n input_dim = mlp_dim\n\n mlp_layers.append(nn.Linear(mlp_dims[-1], 1))\n \n self.dnn = nn.Sequential(*mlp_layers)\n\n def forward(self,x):\n embeds = torch.stack([emb(x[:, i]) for i, emb in enumerate(self.embeddings)], dim=1)\n #Linear Part\n linear_terms = torch.cat([\n self.linear_weights[i](x[:, i]) # (batch, 1)\n for i in range(self.num_fields)\n ], dim=1)\n linear_output = torch.sum(linear_terms, dim=1).squeeze()\n \n # Interaction Layers\n interactions = []\n for i in range(self.num_fields):\n for j in range(i + 1, self.num_fields):\n # Element-wise product\n interaction = embeds[:, i] * embeds[:, j]\n interactions.append(interaction)\n interactions = torch.stack(interactions, dim=1)\n num_pairs = interactions.size(1)\n \n # Attention Layer\n attention_input = interactions \n attention_input_flat = attention_input.view(-1, self.embed_dims)\n attention_hidden = torch.relu(self.attention_W(attention_input_flat)) \n attention_scores = self.attention_h(attention_hidden) \n attention_scores = attention_scores.view(x.size(0), num_pairs, 1)\n attention_weights = torch.softmax(attention_scores, dim=1)\n weighted_interactions = attention_weights * interactions\n afm_output = torch.sum(weighted_interactions, dim=1)\n afm_output = self.projection(afm_output).squeeze(1)\n\n # DNN Part\n embeds_flat = embeds.view(embeds.size(0), -1)\n dnn_output = self.dnn(embeds_flat).squeeze()\n\n # Combining all components\n output = self.bias + linear_output + afm_output + dnn_output\n # prediction = torch.sigmoid(output)\n \n # return prediction\n return output","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:52:43.677607Z","iopub.status.busy":"2026-02-11T10:52:43.677378Z","iopub.status.idle":"2026-02-11T10:52:43.689608Z","shell.execute_reply":"2026-02-11T10:52:43.688912Z"},"papermill":{"duration":0.021088,"end_time":"2026-02-11T10:52:43.690949","exception":false,"start_time":"2026-02-11T10:52:43.669861","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 2.4 FM (Factorisation Machine)","metadata":{"papermill":{"duration":0.006254,"end_time":"2026-02-11T10:52:43.703696","exception":false,"start_time":"2026-02-11T10:52:43.697442","status":"completed"},"tags":[]}},{"cell_type":"code","source":"class FM(nn.Module):\n def __init__(self, field_dims, embed_dims=10):\n super().__init__()\n self.num_fields = len(field_dims)\n self.embed_dims = embed_dims\n self.embeddings = nn.ModuleList([\n nn.Embedding(dim, embed_dims) for dim in field_dims\n ])\n self.linear_weights = nn.ModuleList([\n nn.Embedding(num_embeddings=dim, embedding_dim=1)\n for dim in field_dims\n ])\n self.bias = nn.Parameter(torch.zeros(1))\n\n def forward(self, x):\n embeds = torch.stack([emb(x[:, i]) for i, emb in enumerate(self.embeddings)], dim=1)\n # Linear Part\n linear_terms = torch.cat([\n self.linear_weights[i](x[:, i])\n for i in range(self.num_fields)\n ], dim=1)\n linear_output = torch.sum(linear_terms, dim=1)\n # FM Part (efficient O(nk) formulation)\n square_of_sum = torch.sum(embeds, dim=1).pow(2)\n sum_of_squares = torch.sum(embeds.pow(2), dim=1)\n cross_term = square_of_sum - sum_of_squares\n fm_output = 0.5 * torch.sum(cross_term, dim=-1)\n # Total\n output = self.bias + linear_output + fm_output\n prediction = torch.sigmoid(output)\n return prediction","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:52:43.717442Z","iopub.status.busy":"2026-02-11T10:52:43.717204Z","iopub.status.idle":"2026-02-11T10:52:43.723696Z","shell.execute_reply":"2026-02-11T10:52:43.723159Z"},"papermill":{"duration":0.014984,"end_time":"2026-02-11T10:52:43.725053","exception":false,"start_time":"2026-02-11T10:52:43.710069","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 2.5 Logistic Regression","metadata":{"papermill":{"duration":0.006302,"end_time":"2026-02-11T10:52:43.737609","exception":false,"start_time":"2026-02-11T10:52:43.731307","status":"completed"},"tags":[]}},{"cell_type":"code","source":"class LogisticRegressionModel(nn.Module):\n def __init__(self, field_dims, embed_dims=10):\n \"\"\"embed_dims accepted but unused, for API consistency.\"\"\"\n super().__init__()\n self.num_fields = len(field_dims)\n self.linear_weights = nn.ModuleList([\n nn.Embedding(num_embeddings=dim, embedding_dim=1)\n for dim in field_dims\n ])\n self.bias = nn.Parameter(torch.zeros(1))\n\n def forward(self, x):\n linear_terms = torch.cat([\n self.linear_weights[i](x[:, i])\n for i in range(self.num_fields)\n ], dim=1)\n linear_output = torch.sum(linear_terms, dim=1)\n output = self.bias + linear_output\n prediction = torch.sigmoid(output)\n return prediction","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:52:43.751473Z","iopub.status.busy":"2026-02-11T10:52:43.751238Z","iopub.status.idle":"2026-02-11T10:52:43.756281Z","shell.execute_reply":"2026-02-11T10:52:43.755692Z"},"papermill":{"duration":0.013289,"end_time":"2026-02-11T10:52:43.757533","exception":false,"start_time":"2026-02-11T10:52:43.744244","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"---\n\n## 3. Data Loading and Preprocessing\n\nLoading and preparing the Avazu CTR dataset (10M sampled observations).","metadata":{"papermill":{"duration":0.006487,"end_time":"2026-02-11T10:52:43.770535","exception":false,"start_time":"2026-02-11T10:52:43.764048","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"### 3.1 Load Raw Data","metadata":{"papermill":{"duration":0.006469,"end_time":"2026-02-11T10:52:43.783353","exception":false,"start_time":"2026-02-11T10:52:43.776884","status":"completed"},"tags":[]}},{"cell_type":"code","source":"train_schema = {\n 'id': pl.UInt64,\n 'click': pl.Int8,\n 'hour': pl.Int32,\n 'C1': pl.Int16,\n 'banner_pos': pl.Int8,\n 'site_id': pl.Categorical,\n 'site_domain': pl.Categorical,\n 'site_category': pl.Categorical,\n 'app_id': pl.Categorical,\n 'app_domain': pl.Categorical,\n 'app_category': pl.Categorical,\n 'device_id': pl.Categorical,\n 'device_ip': pl.Categorical,\n 'device_model': pl.Categorical,\n 'device_type': pl.Int8,\n 'device_conn_type': pl.Int8,\n 'C14': pl.Int16,\n 'C15': pl.Int16,\n 'C16': pl.Int16,\n 'C17': pl.Int16,\n 'C18': pl.Int16,\n 'C19': pl.Int16,\n 'C20': pl.Int32,\n 'C21': pl.Int16\n}\n\ntest_schema = {\n 'id': pl.UInt64,\n 'hour': pl.Int32,\n 'C1': pl.Int16,\n 'banner_pos': pl.Int8,\n 'site_id': pl.Categorical,\n 'site_domain': pl.Categorical,\n 'site_category': pl.Categorical,\n 'app_id': pl.Categorical,\n 'app_domain': pl.Categorical,\n 'app_category': pl.Categorical,\n 'device_id': pl.Categorical,\n 'device_ip': pl.Categorical,\n 'device_model': pl.Categorical,\n 'device_type': pl.Int8,\n 'device_conn_type': pl.Int8,\n 'C14': pl.Int16,\n 'C15': pl.Int16,\n 'C16': pl.Int16,\n 'C17': pl.Int16,\n 'C18': pl.Int16,\n 'C19': pl.Int16,\n 'C20': pl.Int32,\n 'C21': pl.Int16\n}\n\n# Load train and test datasets\ntrain_df = pl.read_csv('/kaggle/input/competitions/avazu-ctr-prediction/train.gz', schema=train_schema)\ntest_df = pl.read_csv('/kaggle/input/competitions/avazu-ctr-prediction/test.gz', schema=test_schema)\n\n","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:52:43.797402Z","iopub.status.busy":"2026-02-11T10:52:43.797156Z","iopub.status.idle":"2026-02-11T10:54:49.308889Z","shell.execute_reply":"2026-02-11T10:54:49.308193Z"},"papermill":{"duration":125.520888,"end_time":"2026-02-11T10:54:49.310751","exception":false,"start_time":"2026-02-11T10:52:43.789863","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Samples\nsample_size = 10000000\ntrain_df_sample = train_df.sample(sample_size, seed=42)\n# test_df_sample = test_df.sample(sample_size, seed=42)","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:54:49.325963Z","iopub.status.busy":"2026-02-11T10:54:49.325284Z","iopub.status.idle":"2026-02-11T10:54:52.380142Z","shell.execute_reply":"2026-02-11T10:54:52.379477Z"},"papermill":{"duration":3.064236,"end_time":"2026-02-11T10:54:52.381910","exception":false,"start_time":"2026-02-11T10:54:49.317674","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Extract hour features\ndef extract_hour_features(df):\n df = df.with_columns([\n # Convert YYMMDDHH to proper datetime immediately\n pl.datetime(\n year=2000 + (pl.col('hour') // 1000000),\n month=(pl.col('hour') // 10000) % 100,\n day=(pl.col('hour') // 100) % 100,\n hour=pl.col('hour') % 100\n ).alias('datetime'),\n \n # Extract components right away\n ((pl.col('hour') // 10000) % 100).cast(pl.Int8).alias('month'),\n ((pl.col('hour') // 100) % 100).cast(pl.Int8).alias('day'),\n (pl.col('hour') % 100).cast(pl.Int8).alias('hour_time'),\n])\n df = df.with_columns([pl.col('datetime').dt.weekday().alias('dayofweek')])\n return df.drop('hour')\n\ntrain = extract_hour_features(train_df_sample)\n# test = extract_hour_features(test_df_sample)","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:54:52.396800Z","iopub.status.busy":"2026-02-11T10:54:52.396570Z","iopub.status.idle":"2026-02-11T10:54:53.650802Z","shell.execute_reply":"2026-02-11T10:54:53.650188Z"},"papermill":{"duration":1.263594,"end_time":"2026-02-11T10:54:53.652485","exception":false,"start_time":"2026-02-11T10:54:52.388891","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"cols = ['C15', 'C16', 'C19', 'C21']\nfor col in cols:\n percentile_98 = train[col].quantile(0.98)\n max_val = train[col].max()\n \n if percentile_98 < 0.5 * max_val:\n # Cap values above 98th percentile\n train = train.with_columns(\n pl.when(pl.col(col) >= percentile_98)\n .then(percentile_98)\n .otherwise(pl.col(col))\n .alias(col)\n )","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:54:53.666929Z","iopub.status.busy":"2026-02-11T10:54:53.666667Z","iopub.status.idle":"2026-02-11T10:54:54.168517Z","shell.execute_reply":"2026-02-11T10:54:54.167887Z"},"papermill":{"duration":0.510808,"end_time":"2026-02-11T10:54:54.170189","exception":false,"start_time":"2026-02-11T10:54:53.659381","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 3.2 Exploratory Data Analysis\n\nExamining class imbalance in the dataset.","metadata":{"papermill":{"duration":0.006404,"end_time":"2026-02-11T10:54:54.183486","exception":false,"start_time":"2026-02-11T10:54:54.177082","status":"completed"},"tags":[]}},{"cell_type":"code","source":"click_counts = train.select(\n pl.col(\"click\").value_counts(normalize=True)\n).unnest(\"click\")\nprint(click_counts)\nprint(\"\\n\")\nclick_stats = train.group_by(\"click\").count().sort(\"click\")\nplt.figure(figsize=(8, 5))\nplt.bar(click_stats[\"click\"], click_stats[\"count\"])\nplt.xlabel(\"Click\")\nplt.ylabel(\"Count\")\nplt.title(\"Click Distribution\")\nplt.xticks([0, 1])\nplt.show()\n","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:54:54.197402Z","iopub.status.busy":"2026-02-11T10:54:54.197132Z","iopub.status.idle":"2026-02-11T10:54:54.589672Z","shell.execute_reply":"2026-02-11T10:54:54.588929Z"},"papermill":{"duration":0.401231,"end_time":"2026-02-11T10:54:54.591071","exception":false,"start_time":"2026-02-11T10:54:54.189840","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 3.3 Feature Selection and Encoding","metadata":{"papermill":{"duration":0.007197,"end_time":"2026-02-11T10:54:54.605550","exception":false,"start_time":"2026-02-11T10:54:54.598353","status":"completed"},"tags":[]}},{"cell_type":"code","source":"\nimportant_columns = [\n 'click', 'C1', 'banner_pos', 'site_id', 'site_domain', 'site_category',\n 'app_id', 'app_domain', 'app_category', 'device_id', 'device_model', \n 'device_type', 'device_conn_type', 'C14', 'C15', 'C16', 'C18', 'C19', 'dayofweek', 'day', 'hour_time'\n]\nfeature_cols = [col for col in important_columns if col != 'click']\n\ntrain = train.select(important_columns)\n# train = train.with_columns([\n# pl.col(c).cast(pl.Categorical).to_physical().cast(pl.Int64).alias(c)\n# for c in feature_cols\n# ])\n\ntrain = train.with_columns([\n (pl.col(c).rank(\"dense\") - 1).cast(pl.Int64).alias(c)\n for c in feature_cols\n])\n\nfield_dims = list(\n train\n .select([pl.col(c).n_unique().alias(c) for c in feature_cols])\n .row(0)\n)","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:54:54.620795Z","iopub.status.busy":"2026-02-11T10:54:54.620522Z","iopub.status.idle":"2026-02-11T10:55:02.808705Z","shell.execute_reply":"2026-02-11T10:55:02.807920Z"},"papermill":{"duration":8.198081,"end_time":"2026-02-11T10:55:02.810458","exception":false,"start_time":"2026-02-11T10:54:54.612377","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 3.4 Train/Validation/Test Split","metadata":{"papermill":{"duration":0.006716,"end_time":"2026-02-11T10:55:02.824530","exception":false,"start_time":"2026-02-11T10:55:02.817814","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Splitting into different datasets (50,25,25)\nX = train.drop(\"click\")\ny = train.select(\"click\")\n# Split test set\nX_train, X_test_val, y_train, y_test_val = train_test_split(X, y, stratify=y, test_size= 0.5, random_state= 42)\n\n# Split training set and validation set\nX_val, X_test, y_val, y_test = train_test_split(X_test_val, y_test_val, stratify=y_test_val, test_size = 0.5, random_state= 42)\n\nprint(\"X_train number of samples:\", len(X_train))\nprint(\"X_val number of samples:\", len(X_val))\nprint(\"X_test number of samples:\", len(X_test))","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:55:02.839342Z","iopub.status.busy":"2026-02-11T10:55:02.838767Z","iopub.status.idle":"2026-02-11T10:55:55.605473Z","shell.execute_reply":"2026-02-11T10:55:55.604705Z"},"papermill":{"duration":52.782305,"end_time":"2026-02-11T10:55:55.613557","exception":false,"start_time":"2026-02-11T10:55:02.831252","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Fix tensor creation - specify dtypes\nX_train = torch.tensor(X_train.to_numpy(), dtype=torch.long) # Important: long for embeddings\ny_train = torch.tensor(y_train.to_numpy(), dtype=torch.float32).squeeze() # Remove extra dim\nX_val = torch.tensor(X_val.to_numpy(), dtype=torch.long)\ny_val = torch.tensor(y_val.to_numpy(), dtype=torch.float32).squeeze()\nX_test = torch.tensor(X_test.to_numpy(), dtype=torch.long)\ny_test = torch.tensor(y_test.to_numpy(), dtype=torch.float32).squeeze()\n\ntrain_dataset = TensorDataset(X_train, y_train)\nval_dataset = TensorDataset(X_val, y_val)\ntest_dataset = TensorDataset(X_test, y_test)","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:55:55.629355Z","iopub.status.busy":"2026-02-11T10:55:55.628821Z","iopub.status.idle":"2026-02-11T10:55:56.943778Z","shell.execute_reply":"2026-02-11T10:55:56.942973Z"},"papermill":{"duration":1.324906,"end_time":"2026-02-11T10:55:56.945655","exception":false,"start_time":"2026-02-11T10:55:55.620749","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def debug_embedding_indices(X_train, X_test, field_dims):\n \"\"\"\n Debug function to check if indices match embedding sizes\n \"\"\"\n print(\"Debugging embedding indices...\")\n print(f\"Number of fields: {X_train.shape[1]}\")\n print(f\"Field dimensions: {field_dims}\")\n print(f\"Length of field_dims: {len(field_dims)}\")\n \n # Check each feature column\n for i in range(X_train.shape[1]):\n train_min = X_train[:, i].min().item()\n train_max = X_train[:, i].max().item()\n test_min = X_test[:, i].min().item() if X_test is not None else train_min\n test_max = X_test[:, i].max().item() if X_test is not None else train_max\n \n print(f\"\\nFeature {i}:\")\n print(f\" Train range: [{train_min}, {train_max}]\")\n print(f\" Test range: [{test_min}, {test_max}]\")\n print(f\" Embedding size: {field_dims[i] if i < len(field_dims) else 'MISSING'}\")\n \n # Check for issues\n if i >= len(field_dims):\n print(f\" ❌ ERROR: No field_dim for feature {i}\")\n elif train_max >= field_dims[i]:\n print(f\" ❌ ERROR: Max value {train_max} >= embedding size {field_dims[i]}\")\n elif test_max >= field_dims[i]:\n print(f\" ❌ ERROR: Test max value {test_max} >= embedding size {field_dims[i]}\")\n elif train_min < 0:\n print(f\" ⚠️ WARNING: Negative values found (min={train_min})\")\n\n# Run debug\n# Assuming you have X_train_tensor and X_test_tensor from your pipeline\ndebug_embedding_indices(X_train, X_test, field_dims)","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:55:56.963356Z","iopub.status.busy":"2026-02-11T10:55:56.962724Z","iopub.status.idle":"2026-02-11T10:55:57.122039Z","shell.execute_reply":"2026-02-11T10:55:57.121131Z"},"papermill":{"duration":0.169202,"end_time":"2026-02-11T10:55:57.123818","exception":false,"start_time":"2026-02-11T10:55:56.954616","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"---\n\n## 4. Training Infrastructure\n\nCustom training loop with early stopping, checkpointing, and metric tracking.","metadata":{"papermill":{"duration":0.007199,"end_time":"2026-02-11T10:55:57.140069","exception":false,"start_time":"2026-02-11T10:55:57.132870","status":"completed"},"tags":[]}},{"cell_type":"code","source":"class RichEpochDashboard:\n def __init__(self, total_epochs: int, console=None, max_rows: int = 12):\n self.total_epochs = total_epochs\n self.console = console or Console()\n self.max_rows = max_rows\n self.rows = [] # list of dicts\n\n def _render(self):\n table = Table(title=\"Training Monitor\", expand=True)\n table.add_column(\"Epoch\", justify=\"right\")\n table.add_column(\"Train Loss\", justify=\"right\", style=\"cyan\")\n table.add_column(\"Val Loss\", justify=\"right\", style=\"red\")\n table.add_column(\"Train AUC\", justify=\"right\", style=\"cyan\")\n table.add_column(\"Val AUC\", justify=\"right\", style=\"red\")\n table.add_column(\"LR\", justify=\"right\", style=\"magenta\")\n table.add_column(\"Time (s)\", justify=\"right\")\n\n # only last max_rows\n for r in self.rows[-self.max_rows:]:\n table.add_row(\n str(r[\"epoch\"]),\n f\"{r['train_loss']:.4f}\",\n f\"{r['val_loss']:.4f}\",\n f\"{r['train_auc']:.4f}\",\n f\"{r['val_auc']:.4f}\",\n f\"{r['lr']:.2e}\",\n f\"{r['time_s']:.1f}\",\n )\n\n best_loss = getattr(self, \"best_val_loss\", None)\n bad = getattr(self, \"bad_epochs\", 0)\n pat = getattr(self, \"patience\", 0)\n \n subtitle = \"\"\n if best_loss is not None and best_loss < float(\"inf\"):\n subtitle = f\"Best val loss: {best_loss:.4f} | EarlyStop: {bad}/{pat}\"\n\n header = Text(f\"Epoch {self.rows[-1]['epoch']}/{self.total_epochs}\" if self.rows else \"Starting…\", style=\"bold green\")\n return Panel(table, title=header, subtitle=subtitle, border_style=\"green\")\n\n def add_epoch(self, epoch, train_loss, val_loss, train_auc, val_auc, lr, time_s,\n best_val_loss=None, improved=False, bad_epochs=0, patience=0):\n self.rows.append({\n \"epoch\": epoch,\n \"train_loss\": train_loss,\n \"val_loss\": val_loss,\n \"train_auc\": train_auc,\n \"val_auc\": val_auc,\n \"lr\": lr,\n \"time_s\": time_s,\n \"improved\": improved,\n })\n self.best_val_loss = best_val_loss\n self.bad_epochs = bad_epochs\n self.patience = patience\n return self._render()\n\n\nclass RichTerminalDashboard:\n def __init__(self, total_epochs: int, max_rows: int = 10):\n self.console = Console()\n self.total_epochs = total_epochs\n self.max_rows = max_rows\n self.rows = []\n self.train_task = None\n self.val_task = None\n self._last_lr = 0.0\n self._build_progress()\n\n def _build_progress(self):\n self.progress = Progress(\n TextColumn(\"[bold cyan]{task.description}\"),\n BarColumn(),\n TextColumn(\"{task.completed}/{task.total}\"),\n TextColumn(\"•\"),\n TextColumn(\"[yellow]loss[/yellow]={task.fields[loss]:.4f}\"),\n TextColumn(\"avg={task.fields[avg]:.4f}\"),\n TextColumn(\"lr={task.fields[lr]:.2e}\"),\n TextColumn(\"•\"),\n TimeElapsedColumn(),\n TimeRemainingColumn(),\n expand=True,\n )\n\n def start_epoch(self, epoch_idx: int, train_steps: int, val_steps: int):\n self._build_progress()\n\n self.train_task = self.progress.add_task(\n f\"Epoch {epoch_idx+1}/{self.total_epochs} [Train]\",\n total=train_steps,\n loss=0.0, avg=0.0, lr=self._last_lr\n )\n self.val_task = self.progress.add_task(\n f\"Epoch {epoch_idx+1}/{self.total_epochs} [Val] \",\n total=val_steps,\n loss=0.0, avg=0.0, lr=self._last_lr\n )\n\n def update_train(self, step_loss: float, avg_loss: float, lr: float):\n self._last_lr = float(lr)\n if self.train_task is not None:\n self.progress.update(self.train_task, advance=1, loss=float(step_loss), avg=float(avg_loss), lr=float(lr))\n if self.val_task is not None:\n self.progress.update(self.val_task, lr=float(lr))\n\n def update_val(self, step_loss: float, avg_loss: float, lr: float = None):\n if lr is None:\n lr = self._last_lr\n if self.val_task is not None:\n self.progress.update(self.val_task, advance=1, loss=float(step_loss), avg=float(avg_loss), lr=float(lr))\n\n def add_epoch_row(self, **kwargs):\n self.rows.append(dict(kwargs))\n\n def render(self):\n layout = Layout()\n layout.split_column(Layout(name=\"top\", size=12), Layout(name=\"bottom\"))\n\n table = Table(title=\"Epoch Summary\", expand=True)\n table.add_column(\"Epoch\", justify=\"right\")\n table.add_column(\"Train Loss\", justify=\"right\", style=\"cyan\")\n table.add_column(\"Val Loss\", justify=\"right\", style=\"red\")\n table.add_column(\"Train AUC\", justify=\"right\", style=\"cyan\")\n table.add_column(\"Val AUC\", justify=\"right\", style=\"red\")\n table.add_column(\"Best\", justify=\"left\", style=\"green\")\n table.add_column(\"LR\", justify=\"right\", style=\"magenta\")\n\n for r in self.rows[-self.max_rows:]:\n epoch = r.get(\"epoch\", \"\")\n tr_loss = r.get(\"train_loss\", float(\"nan\"))\n va_loss = r.get(\"val_loss\", float(\"nan\"))\n tr_auc = r.get(\"train_auc\", float(\"nan\"))\n va_auc = r.get(\"val_auc\", float(\"nan\"))\n lr = r.get(\"lr\", self._last_lr)\n\n best_tag = \"\"\n best_epoch = r.get(\"best_epoch\", None)\n best_metric = r.get(\"best_metric\", None)\n monitor = r.get(\"monitor\", \"metric\")\n if best_epoch is not None and epoch == best_epoch:\n if best_metric is not None and not math.isnan(best_metric):\n best_tag = f\"best ({monitor}) {best_metric:.4f}\"\n\n table.add_row(\n str(epoch),\n f\"{tr_loss:.4f}\",\n f\"{va_loss:.4f}\",\n f\"{tr_auc:.4f}\",\n f\"{va_auc:.4f}\",\n best_tag,\n f\"{lr:.2e}\",\n )\n\n layout[\"top\"].update(Panel(table, border_style=\"green\"))\n layout[\"bottom\"].update(Panel(self.progress, border_style=\"blue\", title=\"Progress\"))\n return layout\n\n\nclass NotebookTrainingMonitor:\n def __init__(self):\n self.console = None\n self.history = {\"train_loss\": [], \"val_loss\": [], \"train_auc\": [], \"val_auc\": [], \"lr\": []}\n\n def start_epoch(self, epoch_idx: int, train_steps: int, val_steps: int):\n pass\n\n def update_train(self, step_loss: float, avg_loss: float, lr: float):\n # keep notebook fast: no per-batch updates by default\n pass\n\n def update_val(self, step_loss: float, avg_loss: float, lr: float = None):\n pass\n\n def add_epoch_row(self, **kwargs):\n self.history[\"train_loss\"].append(float(kwargs.get(\"train_loss\", float(\"nan\"))))\n self.history[\"val_loss\"].append(float(kwargs.get(\"val_loss\", float(\"nan\"))))\n self.history[\"train_auc\"].append(float(kwargs.get(\"train_auc\", float(\"nan\"))))\n self.history[\"val_auc\"].append(float(kwargs.get(\"val_auc\", float(\"nan\"))))\n self.history[\"lr\"].append(float(kwargs.get(\"lr\", 0.0)))\n\n epoch = kwargs.get(\"epoch\", len(self.history[\"train_loss\"]))\n total_epochs = kwargs.get(\"total_epochs\", None)\n\n clear_output(wait=True)\n if total_epochs is not None:\n print(f\"Epoch {epoch}/{total_epochs}\")\n else:\n print(f\"Epoch {epoch}\")\n\n xs = range(1, len(self.history[\"train_loss\"]) + 1)\n fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n\n axes[0].plot(xs, self.history[\"train_loss\"], label=\"Train\")\n axes[0].plot(xs, self.history[\"val_loss\"], label=\"Val\")\n axes[0].set_title(\"Loss\")\n axes[0].grid(True, alpha=0.3)\n axes[0].legend()\n\n axes[1].plot(xs, self.history[\"train_auc\"], label=\"Train\")\n axes[1].plot(xs, self.history[\"val_auc\"], label=\"Val\")\n axes[1].set_title(\"AUC\")\n axes[1].grid(True, alpha=0.3)\n axes[1].legend()\n\n plt.tight_layout()\n display(fig)\n plt.close(fig)\n\n print(f\"Train Loss: {self.history['train_loss'][-1]:.4f} | Val Loss: {self.history['val_loss'][-1]:.4f}\")\n print(f\"Train AUC : {self.history['train_auc'][-1]:.4f} | Val AUC : {self.history['val_auc'][-1]:.4f}\")\n print(f\"LR : {self.history['lr'][-1]:.2e}\")\n\n def render(self):\n return None\n\ndef in_notebook() -> bool:\n try:\n from IPython import get_ipython\n ip = get_ipython()\n if ip is None:\n return False\n # Jupyter / Colab kernels typically have this\n return \"IPKernelApp\" in ip.config\n except Exception:\n return False\n\nclass NullDashboard:\n def __init__(self, total_epochs: int):\n self.total_epochs = total_epochs\n self.console = None # to match RichTerminalDashboard attribute\n\n def start_epoch(self, epoch_idx: int, train_steps: int, val_steps: int):\n pass\n\n def update_train(self, step_loss: float, avg_loss: float, lr: float):\n pass\n\n def update_val(self, step_loss: float, avg_loss: float, lr: float = None):\n pass\n\n def add_epoch_row(self, **kwargs):\n pass\n\n def render(self):\n return None\n\nclass _NullLive:\n def __enter__(self): return self\n def __exit__(self, exc_type, exc, tb): return False\n def update(self, *args, **kwargs): pass\n\n\ndef try_make_dashboard(total_epochs: int, use_rich: bool = True, prefer_notebook: bool = True):\n \"\"\"\n Returns (LiveClassOrNone, dashboardInstance)\n\n Priority:\n 1) RichTerminalDashboard + Live (if use_rich and rich is importable)\n 2) NotebookTrainingMonitor (if in notebook and prefer_notebook)\n 3) NullDashboard\n \"\"\"\n # 1) Try Rich\n if use_rich:\n try:\n from rich.live import Live\n return Live, RichTerminalDashboard(total_epochs)\n except Exception:\n pass\n\n # 2) Notebook fallback\n if prefer_notebook and in_notebook():\n return None, NotebookTrainingMonitor()\n\n # 3) No UI\n return None, NullDashboard(total_epochs)","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:55:57.155004Z","iopub.status.busy":"2026-02-11T10:55:57.154756Z","iopub.status.idle":"2026-02-11T10:55:57.185051Z","shell.execute_reply":"2026-02-11T10:55:57.184480Z"},"papermill":{"duration":0.039701,"end_time":"2026-02-11T10:55:57.186608","exception":false,"start_time":"2026-02-11T10:55:57.146907","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def seed_everything(seed: int = 42, deterministic: bool = False) -> None:\n random.seed(seed)\n np.random.seed(seed)\n torch.manual_seed(seed)\n torch.cuda.manual_seed_all(seed)\n\n if deterministic:\n torch.backends.cudnn.deterministic = True\n torch.backends.cudnn.benchmark = False\n else:\n torch.backends.cudnn.benchmark = True\n\n\ndef safe_auc(y_true: np.ndarray, y_prob: np.ndarray) -> float:\n # roc_auc_score errors if only one class present; return NaN in that case\n if y_true.size == 0 or len(np.unique(y_true)) <= 1:\n return float(\"nan\")\n return float(roc_auc_score(y_true, y_prob))","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:55:57.201607Z","iopub.status.busy":"2026-02-11T10:55:57.201365Z","iopub.status.idle":"2026-02-11T10:55:57.206204Z","shell.execute_reply":"2026-02-11T10:55:57.205611Z"},"papermill":{"duration":0.01384,"end_time":"2026-02-11T10:55:57.207476","exception":false,"start_time":"2026-02-11T10:55:57.193636","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 4.1 Training Configuration","metadata":{"papermill":{"duration":0.006648,"end_time":"2026-02-11T10:55:57.220959","exception":false,"start_time":"2026-02-11T10:55:57.214311","status":"completed"},"tags":[]}},{"cell_type":"code","source":"@dataclass\nclass TrainConfig:\n # optimization\n epochs: int = 5\n lr: float = 1e-3\n weight_decay: float = 0.0\n grad_clip: Optional[float] = 1.0\n\n # regularization policy: use ONE of these\n use_manual_l2: bool = False # if True, will call model.l2_regularization_loss() when available\n # if use_manual_l2=True, recommend weight_decay=0.0 to avoid double regularization\n\n # logging\n log_every_steps: int = 10\n\n # early stopping / selection\n monitor: str = \"val_loss\" # \"val_loss\" or \"val_auc\"\n monitor_mode: str = \"min\" # \"min\" for loss, \"max\" for auc\n min_delta: float = 1e-6\n early_stop_patience: int = 2\n\n # model output convention\n model_returns_probs: bool = False # False: logits (recommended) -> BCEWithLogitsLoss\n\n # amp (optional)\n use_amp: bool = False\n\n # scheduler (optional)\n # \"none\" | \"plateau\" | \"cosine\"\n scheduler: str = \"none\"\n plateau_factor: float = 0.5\n plateau_patience: int = 1\n cosine_tmax: Optional[int] = None # defaults to epochs\n\n # reproducibility\n seed: int = 42\n deterministic: bool = False\n\n\nclass Trainer:\n def __init__(self, model, log_dir=\"runs/ctr\", run_name=None, device=None):\n self.model = model\n self.device = device or torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n self.model.to(self.device)\n\n if run_name is None:\n run_name = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n\n self.run_path = os.path.join(log_dir, run_name)\n os.makedirs(self.run_path, exist_ok=True)\n \n self.writer = SummaryWriter(self.run_path)\n print(f\"📊 TensorBoard: tensorboard --logdir {log_dir}\")\n print(f\"🧾 Run path: {self.run_path}\")\n\n def _loss_fn(self, cfg: TrainConfig):\n return nn.BCELoss() if cfg.model_returns_probs else nn.BCEWithLogitsLoss()\n\n def _to_probs(self, outputs, cfg: TrainConfig):\n if cfg.model_returns_probs:\n return outputs\n return torch.sigmoid(outputs)\n\n def _make_scheduler(self, optimizer: torch.optim.Optimizer, cfg: TrainConfig):\n if cfg.scheduler == \"none\":\n return None\n if cfg.scheduler == \"plateau\":\n # NOTE: expects metric; we call scheduler.step(metric) each epoch\n return torch.optim.lr_scheduler.ReduceLROnPlateau(\n optimizer, mode=(\"min\" if cfg.monitor_mode == \"min\" else \"max\"),\n factor=cfg.plateau_factor, patience=cfg.plateau_patience\n )\n if cfg.scheduler == \"cosine\":\n tmax = cfg.cosine_tmax or cfg.epochs\n return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=tmax)\n raise ValueError(f\"Unknown scheduler: {cfg.scheduler}\")\n\n def _is_improvement(self, current: float, best: float, cfg: TrainConfig) -> bool:\n if math.isnan(current):\n return False\n if cfg.monitor_mode == \"min\":\n return current < (best - cfg.min_delta)\n if cfg.monitor_mode == \"max\":\n return current > (best + cfg.min_delta)\n raise ValueError(f\"monitor_mode must be 'min' or 'max', got {cfg.monitor_mode}\")\n\n def _get_monitor_value(self, metrics: Dict[str, float], cfg: TrainConfig) -> float:\n if cfg.monitor not in metrics:\n raise KeyError(f\"cfg.monitor='{cfg.monitor}' not in metrics keys {list(metrics.keys())}\")\n return float(metrics[cfg.monitor])\n\n def save_checkpoint(\n self,\n path: str,\n *,\n optimizer: torch.optim.Optimizer,\n scheduler: Optional[Any],\n scaler: Optional[Any],\n epoch: int,\n best_metric: float,\n best_epoch: int,\n bad_epochs: int,\n cfg: TrainConfig,\n extra_meta: Optional[Dict[str, Any]] = None,\n ) -> None:\n ckpt = {\n \"model\": self.model.state_dict(),\n \"optimizer\": optimizer.state_dict(),\n \"scheduler\": scheduler.state_dict() if scheduler is not None and hasattr(scheduler, \"state_dict\") else None,\n \"scaler\": scaler.state_dict() if scaler is not None and hasattr(scaler, \"state_dict\") else None,\n \"epoch\": int(epoch),\n \"best_metric\": float(best_metric),\n \"best_epoch\": int(best_epoch),\n \"bad_epochs\": int(bad_epochs),\n \"cfg\": asdict(cfg),\n \"meta\": extra_meta or {},\n }\n torch.save(ckpt, path)\n\n def load_checkpoint(\n self,\n path: str,\n optimizer: Optional[torch.optim.Optimizer] = None,\n scheduler: Optional[Any] = None,\n scaler: Optional[Any] = None,\n map_location: Optional[str] = None,\n strict: bool = True,\n ) -> Dict[str, Any]:\n ckpt = torch.load(path, map_location=map_location or self.device)\n self.model.load_state_dict(ckpt[\"model\"], strict=strict)\n\n if optimizer is not None and ckpt.get(\"optimizer\") is not None:\n optimizer.load_state_dict(ckpt[\"optimizer\"])\n\n if scheduler is not None and ckpt.get(\"scheduler\") is not None:\n try:\n scheduler.load_state_dict(ckpt[\"scheduler\"])\n except Exception:\n pass\n\n if scaler is not None and ckpt.get(\"scaler\") is not None:\n try:\n scaler.load_state_dict(ckpt[\"scaler\"])\n except Exception:\n pass\n\n return ckpt\n\n def train_one_epoch(self, train_loader, optimizer, loss_fn, cfg: TrainConfig, epoch: int, scaler: Optional[Any] = None, dashboard=None):\n self.model.train()\n running_loss = 0.0\n all_probs = []\n all_y = []\n n = 0\n start = time.time()\n\n \n pbar = tqdm(train_loader, desc=f\"Epoch {epoch+1} [Train]\", leave=False)\n for step, (X, y) in enumerate(pbar):\n X = X.to(self.device, non_blocking=True)\n y = y.to(self.device, non_blocking=True).float().view(-1)\n \n optimizer.zero_grad(set_to_none=True)\n \n if cfg.use_amp and scaler is not None and self.device.type == \"cuda\":\n with torch.cuda.amp.autocast():\n out = self.model(X).view(-1)\n loss = loss_fn(out, y)\n if cfg.use_manual_l2 and hasattr(self.model, \"l2_regularization_loss\"):\n loss = loss + self.model.l2_regularization_loss()\n\n scaler.scale(loss).backward()\n if cfg.grad_clip is not None:\n scaler.unscale_(optimizer)\n nn.utils.clip_grad_norm_(self.model.parameters(), cfg.grad_clip)\n scaler.step(optimizer)\n scaler.update()\n else:\n out = self.model(X).view(-1)\n loss = loss_fn(out, y)\n if cfg.use_manual_l2 and hasattr(self.model, \"l2_regularization_loss\"):\n loss = loss + self.model.l2_regularization_loss()\n\n loss.backward()\n if cfg.grad_clip is not None:\n nn.utils.clip_grad_norm_(self.model.parameters(), cfg.grad_clip)\n optimizer.step()\n \n bs = X.size(0)\n running_loss += loss.item() * bs\n n += bs\n \n probs = self._to_probs(out.detach(), cfg)\n all_probs.append(probs.cpu())\n all_y.append(y.detach().cpu())\n \n avg_loss = running_loss / max(n, 1)\n lr_now = float(optimizer.param_groups[0][\"lr\"])\n pbar.set_postfix(loss=float(loss.item()), avg_loss=float(avg_loss), lr=lr_now)\n\n # logging\n if step % cfg.log_every_steps == 0:\n global_step = epoch * len(train_loader) + step\n self.writer.add_scalar(\"batch/train_loss\", float(loss.item()), global_step)\n self.writer.add_scalar(\"batch/lr\", lr_now, global_step)\n\n if dashboard is not None:\n dashboard.update_train(step_loss=float(loss.item()), avg_loss=float(avg_loss), lr=lr_now)\n \n all_probs = torch.cat(all_probs).numpy()\n all_y = torch.cat(all_y).numpy()\n auc = roc_auc_score(all_y, all_probs) if len(np.unique(all_y)) > 1 else float(\"nan\")\n avg_loss = running_loss / max(n, 1)\n elapsed = time.time() - start\n return float(avg_loss), float(auc), float(elapsed)\n\n\n @torch.no_grad()\n def validate(self, val_loader, loss_fn, cfg: TrainConfig, epoch: int, split_name=\"val\", dashboard=None):\n self.model.eval()\n running_loss = 0.0\n all_probs = []\n all_y = []\n n = 0\n\n pbar = tqdm(val_loader, desc=f\"Epoch {epoch+1} [{split_name}]\", leave=False)\n for step, (X, y) in enumerate(pbar):\n X = X.to(self.device, non_blocking=True)\n y = y.to(self.device, non_blocking=True).float().view(-1)\n \n out = self.model(X).view(-1)\n loss = loss_fn(out, y)\n if cfg.use_manual_l2 and hasattr(self.model, \"l2_regularization_loss\"):\n loss = loss + self.model.l2_regularization_loss()\n \n bs = X.size(0)\n running_loss += loss.item() * bs\n n += bs\n \n probs = self._to_probs(out, cfg)\n all_probs.append(probs.cpu())\n all_y.append(y.cpu())\n \n avg_loss = running_loss / max(n, 1)\n pbar.set_postfix(loss=float(loss.item()), avg_loss=float(avg_loss))\n\n if dashboard is not None:\n dashboard.update_val(step_loss=float(loss.item()), avg_loss=float(avg_loss))\n\n \n all_probs = torch.cat(all_probs).numpy()\n all_y = torch.cat(all_y).numpy()\n auc = roc_auc_score(all_y, all_probs) if len(np.unique(all_y)) > 1 else float(\"nan\")\n avg_loss = running_loss / max(n, 1)\n \n return float(avg_loss), float(auc), all_y, all_probs\n\n \n def fit(self, train_loader, val_loader, cfg: TrainConfig, test_loader=None, save_dir=None, resume_from: Optional[str] = None, use_rich_dashboard: bool = False, prefer_notebook_dashboard: bool = True, extra_meta: Optional[Dict[str, Any]] = None):\n\n seed_everything(cfg.seed, cfg.deterministic)\n\n if cfg.use_manual_l2 and cfg.weight_decay != 0.0:\n print(\"⚠️ cfg.use_manual_l2=True and weight_decay!=0.0. Consider setting weight_decay=0 to avoid double regularization.\")\n\n loss_fn = self._loss_fn(cfg)\n optimizer = torch.optim.AdamW(self.model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)\n\n scheduler = self._make_scheduler(optimizer, cfg)\n\n use_amp = cfg.use_amp and self.device.type == \"cuda\"\n\n scaler = torch.amp.GradScaler(\"cuda\",enabled=use_amp)\n\n if save_dir is None:\n save_dir = self.run_path\n os.makedirs(save_dir, exist_ok=True)\n\n # paths\n best_ckpt_path = os.path.join(save_dir, \"best.ckpt\")\n last_ckpt_path = os.path.join(save_dir, \"last.ckpt\")\n history_path = os.path.join(save_dir, \"history.json\")\n cfg_path = os.path.join(save_dir, \"cfg.json\")\n meta_path = os.path.join(save_dir, \"meta.json\")\n\n with open(cfg_path, \"w\") as f:\n json.dump(asdict(cfg), f, indent=2)\n with open(meta_path, \"w\") as f:\n json.dump(extra_meta or {}, f, indent=2)\n \n history: Dict[str, Any] = {\n \"train_loss\": [],\n \"val_loss\": [],\n \"train_auc\": [],\n \"val_auc\": [],\n \"epoch_time_sec\": [],\n \"lr\": [],\n \"best\": {\n \"metric_name\": cfg.monitor,\n \"mode\": cfg.monitor_mode,\n \"value\": (float(\"inf\") if cfg.monitor_mode == \"min\" else -float(\"inf\")),\n \"epoch\": 0,\n \"train_loss\": None,\n \"train_auc\": None,\n \"val_loss\": None,\n \"val_auc\": None,\n },\n \"test\": None,\n }\n\n # early stopping state\n best_metric = history[\"best\"][\"value\"]\n best_epoch = 0\n bad_epochs = 0\n start_epoch = 0\n\n # resume\n if resume_from is not None and os.path.exists(resume_from):\n ckpt = self.load_checkpoint(resume_from, optimizer=optimizer, scheduler=scheduler, scaler=scaler)\n start_epoch = int(ckpt.get(\"epoch\", 0))\n best_metric = float(ckpt.get(\"best_metric\", best_metric))\n best_epoch = int(ckpt.get(\"best_epoch\", best_epoch))\n bad_epochs = int(ckpt.get(\"bad_epochs\", bad_epochs))\n history_resume = ckpt.get(\"meta\", {}).get(\"history\")\n if isinstance(history_resume, dict):\n history = history_resume\n print(f\"🔁 Resumed from {resume_from} at epoch {start_epoch} (best_epoch={best_epoch}, bad_epochs={bad_epochs}).\")\n \n Live, dashboard = try_make_dashboard(total_epochs=cfg.epochs, use_rich=use_rich_dashboard, prefer_notebook=prefer_notebook_dashboard)\n \n try:\n if Live is not None:\n with Live(dashboard.render(), console=dashboard.console, refresh_per_second=10) as live:\n for epoch in range(start_epoch, cfg.epochs):\n dashboard.start_epoch(epoch, len(train_loader), len(val_loader))\n live.update(dashboard.render())\n\n tr_loss, tr_auc, tr_time = self.train_one_epoch(\n train_loader, optimizer, loss_fn, cfg, epoch, scaler=scaler, dashboard=dashboard\n )\n va_loss, va_auc, y_val, p_val = self.validate(\n val_loader, loss_fn, cfg, epoch, split_name=\"val\", dashboard=dashboard\n )\n\n lr_now = float(optimizer.param_groups[0][\"lr\"])\n metrics = {\"train_loss\": tr_loss, \"val_loss\": va_loss, \"train_auc\": tr_auc, \"val_auc\": va_auc}\n monitor_value = self._get_monitor_value(metrics, cfg)\n\n # scheduler step\n if scheduler is not None:\n if cfg.scheduler == \"plateau\":\n scheduler.step(monitor_value)\n else:\n scheduler.step()\n\n # tensorboard epoch logs\n self.writer.add_scalars(\"loss\", {\"train\": tr_loss, \"val\": va_loss}, epoch)\n self.writer.add_scalars(\"auc\", {\"train\": tr_auc, \"val\": va_auc}, epoch)\n self.writer.add_scalar(\"time/epoch_seconds\", tr_time, epoch)\n self.writer.add_scalar(\"lr/epoch\", lr_now, epoch)\n\n # record history\n history[\"train_loss\"].append(tr_loss)\n history[\"val_loss\"].append(va_loss)\n history[\"train_auc\"].append(tr_auc)\n history[\"val_auc\"].append(va_auc)\n history[\"epoch_time_sec\"].append(tr_time)\n history[\"lr\"].append(lr_now)\n\n improved = self._is_improvement(monitor_value, best_metric, cfg)\n if improved:\n best_metric = monitor_value\n best_epoch = epoch + 1\n bad_epochs = 0\n history[\"best\"][\"value\"] = float(best_metric)\n history[\"best\"][\"epoch\"] = int(best_epoch)\n history[\"best\"][\"train_loss\"] = float(tr_loss)\n history[\"best\"][\"train_auc\"] = float(tr_auc)\n history[\"best\"][\"val_loss\"] = float(va_loss)\n history[\"best\"][\"val_auc\"] = float(va_auc)\n\n self.save_checkpoint(\n best_ckpt_path,\n optimizer=optimizer,\n scheduler=scheduler,\n scaler=scaler,\n epoch=epoch + 1,\n best_metric=float(best_metric),\n best_epoch=int(best_epoch),\n bad_epochs=int(bad_epochs),\n cfg=cfg,\n extra_meta={\"history\": history, **(extra_meta or {})},\n )\n else:\n bad_epochs += 1\n\n # always save last\n self.save_checkpoint(\n last_ckpt_path,\n optimizer=optimizer,\n scheduler=scheduler,\n scaler=scaler,\n epoch=epoch + 1,\n best_metric=float(best_metric),\n best_epoch=int(best_epoch),\n bad_epochs=int(bad_epochs),\n cfg=cfg,\n extra_meta={\"history\": history, **(extra_meta or {})},\n )\n\n dashboard.add_epoch_row(\n epoch=epoch + 1,\n total_epochs=cfg.epochs,\n train_loss=tr_loss,\n val_loss=va_loss,\n train_auc=tr_auc,\n val_auc=va_auc,\n lr=lr_now,\n time_s=tr_time,\n best_epoch=best_epoch,\n best_metric=float(best_metric),\n bad_epochs=bad_epochs,\n patience=cfg.early_stop_patience,\n monitor=cfg.monitor,\n )\n live.update(dashboard.render())\n\n if bad_epochs >= cfg.early_stop_patience:\n print(f\"⏹️ Early stop at epoch {epoch+1}. Best {cfg.monitor}={best_metric:.6g} at epoch {best_epoch}.\")\n break\n else:\n for epoch in range(start_epoch, cfg.epochs):\n dashboard.start_epoch(epoch, len(train_loader), len(val_loader))\n\n tr_loss, tr_auc, tr_time = self.train_one_epoch(\n train_loader, optimizer, loss_fn, cfg, epoch, scaler=scaler, dashboard=dashboard\n )\n va_loss, va_auc, y_val, p_val = self.validate(\n val_loader, loss_fn, cfg, epoch, split_name=\"val\", dashboard=dashboard\n )\n\n lr_now = float(optimizer.param_groups[0][\"lr\"])\n metrics = {\"train_loss\": tr_loss, \"val_loss\": va_loss, \"train_auc\": tr_auc, \"val_auc\": va_auc}\n monitor_value = self._get_monitor_value(metrics, cfg)\n\n if scheduler is not None:\n if cfg.scheduler == \"plateau\":\n scheduler.step(monitor_value)\n else:\n scheduler.step()\n\n self.writer.add_scalars(\"loss\", {\"train\": tr_loss, \"val\": va_loss}, epoch)\n self.writer.add_scalars(\"auc\", {\"train\": tr_auc, \"val\": va_auc}, epoch)\n self.writer.add_scalar(\"time/epoch_seconds\", tr_time, epoch)\n self.writer.add_scalar(\"lr/epoch\", lr_now, epoch)\n\n history[\"train_loss\"].append(tr_loss)\n history[\"val_loss\"].append(va_loss)\n history[\"train_auc\"].append(tr_auc)\n history[\"val_auc\"].append(va_auc)\n history[\"epoch_time_sec\"].append(tr_time)\n history[\"lr\"].append(lr_now)\n\n improved = self._is_improvement(monitor_value, best_metric, cfg)\n if improved:\n best_metric = monitor_value\n best_epoch = epoch + 1\n bad_epochs = 0\n history[\"best\"][\"value\"] = float(best_metric)\n history[\"best\"][\"epoch\"] = int(best_epoch)\n history[\"best\"][\"train_loss\"] = float(tr_loss)\n history[\"best\"][\"train_auc\"] = float(tr_auc)\n history[\"best\"][\"val_loss\"] = float(va_loss)\n history[\"best\"][\"val_auc\"] = float(va_auc)\n\n self.save_checkpoint(\n best_ckpt_path,\n optimizer=optimizer,\n scheduler=scheduler,\n scaler=scaler,\n epoch=epoch + 1,\n best_metric=float(best_metric),\n best_epoch=int(best_epoch),\n bad_epochs=int(bad_epochs),\n cfg=cfg,\n extra_meta={\"history\": history, **(extra_meta or {})},\n )\n else:\n bad_epochs += 1\n\n self.save_checkpoint(\n last_ckpt_path,\n optimizer=optimizer,\n scheduler=scheduler,\n scaler=scaler,\n epoch=epoch + 1,\n best_metric=float(best_metric),\n best_epoch=int(best_epoch),\n bad_epochs=int(bad_epochs),\n cfg=cfg,\n extra_meta={\"history\": history, **(extra_meta or {})},\n )\n\n dashboard.add_epoch_row(\n epoch=epoch + 1,\n total_epochs=cfg.epochs,\n train_loss=tr_loss,\n val_loss=va_loss,\n train_auc=tr_auc,\n val_auc=va_auc,\n lr=lr_now,\n time_s=tr_time,\n best_epoch=best_epoch,\n best_metric=float(best_metric),\n bad_epochs=bad_epochs,\n patience=cfg.early_stop_patience,\n monitor=cfg.monitor,\n )\n\n if bad_epochs >= cfg.early_stop_patience:\n print(f\"⏹️ Early stop at epoch {epoch+1}. Best {cfg.monitor}={best_metric:.6g} at epoch {best_epoch}.\")\n break\n\n finally:\n with open(history_path, \"w\") as f:\n json.dump(history, f, indent=2)\n\n # load best weights before test evaluation\n if os.path.exists(best_ckpt_path):\n self.load_checkpoint(best_ckpt_path, strict=True)\n \n # ---- test evaluation ONCE on best checkpoint ----\n if test_loader is not None:\n te_loss, te_auc, y_te, p_te = self.validate(\n test_loader, loss_fn, cfg, epoch=0, split_name=\"test\", dashboard=None\n )\n history[\"test\"] = {\"test_loss\": float(te_loss), \"test_auc\": float(te_auc)}\n self.writer.add_scalar(\"test/loss\", te_loss, 0)\n self.writer.add_scalar(\"test/auc\", te_auc, 0)\n\n np.savez_compressed(os.path.join(save_dir, \"test_preds.npz\"), y=y_te, p=p_te)\n print(f\"🧪 Test: loss={te_loss:.6f} | auc={te_auc:.6f}\")\n\n with open(history_path, \"w\") as f:\n json.dump(history, f, indent=2)\n \n return history\n\n @torch.no_grad()\n def predict_for_roc(self, loader, cfg: TrainConfig):\n loss_fn = self._loss_fn(cfg) # not used, but keeps symmetry\n self.model.eval()\n all_probs, all_y = [], []\n for X, y in loader:\n X = X.to(self.device, non_blocking=True)\n y = y.to(self.device, non_blocking=True).float().view(-1)\n out = self.model(X).view(-1)\n probs = self._to_probs(out, cfg)\n all_probs.append(probs.cpu())\n all_y.append(y.cpu())\n return torch.cat(all_y).numpy(), torch.cat(all_probs).numpy()\n\n def close(self):\n self.writer.close()\n","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:55:57.236335Z","iopub.status.busy":"2026-02-11T10:55:57.236110Z","iopub.status.idle":"2026-02-11T10:55:57.289576Z","shell.execute_reply":"2026-02-11T10:55:57.288948Z"},"papermill":{"duration":0.06342,"end_time":"2026-02-11T10:55:57.291139","exception":false,"start_time":"2026-02-11T10:55:57.227719","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 4.2 Visualization Utilities","metadata":{"papermill":{"duration":0.006946,"end_time":"2026-02-11T10:55:57.305390","exception":false,"start_time":"2026-02-11T10:55:57.298444","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def plot_history(history: Dict[str, Any], title: str = \"Training Curves\"):\n epochs = range(1, len(history[\"train_loss\"]) + 1)\n\n plt.figure(figsize=(10, 4))\n plt.subplot(1, 2, 1)\n plt.plot(epochs, history[\"train_loss\"], label=\"Train\")\n plt.plot(epochs, history[\"val_loss\"], label=\"Val\")\n plt.xlabel(\"Epoch\"); plt.ylabel(\"Loss\"); plt.title(f\"{title} - Loss\")\n plt.legend()\n\n plt.subplot(1, 2, 2)\n plt.plot(epochs, history[\"train_auc\"], label=\"Train\")\n plt.plot(epochs, history[\"val_auc\"], label=\"Val\")\n plt.xlabel(\"Epoch\"); plt.ylabel(\"AUC\"); plt.title(f\"{title} - AUC\")\n plt.legend()\n\n plt.tight_layout()\n plt.show()\n\n\ndef plot_roc_curve(y_true, y_prob, title=\"ROC Curve\"):\n \"\"\"\n Plots ROC curve + prints AUC.\n y_true: array-like of {0,1}\n y_prob: array-like of probabilities in [0,1]\n \"\"\"\n y_true = np.asarray(y_true).astype(int)\n y_prob = np.asarray(y_prob).astype(float)\n\n # Guard against degenerate labels\n if len(np.unique(y_true)) < 2:\n raise ValueError(\"ROC/AUC is undefined: y_true has only one class.\")\n\n fpr, tpr, _ = roc_curve(y_true, y_prob)\n auc_val = roc_auc_score(y_true, y_prob)\n\n plt.figure(figsize=(6, 5))\n plt.plot(fpr, tpr, label=f\"AUC = {auc_val:.4f}\")\n plt.plot([0, 1], [0, 1], linestyle=\"--\", label=\"Random\")\n plt.xlabel(\"False Positive Rate\")\n plt.ylabel(\"True Positive Rate\")\n plt.title(title)\n plt.grid(True, alpha=0.3)\n plt.legend()\n plt.tight_layout()\n plt.show()\n\n return auc_val","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:55:57.320759Z","iopub.status.busy":"2026-02-11T10:55:57.320466Z","iopub.status.idle":"2026-02-11T10:55:57.328233Z","shell.execute_reply":"2026-02-11T10:55:57.327623Z"},"papermill":{"duration":0.017216,"end_time":"2026-02-11T10:55:57.329790","exception":false,"start_time":"2026-02-11T10:55:57.312574","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"---\n\n## 5. Hyperparameter Tuning","metadata":{"papermill":{"duration":0.006775,"end_time":"2026-02-11T10:55:57.343623","exception":false,"start_time":"2026-02-11T10:55:57.336848","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"### 5.1 Grid Search Implementation","metadata":{"papermill":{"duration":0.006619,"end_time":"2026-02-11T10:55:57.357021","exception":false,"start_time":"2026-02-11T10:55:57.350402","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def grid_search(\n build_model_fn,\n base_model_kwargs,\n param_grid,\n train_loader,\n val_loader,\n cfg_template,\n test_loader=None,\n log_dir=\"runs/grid\",\n model_name=\"model\",\n score_key=None,\n resume=False,\n cfg_keys=(\"lr\", \"weight_decay\", \"early_stop_patience\"), # keys that belong to cfg\n):\n keys = list(param_grid.keys())\n combos = list(itertools.product(*[param_grid[k] for k in keys]))\n\n results = []\n best_row = None\n\n sel_metric = score_key or cfg_template.monitor\n sel_mode = cfg_template.monitor_mode\n\n def better(a, b):\n if math.isnan(a): return False\n if math.isnan(b): return True\n return (a < b) if sel_mode == \"min\" else (a > b)\n\n os.makedirs(log_dir, exist_ok=True)\n\n for values in combos:\n params = dict(zip(keys, values))\n\n # split params: cfg vs model\n cfg_params = {k: params[k] for k in params if k in cfg_keys}\n model_params = {k: params[k] for k in params if k not in cfg_keys}\n\n model_kwargs = {**base_model_kwargs, **model_params}\n\n # clone cfg and apply cfg params\n cfg = TrainConfig(**vars(cfg_template))\n for k, v in cfg_params.items():\n setattr(cfg, k, v)\n\n run_name = (\n f\"{model_name}_\"\n + \"_\".join([f\"{k}={v}\" for k, v in params.items()])\n + \"_\"\n + datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n )\n\n trainer = None\n try:\n model = build_model_fn(**model_kwargs)\n trainer = Trainer(model, log_dir=log_dir, run_name=run_name)\n\n save_dir = trainer.run_path\n last_ckpt = os.path.join(save_dir, \"last.ckpt\")\n resume_from = last_ckpt if (resume and os.path.exists(last_ckpt)) else None\n\n history = trainer.fit(\n train_loader=train_loader,\n val_loader=val_loader,\n test_loader=test_loader,\n cfg=cfg, # IMPORTANT: per-trial cfg\n save_dir=save_dir,\n resume_from=resume_from,\n use_rich_dashboard=False,\n prefer_notebook_dashboard=True,\n extra_meta={\"model_kwargs\": model_kwargs, \"params\": params, \"cfg_params\": cfg_params},\n )\n\n best_info = history.get(\"best\", {})\n best_epoch = int(best_info.get(\"epoch\", 0))\n best_train_loss = float(best_info.get(\"train_loss\", float(\"nan\")))\n best_train_auc = float(best_info.get(\"train_auc\", float(\"nan\")))\n best_val_loss = float(best_info.get(\"val_loss\", float(\"nan\")))\n best_val_auc = float(best_info.get(\"val_auc\", float(\"nan\")))\n\n if sel_metric == cfg.monitor:\n sel_value = float(best_info.get(\"value\", float(\"nan\")))\n else:\n idx = max(best_epoch - 1, 0)\n if sel_metric == \"val_loss\":\n sel_value = float(history[\"val_loss\"][idx])\n elif sel_metric == \"val_auc\":\n sel_value = float(history[\"val_auc\"][idx])\n else:\n raise ValueError(f\"Unsupported score_key: {sel_metric}\")\n\n test_info = history.get(\"test\") or {}\n row = {\n \"params\": params,\n \"selection_metric\": sel_metric,\n \"selection_value\": sel_value,\n \"best_epoch\": best_epoch,\n \"best_train_loss\": best_train_loss,\n \"best_train_auc\": best_train_auc,\n \"best_val_loss\": best_val_loss,\n \"best_val_auc\": best_val_auc,\n \"test_loss\": test_info.get(\"test_loss\", float(\"nan\")),\n \"test_auc\": test_info.get(\"test_auc\", float(\"nan\")),\n \"run_path\": save_dir,\n \"status\": \"ok\",\n }\n results.append(row)\n\n if best_row is None or better(sel_value, best_row[\"selection_value\"]):\n best_row = row\n\n finally:\n if trainer is not None:\n trainer.close()\n try:\n del model\n except Exception:\n pass\n if torch.cuda.is_available():\n torch.cuda.empty_cache()\n\n summary_path = os.path.join(log_dir, f\"{model_name}_grid_summary_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json\")\n with open(summary_path, \"w\") as f:\n json.dump({\"results\": results, \"best\": best_row}, f, indent=2)\n\n return results, best_row","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:55:57.372803Z","iopub.status.busy":"2026-02-11T10:55:57.372528Z","iopub.status.idle":"2026-02-11T10:55:57.386372Z","shell.execute_reply":"2026-02-11T10:55:57.385582Z"},"papermill":{"duration":0.023886,"end_time":"2026-02-11T10:55:57.387882","exception":false,"start_time":"2026-02-11T10:55:57.363996","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 5.2 One-Factor-at-a-Time (OFAT) Sweep","metadata":{"papermill":{"duration":0.007663,"end_time":"2026-02-11T10:55:57.403352","exception":false,"start_time":"2026-02-11T10:55:57.395689","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def ofat_sweep(\n build_model_fn: Callable[..., nn.Module],\n baseline_kwargs: Dict[str, Any],\n sweep_param: str,\n sweep_values: List[Any],\n train_loader,\n val_loader,\n cfg_template: TrainConfig,\n log_dir: str = \"runs/ofat\",\n model_name: str = \"model\",\n):\n points = []\n for v in sweep_values:\n kwargs = dict(baseline_kwargs)\n kwargs[sweep_param] = v\n\n run_name = f\"{model_name}_OFAT_{sweep_param}={v}_{datetime.now().strftime('%Y%m%d_%H%M%S')}\"\n trainer = None\n try:\n model = build_model_fn(**kwargs)\n trainer = Trainer(model, log_dir=log_dir, run_name=run_name)\n history = trainer.fit(\n train_loader=train_loader,\n val_loader=val_loader,\n cfg=cfg_template,\n save_dir=trainer.run_path,\n resume_from=None,\n use_rich_dashboard=False,\n extra_meta={\"model_kwargs\": kwargs, \"sweep_param\": sweep_param, \"sweep_value\": v},\n )\n\n best_info = history.get(\"best\", {})\n points.append({\n \"value\": v,\n \"best_epoch\": int(best_info.get(\"epoch\", 0)),\n \"best_metric\": float(best_info.get(\"value\", float(\"nan\"))),\n \"best_train_loss\": float(best_info.get(\"train_loss\", float(\"nan\"))),\n \"best_train_auc\": float(best_info.get(\"train_auc\", float(\"nan\"))),\n \"best_val_loss\": float(best_info.get(\"val_loss\", float(\"nan\"))),\n \"best_val_auc\": float(best_info.get(\"val_auc\", float(\"nan\"))),\n \"best_test_loss\": float(best_info.get(\"test_loss\", float(\"nan\"))),\n \"best_test_auc\": float(best_info.get(\"test_auc\", float(\"nan\"))),\n \"run_path\": trainer.run_path,\n \"history\": history, # keep if you want; can be big\n })\n\n finally:\n if trainer is not None:\n trainer.close()\n try:\n del model\n except Exception:\n pass\n if torch.cuda.is_available():\n torch.cuda.empty_cache()\n\n return points\n\n\ndef plot_ofat(points: List[Dict[str, Any]], param_name: str, metric: str = \"best_val_loss\"):\n xs = [p[\"value\"] for p in points]\n ys = [p.get(metric, float(\"nan\")) for p in points]\n plt.figure(figsize=(6, 4))\n plt.plot(xs, ys, marker=\"o\")\n plt.xlabel(param_name)\n plt.ylabel(metric)\n plt.title(f\"OFAT Sweep: {param_name}\")\n plt.show()\n","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:55:57.419617Z","iopub.status.busy":"2026-02-11T10:55:57.418793Z","iopub.status.idle":"2026-02-11T10:55:57.428996Z","shell.execute_reply":"2026-02-11T10:55:57.428267Z"},"papermill":{"duration":0.01984,"end_time":"2026-02-11T10:55:57.430434","exception":false,"start_time":"2026-02-11T10:55:57.410594","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"---\n\n## 6. Experiments and Results","metadata":{"papermill":{"duration":0.007023,"end_time":"2026-02-11T10:55:57.444513","exception":false,"start_time":"2026-02-11T10:55:57.437490","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"### 6.1 Data Loaders","metadata":{"papermill":{"duration":0.007005,"end_time":"2026-02-11T10:55:57.458460","exception":false,"start_time":"2026-02-11T10:55:57.451455","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Create data loaders\nbatch_size = 1024\ntrain_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\nval_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)\ntest_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:55:57.474342Z","iopub.status.busy":"2026-02-11T10:55:57.473556Z","iopub.status.idle":"2026-02-11T10:55:57.477935Z","shell.execute_reply":"2026-02-11T10:55:57.477374Z"},"papermill":{"duration":0.013792,"end_time":"2026-02-11T10:55:57.479340","exception":false,"start_time":"2026-02-11T10:55:57.465548","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 6.2 Hyperparameter Grid Search","metadata":{"papermill":{"duration":0.006997,"end_time":"2026-02-11T10:55:57.493184","exception":false,"start_time":"2026-02-11T10:55:57.486187","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Grid 3: Regularization Search (drop_rate and weight_decay)\ncfg = TrainConfig(\n epochs=4,\n lr=2e-3,\n weight_decay=1e-5,\n use_manual_l2=False,\n monitor=\"val_auc\",\n monitor_mode=\"max\",\n early_stop_patience=1,\n scheduler=\"plateau\",\n use_amp=True,\n model_returns_probs=False,\n)\n\nbase_model_kwargs = dict(\n field_dims=field_dims,\n embed_dims=8,\n attention_dim=8,\n mlp_dims=(400, 400, 400),\n drop_rate=0.1,\n)\n\nparam_grid = {\n \"drop_rate\": [0.05, 0.1, 0.2, 0.3],\n \"weight_decay\": [0.0, 1e-6, 1e-5]\n}\n\nresults, best = grid_search(\n build_model_fn=DeepAFM,\n base_model_kwargs=base_model_kwargs,\n param_grid=param_grid,\n train_loader=train_loader,\n val_loader=val_loader,\n test_loader=test_loader,\n cfg_template=cfg,\n log_dir=\"runs/grid_deepafm\",\n model_name=\"deepafm\",\n score_key=\"val_auc\",\n resume=False,\n)","metadata":{"execution":{"iopub.execute_input":"2026-02-11T10:55:57.508428Z","iopub.status.busy":"2026-02-11T10:55:57.508165Z","iopub.status.idle":"2026-02-11T15:21:36.883340Z","shell.execute_reply":"2026-02-11T15:21:36.882552Z"},"papermill":{"duration":15939.384923,"end_time":"2026-02-11T15:21:36.884947","exception":false,"start_time":"2026-02-11T10:55:57.500024","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 6.3 Final Model Training","metadata":{"papermill":{"duration":0.007669,"end_time":"2026-02-11T15:21:36.900958","exception":false,"start_time":"2026-02-11T15:21:36.893289","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Final best tuned model run\ncfg = TrainConfig(\n epochs=12,\n lr=0.002,\n weight_decay=0,\n use_manual_l2=False,\n monitor=\"val_auc\",\n monitor_mode=\"max\",\n early_stop_patience=2,\n scheduler=\"plateau\",\n use_amp=True,\n model_returns_probs=False,\n)\n\nbase_kwargs = dict(\n field_dims=field_dims,\n embed_dims=8,\n attention_dim=8,\n mlp_dims=(400, 400, 400),\n drop_rate=0.05,\n)\n\nmodel = DeepAFM(**base_kwargs)\ntrainer = Trainer(model, log_dir=\"runs/final\", run_name=\"deepafm_tuned_parameters\")\n\nhistory_afm_mlp = history = trainer.fit(\n train_loader=train_loader,\n val_loader=val_loader,\n cfg=cfg,\n test_loader=test_loader,\n use_rich_dashboard=False,\n prefer_notebook_dashboard=True \n)\n\n# --- 2) Plot training curves (epoch-level) ---\nplot_history(history, title=\"AFM-MLP\")\n\n# --- 3) Get y_true and y_prob for ROC (choose which split you want) ---\n# Use test set for paper-quality evaluation\ny_true, y_prob = trainer.predict_for_roc(test_loader, cfg)\n\n# --- 4) Plot ROC curve ---\nauc_val = plot_roc_curve(y_true, y_prob, title=\"AFM-MLP - Test ROC\")\nprint(\"Test AUC:\", auc_val)\n\n# (Optional) Plot PR curve + calibration etc if you included evaluate_and_plot_all earlier:\n# evaluate_and_plot_all(y_true, y_prob, title_prefix=\"DeepFM Test\")\n\n","metadata":{"execution":{"iopub.execute_input":"2026-02-11T15:21:36.918088Z","iopub.status.busy":"2026-02-11T15:21:36.917303Z","iopub.status.idle":"2026-02-11T15:54:55.189578Z","shell.execute_reply":"2026-02-11T15:54:55.188891Z"},"papermill":{"duration":1998.293162,"end_time":"2026-02-11T15:54:55.201755","exception":false,"start_time":"2026-02-11T15:21:36.908593","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 6.4 Over-parameterized Weight Decay Analysis","metadata":{"papermill":{"duration":0.009443,"end_time":"2026-02-11T15:54:55.220640","exception":false,"start_time":"2026-02-11T15:54:55.211197","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Over-parameterized model: Weight Decay Analysis (embed=32, attn=32)\ncfg = TrainConfig(\n epochs=4,\n lr=2e-3,\n weight_decay=1e-5,\n use_manual_l2=False,\n monitor=\"val_auc\",\n monitor_mode=\"max\",\n early_stop_patience=1,\n scheduler=\"plateau\",\n use_amp=True,\n model_returns_probs=False,\n)\n\nbase_model_kwargs = dict(\n field_dims=field_dims,\n embed_dims=32,\n attention_dim=32,\n mlp_dims=(400, 400, 400),\n drop_rate=0,\n)\n\nparam_grid = {\"weight_decay\": [0, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 0.1]}\n\nresults, best = grid_search(\n build_model_fn=DeepAFM,\n base_model_kwargs=base_model_kwargs,\n param_grid=param_grid,\n train_loader=train_loader,\n val_loader=val_loader,\n test_loader=test_loader,\n cfg_template=cfg,\n log_dir=\"runs/grid_deepafm\",\n model_name=\"deepafm\",\n score_key=\"val_auc\",\n resume=False,\n)","metadata":{"execution":{"iopub.execute_input":"2026-02-11T15:54:55.241411Z","iopub.status.busy":"2026-02-11T15:54:55.240773Z","iopub.status.idle":"2026-02-11T18:21:34.675303Z","shell.execute_reply":"2026-02-11T18:21:34.674507Z"},"papermill":{"duration":8799.446961,"end_time":"2026-02-11T18:21:34.676930","exception":false,"start_time":"2026-02-11T15:54:55.229969","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Print results\nfor r in results:\n print(\n r[\"params\"],\n \"train_auc:\", f\"{r['best_train_auc']:.4f}\",\n \"val_auc:\", f\"{r['best_val_auc']:.4f}\",\n \"test_auc:\", f\"{r['test_auc']:.4f}\",\n )\n\nprint(\"\\nBEST RUN:\")\nprint(best)","metadata":{"execution":{"iopub.execute_input":"2026-02-11T18:21:34.699574Z","iopub.status.busy":"2026-02-11T18:21:34.699316Z","iopub.status.idle":"2026-02-11T18:21:34.703931Z","shell.execute_reply":"2026-02-11T18:21:34.703330Z"},"papermill":{"duration":0.017501,"end_time":"2026-02-11T18:21:34.705497","exception":false,"start_time":"2026-02-11T18:21:34.687996","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 7. Baseline Model Comparison\n\nTraining Logistic Regression, FM, AFM, and DeepFM with matched hyperparameters (embed_dims=8, 6 epochs) for controlled comparison against AFM-MLP.","metadata":{"papermill":{"duration":0.01039,"end_time":"2026-02-11T18:21:34.726291","exception":false,"start_time":"2026-02-11T18:21:34.715901","status":"completed"},"tags":[]}},{"cell_type":"code","source":"baseline_cfg = TrainConfig(\n epochs=6,\n lr=1e-3,\n weight_decay=0.0,\n use_manual_l2=False,\n monitor=\"val_auc\",\n monitor_mode=\"max\",\n early_stop_patience=6,\n scheduler=\"none\",\n use_amp=False,\n model_returns_probs=True,\n)\n\nbaseline_models = {\n \"Logistic Regression\": {\n \"build_fn\": LogisticRegressionModel,\n \"kwargs\": dict(field_dims=field_dims),\n },\n \"FM\": {\n \"build_fn\": FM,\n \"kwargs\": dict(field_dims=field_dims, embed_dims=8),\n },\n \"AFM\": {\n \"build_fn\": AFM,\n \"kwargs\": dict(field_dims=field_dims, embed_dims=8, attention_dims=8),\n },\n \"DeepFM\": {\n \"build_fn\": DeepFM,\n \"kwargs\": dict(field_dims=field_dims, embed_dims=8, mlp_dims=[400, 400, 400], drop_rate=0.1),\n },\n}\n\nbaseline_results = {}\n\nfor name, spec in baseline_models.items():\n print(f\"\\n{'='*60}\")\n print(f\"Training: {name}\")\n print(f\"{'='*60}\")\n \n model = spec[\"build_fn\"](**spec[\"kwargs\"])\n trainer = Trainer(model, log_dir=\"runs/baselines\", run_name=f\"baseline_{name.replace(' ', '_')}\")\n \n history = trainer.fit(\n train_loader=train_loader,\n val_loader=val_loader,\n test_loader=test_loader,\n cfg=baseline_cfg,\n use_rich_dashboard=False,\n prefer_notebook_dashboard=True,\n )\n \n test_info = history.get(\"test\", {})\n best_info = history.get(\"best\", {})\n \n baseline_results[name] = {\n \"test_auc\": test_info.get(\"test_auc\", float(\"nan\")),\n \"test_loss\": test_info.get(\"test_loss\", float(\"nan\")),\n \"best_val_auc\": best_info.get(\"val_auc\", float(\"nan\")),\n \"best_val_loss\": best_info.get(\"val_loss\", float(\"nan\")),\n \"history\": history,\n }\n \n trainer.close()\n del model\n if torch.cuda.is_available():\n torch.cuda.empty_cache()\n \n print(f\"{name} -> Test AUC: {test_info.get('test_auc', 'N/A'):.4f} | Test Log-loss: {test_info.get('test_loss', 'N/A'):.4f}\")","metadata":{"papermill":{"duration":0.503836,"end_time":"2026-02-11T18:21:35.240336","exception":true,"start_time":"2026-02-11T18:21:34.736500","status":"failed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"baseline_results[\"AFM-MLP\"] = {\n \"test_auc\": history_afm_mlp[\"test\"][\"test_auc\"],\n \"test_loss\": history_afm_mlp[\"test\"][\"test_loss\"],\n}\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"FINAL MODEL COMPARISON\")\nprint(\"=\"*60)\nprint(f\"{'Model':<25} {'AUC':>10} {'Log-loss':>10}\")\nprint(\"-\"*45)\nfor name in [\"Logistic Regression\", \"FM\", \"AFM\", \"DeepFM\", \"AFM-MLP\"]:\n r = baseline_results[name]\n print(f\"{name:<25} {r['test_auc']:>10.4f} {r['test_loss']:>10.4f}\")","metadata":{"papermill":{"duration":null,"end_time":null,"exception":null,"start_time":null,"status":"pending"},"tags":[]},"outputs":[],"execution_count":null}]}