{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":84493,"databundleVersionId":11305158,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":11988796,"sourceType":"datasetVersion","datasetId":7539809},{"sourceId":203900450,"sourceType":"kernelVersion"},{"sourceId":224157538,"sourceType":"kernelVersion"}],"dockerImageVersionId":31040,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!git clone https://github.com/yandex-research/tabm\n!pip install rtdl_num_embeddings","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-05-28T22:50:07.133769Z","iopub.execute_input":"2025-05-28T22:50:07.134041Z","iopub.status.idle":"2025-05-28T22:51:46.743793Z","shell.execute_reply.started":"2025-05-28T22:50:07.134023Z","shell.execute_reply":"2025-05-28T22:51:46.742676Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, sys, gc\nimport pickle\nimport dill\nimport numpy as np\nimport pandas as pd\nimport polars as pl\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pytorch_lightning import (LightningDataModule, LightningModule, Trainer)\n\nfrom sklearn.metrics import r2_score\n\nimport torch.optim\nfrom torch.utils.data import Dataset, DataLoader, TensorDataset\nfrom sklearn.model_selection import train_test_split\nimport math\nfrom tqdm import tqdm\nfrom collections import OrderedDict\nfrom tabm_reference import Model, make_parameter_groups\n\nimport warnings\nimport joblib\nfrom pytorch_lightning.callbacks import Callback\nimport gc\n\nimport lightgbm as lgb\nfrom lightgbm import LGBMRegressor, Booster\nfrom xgboost import XGBRegressor\nfrom catboost import CatBoostRegressor\n\nimport warnings\nwarnings.filterwarnings('ignore')\npd.options.display.max_columns = None","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T22:52:06.692364Z","iopub.execute_input":"2025-05-28T22:52:06.692865Z","iopub.status.idle":"2025-05-28T22:52:06.698277Z","shell.execute_reply.started":"2025-05-28T22:52:06.692850Z","shell.execute_reply":"2025-05-28T22:52:06.697507Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 1. Fully-Connected NN + XGB (5-fold)  ","metadata":{}},{"cell_type":"markdown","source":"## 1.1 NN","metadata":{}},{"cell_type":"code","source":"class CFG:\n    seed = 42  \n    target_col = \"responder_6\"  \n    feature_cols = [f\"feature_{idx:02d}\" for idx in range(79)] + [\n        f\"responder_{idx}_lag_1\" for idx in range(9)]\n    model_path = \"/kaggle/input/trained\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T22:52:15.579069Z","iopub.execute_input":"2025-05-28T22:52:15.579345Z","iopub.status.idle":"2025-05-28T22:52:15.585286Z","shell.execute_reply.started":"2025-05-28T22:52:15.579320Z","shell.execute_reply":"2025-05-28T22:52:15.584213Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load validation data\nvalid = pl.scan_parquet(\n    f\"/kaggle/input/js24-preprocessing-create-lags/validation.parquet/\"\n).collect().to_pandas()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T22:53:06.833702Z","iopub.execute_input":"2025-05-28T22:53:06.834012Z","iopub.status.idle":"2025-05-28T22:53:11.136818Z","shell.execute_reply.started":"2025-05-28T22:53:06.833993Z","shell.execute_reply":"2025-05-28T22:53:11.135903Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# r2 val eval function\ndef r2_val(y_true, y_pred, sample_weight):\n    nom = np.average((y_pred - y_true) ** 2, weights=sample_weight)\n    denom = (np.average((y_true) ** 2, weights=sample_weight) + 1e-38)\n    r2 = 1 - nom/denom \n    return r2\n\nclass NN(LightningModule):\n    # define NN specs to load saved params (checkpoint)\n    def __init__(self, input_dim, hidden_dims, dropouts, lr, weight_decay):\n        super().__init__()\n        self.save_hyperparameters()\n        \n        # Build network architecture\n        layers = []\n        in_dim = input_dim\n        for i, hidden_dim in enumerate(hidden_dims):\n            layers.append(nn.BatchNorm1d(in_dim))  # Batch normalization\n            if i > 0:\n                layers.append(nn.SiLU())  # SiLU activation (except first layer)\n            if i < len(dropouts):\n                layers.append(nn.Dropout(dropouts[i]))  # Dropout for regularization\n            layers.append(nn.Linear(in_dim, hidden_dim))  # Linear layer\n            in_dim = hidden_dim\n            \n        # Output layer\n        layers.append(nn.Linear(in_dim, 1))\n        layers.append(nn.Tanh())  # Tanh activation for bounded output\n        \n        self.model = nn.Sequential(*layers)\n        self.lr = lr\n        self.weight_decay = weight_decay\n        self.validation_step_outputs = []\n\n    def forward(self, x):\n        return 5 * self.model(x).squeeze(-1)  # Scale output to [-5, 5] range\n\n    def training_step(self, batch):\n        x, y, w = batch\n        y_hat = self(x)\n        loss = F.mse_loss(y_hat, y, reduction='none') * w  # Weighted MSE loss\n        loss = loss.mean()\n        self.log('train_loss', loss, on_step=False, on_epoch=True, batch_size=x.size(0))\n        return loss\n\n    def validation_step(self, batch):\n        x, y, w = batch\n        y_hat = self(x)\n        loss = F.mse_loss(y_hat, y, reduction='none') * w\n        loss = loss.mean()\n        self.log('val_loss', loss, on_step=False, on_epoch=True, batch_size=x.size(0))\n        self.validation_step_outputs.append((y_hat, y, w))\n        return loss\n\n    def on_validation_epoch_end(self):\n        if not self.trainer.sanity_checking:\n            y = torch.cat([x[1] for x in self.validation_step_outputs]\n                         ).cpu().numpy()\n            prob = torch.cat([x[0] for x in self.validation_step_outputs]\n                            ).cpu().numpy()\n            weights = torch.cat([x[2] for x in self.validation_step_outputs]\n                               ).cpu().numpy()\n            val_r_square = r2_val(y, prob, weights)\n            self.log(\n                \"val_r_square\", val_r_square, prog_bar=True, \n                on_step=False, on_epoch=True)\n        self.validation_step_outputs.clear()\n\n    def configure_optimizers(self):\n        optimizer = torch.optim.Adam(\n            self.parameters(), lr=self.lr, \n            weight_decay=self.weight_decay)\n        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n            optimizer, \n            mode='min', \n            factor=0.5, \n            patience=5, \n            verbose=True\n        )\n        return {\n            'optimizer': optimizer,\n            'lr_scheduler': {\n                'scheduler': scheduler,\n                'monitor': 'val_loss',\n            }\n        }\n\n    def on_train_epoch_end(self):\n        if not self.trainer.sanity_checking:\n            epoch = self.trainer.current_epoch\n            metrics = {k: v.item() if isinstance(v, torch.Tensor) else v \n                      for k, v in self.trainer.logged_metrics.items()}\n            formatted_metrics = {k: f\"{v:.5f}\" for k, v in metrics.items()}\n            print(f\"Epoch {epoch}: {formatted_metrics}\")\n\n# Load the trained models (5-fold cv) from the '/input/trained'\n# saved models were uploaded to Kaggle as the Dataset (Public)\n# this is a common trick for API submit like this competition\n\nN_folds = 5\nnn_models = []\nfor fold in range(N_folds):\n    checkpoint_path = f\"{CFG.model_path}/nn_{fold}.model\"\n    model = NN.load_from_checkpoint(checkpoint_path)\n    nn_models.append(model)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T22:53:56.293790Z","iopub.execute_input":"2025-05-28T22:53:56.294073Z","iopub.status.idle":"2025-05-28T22:53:57.045958Z","shell.execute_reply.started":"2025-05-28T22:53:56.294056Z","shell.execute_reply":"2025-05-28T22:53:57.045111Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 1.2 XGB","metadata":{}},{"cell_type":"code","source":"# Load XGBoost model\nxgb_models = []\nfor fold in range(N_folds): \n    with open(f\"{CFG.model_path}/xgb_result{fold}.pkl\", \"rb\") as fp:\n        result = pickle.load(fp)\n        xgb_models.append(result[\"model\"])\n    xgb_feature_cols = [\"symbol_id\", \"time_id\"] + CFG.feature_cols\n\ndisplay(xgb_models[0]) # show one","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T22:56:19.149477Z","iopub.execute_input":"2025-05-28T22:56:19.149764Z","iopub.status.idle":"2025-05-28T22:56:19.234686Z","shell.execute_reply.started":"2025-05-28T22:56:19.149746Z","shell.execute_reply":"2025-05-28T22:56:19.233681Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 1.3 Blend NN with XGB","metadata":{}},{"cell_type":"code","source":"X_valid = valid[ CFG.feature_cols ]\ny_valid = valid[ CFG.target_col ]\nw_valid = valid[ \"weight\" ]\nX_valid = X_valid.fillna(method = 'ffill').fillna(0)\nX_valid.shape, y_valid.shape, w_valid.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T22:56:44.208648Z","iopub.execute_input":"2025-05-28T22:56:44.208977Z","iopub.status.idle":"2025-05-28T22:56:45.516417Z","shell.execute_reply.started":"2025-05-28T22:56:44.208952Z","shell.execute_reply":"2025-05-28T22:56:45.515669Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"y_pred_valid_nn = np.zeros(y_valid.shape)\n\nwith torch.no_grad():\n    for nn_model in nn_models:\n        nn_model.eval()\n        y_pred_valid_nn += model(\n            torch.FloatTensor(X_valid.values)\n        ).cpu().numpy() / len(nn_models)\n\nvalid_score = r2_score( y_valid, y_pred_valid_nn, sample_weight=w_valid )\nvalid_score\nprint(f'NN R2 score on valid dataset: {valid_score}')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T22:58:53.479499Z","iopub.execute_input":"2025-05-28T22:58:53.479770Z","iopub.status.idle":"2025-05-28T23:00:17.764582Z","shell.execute_reply.started":"2025-05-28T22:58:53.479752Z","shell.execute_reply":"2025-05-28T23:00:17.763903Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X_valid = valid[ [\"symbol_id\", \"time_id\"] + CFG.feature_cols ]\ny_valid = valid[ CFG.target_col ]\nw_valid = valid[ \"weight\" ]\n\ny_pred_valid_xgb = xgb_models[0].predict(X_valid)\n\nvalid_score = r2_score( y_valid, y_pred_valid_xgb, sample_weight=w_valid )\nprint(f'XGB R2 score on valid dataset: {valid_score}')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T23:03:34.428963Z","iopub.execute_input":"2025-05-28T23:03:34.429264Z","iopub.status.idle":"2025-05-28T23:03:36.906598Z","shell.execute_reply.started":"2025-05-28T23:03:34.429246Z","shell.execute_reply":"2025-05-28T23:03:36.905991Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 2. Ridge ","metadata":{}},{"cell_type":"code","source":"import dill\nis_local = os.environ.get(\"DOCKER_USING\", \"\") == \"LOCAL\"\ndef load_from_dill():\n    model_object = None\n    with open(\"/kaggle/input/jsridgev01011635/Ridge.dill\", \"rb\") as file_handle:\n        model_object = dill.load(file_handle)\n    return model_object","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T23:06:38.731176Z","iopub.execute_input":"2025-05-28T23:06:38.731593Z","iopub.status.idle":"2025-05-28T23:06:38.736621Z","shell.execute_reply.started":"2025-05-28T23:06:38.731568Z","shell.execute_reply":"2025-05-28T23:06:38.735762Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def predict_ridge(test, lags):\n    cols = [f'feature_{i:02}' for i in range(79)]\n    predictions = test.select(\n        'row_id',\n        pl.lit(0.0).alias('responder_6'),\n    )\n    ######################\n    rdg_model = load_from_dill()\n    test_preds = rdg_model.predict(test[cols].to_pandas().fillna(3).values)\n    return test_preds\n####################","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T23:06:36.609500Z","iopub.execute_input":"2025-05-28T23:06:36.609801Z","iopub.status.idle":"2025-05-28T23:06:36.615057Z","shell.execute_reply.started":"2025-05-28T23:06:36.609783Z","shell.execute_reply":"2025-05-28T23:06:36.614101Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 3. TabM","metadata":{}},{"cell_type":"code","source":"# Create list of feature names from 0-78, excluding feature_61\nfeature_list = [f\"feature_{idx:02d}\" for idx in range(79) if idx != 61]\n\n# Define target column name\ntarget_col = \"responder_6\" \n\n# Create list of features for testing, combining feature_list with lagged responder features\nfeature_test = feature_list + [f\"responder_{idx}_lag_1\" for idx in range(9)] \n\n# Define categorical features\nfeature_cat = [\"feature_09\", \"feature_10\", \"feature_11\"]\n\n# Define continuous features by excluding categorical ones from feature_test\nfeature_cont = [item for item in feature_test if item not in feature_cat]\n\n# Set batch size for model training\nbatch_size = 8192\n\n# Create list of features to standardize (continuous features + lagged responder features)\nstd_feature = [i for i in feature_list if i not in feature_cat] + [\n    f\"responder_{idx}_lag_1\" for idx in range(9)]\n\n# Load pre-computed statistics for standard scaling\ndata_stats = joblib.load(\"/kaggle/input/trained/data_stats.pkl\")\nmeans = data_stats['mean']\nstds = data_stats['std']","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T23:42:43.299561Z","iopub.execute_input":"2025-05-28T23:42:43.299885Z","iopub.status.idle":"2025-05-28T23:42:43.309206Z","shell.execute_reply.started":"2025-05-28T23:42:43.299863Z","shell.execute_reply":"2025-05-28T23:42:43.308435Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def standardize(df, feature_cols, means, stds):\n    return df.with_columns([\n        ((pl.col(col) - means[col]) / stds[col]).alias(col) for col in feature_cols\n    ])\n\n# Dictionary mappings for categorical variables encoding\ncategory_mappings = {\n    'feature_09': {\n        2: 0, 4: 1, 9: 2, 11: 3, 12: 4, 14: 5, 15: 6, \n        25: 7, 26: 8, 30: 9, 34: 10, 42: 11, 44: 12, 46: 13, \n        49: 14, 50: 15, 57: 16, 64: 17, 68: 18, 70: 19, 81: 20, 82: 21},\n    \n    'feature_10': {\n     1: 0, 2: 1, 3: 2, 4: 3, 5: 4, \n     6: 5, 7: 6, 10: 7, 12: 8},\n    \n    'feature_11': {9: 0, 11: 1, 13: 2, 16: 3, 24: 4, 25: 5, 34: 6, \n                40: 7, 48: 8, 50: 9, 59: 10, 62: 11, 63: 12, 66: 13,\n                76: 14, 150: 15, 158: 16, 159: 17, 171: 18, 195: 19, \n                214: 20, 230: 21, 261: 22, 297: 23, 336: 24, 376: 25, \n                388: 26, 410: 27, 522: 28, 534: 29, 539: 30},\n    \n    'symbol_id': {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, \n               6: 6, 7: 7, 8: 8, 9: 9, 10: 10, 11: 11, \n               12: 12, 13: 13, 14: 14, 15: 15, 16: 16, 17: 17, 18: 18, 19: 19,\n               20: 20, 21: 21, 22: 22, 23: 23, 24: 24, 25: 25, 26: 26, \n               27: 27, 28: 28, 29: 29, 30: 30, 31: 31, 32: 32, 33: 33, 34: 34, \n               35: 35, 36: 36, 37: 37, 38: 38},\n \n    'time_id' : {i : i for i in range(968)}}\n\ndef encode_column(df, column, mapping):\n\n    max_value = max(mapping.values())  \n    \n    def encode_category(category):\n        # Return max_value + 1 for any unseen categories\n        return mapping.get(category, max_value + 1)  \n    \n    return df.with_columns(\n        pl.col(column).map_elements(encode_category).alias(column)\n    )\n\nclass R2Loss(nn.Module):\n\n    def __init__(self):\n        super(R2Loss, self).__init__()\n\n    def forward(self, y_pred, y_true):\n        # Calculate MSE\n        mse_loss = torch.sum((y_pred - y_true) ** 2)\n        # Calculate variance of true values\n        var_y = torch.sum(y_true ** 2)\n        # Calculate R² loss (adding small epsilon to avoid division by zero)\n        loss = mse_loss / (var_y + 1e-38)\n        return loss\n\nclass NN(LightningModule):\n\n    def __init__(self, n_cont_features, cat_cardinalities, n_classes, lr, weight_decay):\n        super().__init__()\n        self.save_hyperparameters()\n        self.k = 16  # Number of ensemble members\n\n        # Initialize the main model architecture\n        self.model = Model(\n                n_num_features=n_cont_features,\n                cat_cardinalities=cat_cardinalities,\n                n_classes=n_classes,\n                backbone={\n                    'type': 'MLP',\n                    'n_blocks': 3,\n                    'd_block': [512]*3,\n                    'dropout': 0.25,\n                },\n                bins=None,\n                num_embeddings=None,\n                arch_type='tabm',\n                k=self.k,\n            )\n        \n        # Set learning parameters\n        self.lr = lr\n        self.weight_decay = weight_decay\n        \n        # Initialize lists to store outputs during training and validation\n        self.training_step_outputs = []\n        self.validation_step_outputs = []\n        \n        # Define loss function\n        self.loss_fn = R2Loss()\n\n    def forward(self, x_cont, x_cat):\n\n        return self.model(x_cont, x_cat).squeeze(-1)\n\n    def training_step(self, batch):\n        \n        x_cont, x_cat, y, w, w_y = batch\n        \n        # Add random noise to continuous features for regularization\n        x_cont = x_cont + torch.randn_like(x_cont) * 0.02\n        \n        # Get model predictions\n        y_hat = self(x_cont, x_cat)\n        \n        # Calculate loss\n        loss = self.loss_fn(y_hat.flatten(0, 1), y.repeat_interleave(self.k))\n        \n        # Log training loss\n        self.log('train_loss', loss, on_step=True, on_epoch=True, \n                prog_bar=True, logger=True, batch_size=x_cont.size(0))\n        \n        # Store outputs for epoch-end calculations\n        self.training_step_outputs.append((y_hat.mean(1), y, w))\n        \n        return loss\n\n    def validation_step(self, batch):\n\n        x_cont, x_cat, y, w, w_y = batch\n        x_cont = x_cont + torch.randn_like(x_cont) * 0.02\n        y_hat = self(x_cont, x_cat)\n        \n        loss = self.loss_fn(y_hat.flatten(0, 1), y.repeat_interleave(self.k))\n        \n        self.log('val_loss', loss, on_step=False, on_epoch=True, \n                prog_bar=True, logger=True, batch_size=x_cont.size(0))\n        \n        self.validation_step_outputs.append((y_hat.mean(1), y, w))\n        return loss\n\n    def on_validation_epoch_end(self):\n        \n        y = torch.cat([x[1] for x in self.validation_step_outputs]).cpu().numpy()\n        \n        if self.trainer.sanity_checking:\n            prob = torch.cat(\n                [x[0] for x in self.validation_step_outputs]).cpu().numpy()\n        else:\n            prob = torch.cat(\n                [x[0] for x in self.validation_step_outputs]).cpu().numpy()\n            weights = torch.cat(\n                [x[2] for x in self.validation_step_outputs]).cpu().numpy()\n            \n            # Calculate R² score for validation\n            val_r_square = r2_val(y, prob, weights)\n            self.log(\"val_r_square\", val_r_square, \n                     prog_bar=True, on_step=False, on_epoch=True)\n        \n        self.validation_step_outputs.clear()\n\n    def configure_optimizers(self):\n        \n        optimizer = torch.optim.AdamW(\n            make_parameter_groups(self.model), \n            lr=self.lr, \n            weight_decay=self.weight_decay\n        )\n        return {'optimizer': optimizer}\n\n    def on_train_epoch_end(self):\n        \n        if self.trainer.sanity_checking:\n            return\n            \n        # Gather all outputs from training steps\n        y = torch.cat([x[1] for x in self.training_step_outputs]).cpu().numpy()\n        prob = torch.cat([x[0] for x in self.training_step_outputs]).detach().cpu().numpy()\n        weights = torch.cat([x[2] for x in self.training_step_outputs]).cpu().numpy()\n        \n        # Calculate R² score for training\n        train_r_square = r2_val(y, prob, weights)\n        self.log(\"train_r_square\", train_r_square, \n                 prog_bar=True, on_step=False, on_epoch=True)\n        \n        self.training_step_outputs.clear()\n        \n        # Print epoch metrics\n        epoch = self.trainer.current_epoch\n        metrics = {k: v.item() if isinstance(v, torch.Tensor) else v \n                  for k, v in self.trainer.logged_metrics.items()}\n        formatted_metrics = {k: f\"{v:.5f}\" for k, v in metrics.items()}\n        print(f\"Epoch {epoch}: {formatted_metrics}\")\n        ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T23:25:40.496645Z","iopub.execute_input":"2025-05-28T23:25:40.496910Z","iopub.status.idle":"2025-05-28T23:25:40.515424Z","shell.execute_reply.started":"2025-05-28T23:25:40.496894Z","shell.execute_reply":"2025-05-28T23:25:40.514516Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class custom_args:\n\n    def __init__(self):\n        # GPU Configuration\n        self.usegpu = True\n        self.gpuid = 0\n        \n        # Random seed for reproducibility\n        self.seed = 42\n        \n        # Model Configuration\n        self.model = 'nn'  # Neural network model type\n        \n        # Wandb logging configuration\n        self.use_wandb = False\n        self.project = 'js-tabm-with-lags'\n        \n        # Data and loading configuration\n        self.dname = \"./input_df/\"  # Data directory\n        self.loader_workers = 10    # Number of workers for data loading\n        self.bs = 8192             # Batch size\n        \n        # Model hyperparameters\n        self.lr = 1e-3             # Learning rate\n        self.weight_decay = 8e-4    # Weight decay for regularization\n        \n        # Feature configuration\n        self.n_cont_features = 84   # Number of continuous features\n        self.n_cat_features = 5     # Number of categorical features\n        self.n_classes = None       # Number of classes (None for regression)\n        \n        # Categorical feature cardinalities\n        # [feature_09, feature_10, feature_11, symbol_id, time_id]\n        self.cat_cardinalities = [23, 10, 32, 40, 969]\n        \n        # Training configuration\n        self.patience = 7           # Early stopping patience\n        self.max_epochs = 10        # Maximum training epochs\n        self.N_fold = 5            # Number of cross-validation folds\n\n# Create instance of custom arguments\nCFG = custom_args()\n\n# Set up device (GPU if available, else CPU)\ndevice = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n\n# Load pre-trained model from checkpoint and move to appropriate device\ntabm_model = NN.load_from_checkpoint('/kaggle/input/trained/tabm.ckpt').to(device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T23:43:34.828732Z","iopub.execute_input":"2025-05-28T23:43:34.829023Z","iopub.status.idle":"2025-05-28T23:43:35.200477Z","shell.execute_reply.started":"2025-05-28T23:43:34.829006Z","shell.execute_reply":"2025-05-28T23:43:35.199320Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Here we write a predict function for tabm models to accomadate the JS submission API**","metadata":{}},{"cell_type":"code","source":"# Global variables to store lag features\nlags_ : pl.DataFrame | None = None\n\nlags_history = None\n\ndef predict_tabm(test: pl.DataFrame, lags: pl.DataFrame | None) -> pl.DataFrame | pd.DataFrame:\n    \"\"\"\n    Make predictions using the TABM (Tabular Model)\n    \n    Args:\n        test: Input DataFrame containing test features\n        lags: DataFrame containing lagged features\n        \n    Returns:\n        DataFrame with predictions\n    \"\"\"\n    global lags_, lags_history\n    # Update global lags if new ones provided\n    if lags is not None:\n        lags_ = lags\n\n    # Encode categorical features\n    for col in feature_cat + ['symbol_id', 'time_id']:\n        test = encode_column(test, col, category_mappings[col])\n\n    # Initialize predictions DataFrame with row_ids\n    predictions = test.select(\n        'row_id',\n        pl.lit(0.0).alias('responder_6'),\n    )\n\n    # Extract symbol and time information\n    symbol_ids = test.select('symbol_id').to_numpy()[:, 0]\n    time_id = test.select(\"time_id\").to_numpy()[0]\n    timie_id_array = test.select(\"time_id\").to_numpy()[:, 0]\n    \n    # Handle time_id = 0 case (first prediction)\n    if time_id == 0:\n        # Convert time_id and symbol_id to integers\n        lags = lags.with_columns(pl.col('time_id').cast(pl.Int64))\n        lags = lags.with_columns(pl.col('symbol_id').cast(pl.Int64)) \n        # Store full lags history and filter for time_id 0\n        lags_history = lags\n        lags = lags.filter(pl.col(\"time_id\") == 0)  \n        test = test.join(lags, on=[\"time_id\", \"symbol_id\"],  how=\"left\")\n    else:\n        # Filter lags for current time_id\n        lags = lags_history.filter(pl.col(\"time_id\") == time_id)\n        test = test.join(lags, on=[\"time_id\", \"symbol_id\"],  how=\"left\")\n\n    # Fill missing values with 0\n    test = test.with_columns([\n        pl.col(col).fill_null(0) for col in feature_list + [\n            f\"responder_{idx}_lag_1\" for idx in range(9)] \n    ])\n    \n    # Standardize features\n    test = standardize(test, std_feature, means, stds)\n\n    # Convert to numpy array and then to torch tensors\n    X_test = test[feature_test].to_numpy()\n    X_test_tensor = torch.tensor(X_test, dtype=torch.float32).to(device)\n    symbol_tensor = torch.tensor(symbol_ids, dtype=torch.float32).to(device)\n    time_tensor = torch.tensor(timie_id_array, dtype=torch.float32).to(device)\n    \n    # Separate categorical and continuous features\n    X_cat = X_test_tensor[:, [9, 10, 11]]\n    X_cont = X_test_tensor[:, \n    [i for i in range(X_test_tensor.shape[1]) if i not in [9, 10, 11]]]\n\n    # Combine categorical features with symbol and time information\n    X_cat = (torch.concat([\n        X_cat, symbol_tensor.unsqueeze(-1), \n        time_tensor.unsqueeze(-1)\n    ], axis=1)).to(torch.int64)\n\n    # Make predictions\n    model.eval()\n    with torch.no_grad():\n        \n        outputs = model(X_cont, X_cat)\n        # Assuming the model outputs a tensor of shape (batch_size, 1)\n        preds = outputs.squeeze(-1).cpu().numpy()\n        preds = preds.mean(1)\n\n    # Create final predictions DataFrame\n    predictions = \\\n    test.select('row_id').\\\n    with_columns(\n        pl.Series(\n            name   = 'responder_6', \n            values = np.clip(preds, a_min = -5, a_max = 5),\n            dtype  = pl.Float64,\n        )\n    )\n\n    return predictions","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T23:44:33.219188Z","iopub.execute_input":"2025-05-28T23:44:33.219518Z","iopub.status.idle":"2025-05-28T23:44:33.230546Z","shell.execute_reply.started":"2025-05-28T23:44:33.219500Z","shell.execute_reply":"2025-05-28T23:44:33.229788Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 4. Blending","metadata":{}},{"cell_type":"code","source":"lags_ : pl.DataFrame | None = None\n    \ndef predict(test: pl.DataFrame, lags: pl.DataFrame | None) -> pl.DataFrame | pd.DataFrame:\n    global lags_\n    if lags is not None:\n        lags_ = lags\n\n    predictions = test.select(\n        'row_id',\n        pl.lit(0.0).alias('responder_6'),\n    )\n    symbol_ids = test.select('symbol_id').to_numpy()[:, 0]\n\n    lags = lags_.clone().group_by(\n        [\"date_id\", \"symbol_id\"], maintain_order=True\n    ).last() # pick up last record of previous date\n    test = test.join(lags, on=[\"date_id\", \"symbol_id\"],  how=\"left\")\n\n    # ------------------------------------------------- #\n    # Inf\n    # ------------------------------------------------- #\n    preds_xgb = np.zeros((test.shape[0],))\n    preds_nn = np.zeros((test.shape[0],))\n    preds_tabm = np.zeros((test.shape[0],))\n\n    \"\"\" Pred Ridge \"\"\"\n    preds_ridge = predict_ridge(test,lags)\n    \n    \"\"\" Pred XGB \"\"\"\n    for i in range(5):\n        preds_xgb += xgb_models[i].predict(test[xgb_feature_cols].to_pandas())/5\n    \n    \n    \"\"\" Pred NN \"\"\"\n    test_input = test[CONFIG.feature_cols].to_pandas()\n    test_input = test_input.fillna(method = 'ffill').fillna(0)\n    test_input = torch.FloatTensor(test_input.values)\n    with torch.no_grad():\n        for i, nn_model in enumerate(tqdm(nn_models)):\n            nn_model.eval()\n            preds_nn += nn_model(test_input).cpu().numpy()/len(models)\n\n    \"\"\" Pred TabM\"\"\"\n\n    preds_tabm = predict_tabm(test, lags)\n\n    \"\"\" Model Weight \"\"\"\n    _ModelW = [0.55, 0.45] * 0.7 + [0.5, 0.5] * 0.3\n    preds = (preds_xgb*_ModelW[0] + \\\n             preds_nn*_ModelW[1] + \\\n             preds_ridge*_ModelW[2] + \\\n             preds_rigde*_ModelW[3])\n\n\n    \"\"\" Finaly \"\"\"\n    predictions = \\\n    test.select('row_id').\\\n    with_columns(\n        pl.Series(\n            name   = 'responder_6', \n            values = np.clip(preds, a_min = -5, a_max = 5),\n            dtype  = pl.Float64,\n        )\n    )\n\n    assert isinstance(predictions, pl.DataFrame | pd.DataFrame)\n    assert list(predictions.columns) == ['row_id', 'responder_6']\n    assert len(predictions) == len(test)\n\n    return predictions","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T00:00:35.075088Z","iopub.execute_input":"2025-05-29T00:00:35.075344Z","iopub.status.idle":"2025-05-29T00:00:35.083810Z","shell.execute_reply.started":"2025-05-29T00:00:35.075328Z","shell.execute_reply":"2025-05-29T00:00:35.083003Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"SUBMIT = False\nif SUBMIT:\n    import kaggle_evaluation.jane_street_inference_server\n    #########################\n    inference_server = kaggle_evaluation.jane_street_inference_server.JSInferenceServer(predict)\n\n    if os.getenv('KAGGLE_IS_COMPETITION_RERUN'):\n        inference_server.serve()\n    else:\n        inference_server.run_local_gateway(\n            (\n                '/kaggle/input/jane-street-realtime-marketdata-forecasting/test.parquet',\n                '/kaggle/input/jane-street-realtime-marketdata-forecasting/lags.parquet',\n            )\n        )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T00:02:09.710729Z","iopub.execute_input":"2025-05-29T00:02:09.711086Z","iopub.status.idle":"2025-05-29T00:02:09.716588Z","shell.execute_reply.started":"2025-05-29T00:02:09.711065Z","shell.execute_reply":"2025-05-29T00:02:09.715763Z"}},"outputs":[],"execution_count":null}]}