{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":84493,"databundleVersionId":9871156,"sourceType":"competition"},{"sourceId":217842901,"sourceType":"kernelVersion"}],"dockerImageVersionId":30787,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":80.344101,"end_time":"2024-10-26T03:27:42.952247","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2024-10-26T03:26:22.608146","version":"2.6.0"},"widgets":{"application/vnd.jupyter.widget-state+json":{"state":{"252dd2de87de42f4becd877fbbafc26b":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"HTMLModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_ff55584ae0ab4f39ae471f314eaad988","placeholder":"​","style":"IPY_MODEL_b8d338c473aa4dc3ba25f37a997a9037","value":" 1/1 [00:00&lt;00:00, 33.79it/s]"}},"48a2731fb59b4ce8ace5b53d6f0e3337":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"HTMLModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_a8dbc79c7c5a48318466a102c55801bb","placeholder":"​","style":"IPY_MODEL_ebd06aaaa7024a1684ba1b9fe89358cf","value":"100%"}},"56da652f3eeb42aca986e6fd815629dc":{"model_module":"@jupyter-widgets/base","model_module_version":"1.2.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"5bb4e0df01c44716afebab25aafe9f5f":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"FloatProgressModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"FloatProgressModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"ProgressView","bar_style":"success","description":"","description_tooltip":null,"layout":"IPY_MODEL_f4e41234b0cc4f3e9313d971f5aadc9f","max":1,"min":0,"orientation":"horizontal","style":"IPY_MODEL_6408698650f74dd699bcc914b850e396","value":1}},"6408698650f74dd699bcc914b850e396":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"ProgressStyleModel","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"ProgressStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","bar_color":null,"description_width":""}},"9cb493ec04fc4ed391f5ac28cc84500e":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"HBoxModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HBoxModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HBoxView","box_style":"","children":["IPY_MODEL_48a2731fb59b4ce8ace5b53d6f0e3337","IPY_MODEL_5bb4e0df01c44716afebab25aafe9f5f","IPY_MODEL_252dd2de87de42f4becd877fbbafc26b"],"layout":"IPY_MODEL_56da652f3eeb42aca986e6fd815629dc"}},"a8dbc79c7c5a48318466a102c55801bb":{"model_module":"@jupyter-widgets/base","model_module_version":"1.2.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"b8d338c473aa4dc3ba25f37a997a9037":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"DescriptionStyleModel","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"ebd06aaaa7024a1684ba1b9fe89358cf":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"DescriptionStyleModel","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"f4e41234b0cc4f3e9313d971f5aadc9f":{"model_module":"@jupyter-widgets/base","model_module_version":"1.2.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"ff55584ae0ab4f39ae471f314eaad988":{"model_module":"@jupyter-widgets/base","model_module_version":"1.2.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}}},"version_major":2,"version_minor":0}}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# If Internet access is not allowed\n!pip install rtdl_num_embeddings -q --no-index --find-links=/kaggle/input/import-tabm/rtdl_num_embeddings","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:18.323780Z","iopub.execute_input":"2025-01-16T02:50:18.324243Z","iopub.status.idle":"2025-01-16T02:50:26.579926Z","shell.execute_reply.started":"2025-01-16T02:50:18.324209Z","shell.execute_reply":"2025-01-16T02:50:26.578899Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# # If Internet access is allowed\n# !git clone https://github.com/yandex-research/tabm\n# !pip install rtdl_num_embeddings","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:26.582056Z","iopub.execute_input":"2025-01-16T02:50:26.582365Z","iopub.status.idle":"2025-01-16T02:50:26.587271Z","shell.execute_reply.started":"2025-01-16T02:50:26.582336Z","shell.execute_reply":"2025-01-16T02:50:26.585629Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!cp /kaggle/input/import-tabm/tabm/tabm_reference.py .","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:26.588575Z","iopub.execute_input":"2025-01-16T02:50:26.589265Z","iopub.status.idle":"2025-01-16T02:50:27.777298Z","shell.execute_reply.started":"2025-01-16T02:50:26.589236Z","shell.execute_reply":"2025-01-16T02:50:27.775928Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from tabm_reference import Model, make_parameter_groups","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:27.779958Z","iopub.execute_input":"2025-01-16T02:50:27.780296Z","iopub.status.idle":"2025-01-16T02:50:27.785261Z","shell.execute_reply.started":"2025-01-16T02:50:27.780264Z","shell.execute_reply":"2025-01-16T02:50:27.784270Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport polars as pl\nimport numpy as np\nimport pandas as pd\nimport gc\n\nimport warnings\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport pytorch_lightning as plt\nfrom pytorch_lightning import (LightningDataModule, LightningModule, Trainer)\nfrom pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint, Timer, LearningRateMonitor\nfrom pytorch_lightning.loggers import WandbLogger\nimport wandb\nfrom sklearn.metrics import r2_score\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.preprocessing import RobustScaler, normalize\nfrom sklearn.preprocessing import OneHotEncoder\nfrom torch.utils.data import Dataset, DataLoader\n\nimport joblib","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:27.786321Z","iopub.execute_input":"2025-01-16T02:50:27.786601Z","iopub.status.idle":"2025-01-16T02:50:27.803157Z","shell.execute_reply.started":"2025-01-16T02:50:27.786576Z","shell.execute_reply":"2025-01-16T02:50:27.802241Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"plt.seed_everything(42)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:27.804277Z","iopub.execute_input":"2025-01-16T02:50:27.804838Z","iopub.status.idle":"2025-01-16T02:50:27.822328Z","shell.execute_reply.started":"2025-01-16T02:50:27.804810Z","shell.execute_reply":"2025-01-16T02:50:27.821431Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"input_path = './jane-street-real-time-market-data-forecasting/' if os.path.exists('./jane-street-real-time-market-data-forecasting') else '/kaggle/input/jane-street-real-time-market-data-forecasting/'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:27.823339Z","iopub.execute_input":"2025-01-16T02:50:27.823609Z","iopub.status.idle":"2025-01-16T02:50:27.830937Z","shell.execute_reply.started":"2025-01-16T02:50:27.823585Z","shell.execute_reply":"2025-01-16T02:50:27.830027Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"col_target = \"responder_6\"\ncol_weight = \"weight\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:27.831967Z","iopub.execute_input":"2025-01-16T02:50:27.832226Z","iopub.status.idle":"2025-01-16T02:50:27.840456Z","shell.execute_reply.started":"2025-01-16T02:50:27.832190Z","shell.execute_reply":"2025-01-16T02:50:27.839632Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"cat_map = {'feature_09': {2: 0, 4: 1, 9: 2, 11: 3, 12: 4, 14: 5, 15: 6, 25: 7, 26: 8, 30: 9, 34: 10, 42: 11, 44: 12, 46: 13, 49: 14, 50: 15, 57: 16, 64: 17, 68: 18, 70: 19, 81: 20, 82: 21},\n 'feature_10': {1: 0, 2: 1, 3: 2, 4: 3, 5: 4, 6: 5, 7: 6, 10: 7, 12: 8},\n 'feature_11': {9: 0, 11: 1, 13: 2, 16: 3, 24: 4, 25: 5, 34: 6, 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, 214: 20, 230: 21, 261: 22, 297: 23, 336: 24, 376: 25, 388: 26, 410: 27, 522: 28, 534: 29, 539: 30},\n 'symbol_id': {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7, 8: 8, 9: 9, 10: 10, 11: 11, 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, 27: 27, 28: 28, 29: 29, 30: 30, 31: 31, 32: 32, 33: 33, 34: 34, 35: 35, 36: 36, 37: 37, 38: 38},\n 'time_id' : {i : i for i in range(968)}}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:27.841606Z","iopub.execute_input":"2025-01-16T02:50:27.841866Z","iopub.status.idle":"2025-01-16T02:50:27.852976Z","shell.execute_reply.started":"2025-01-16T02:50:27.841841Z","shell.execute_reply":"2025-01-16T02:50:27.852251Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"col_feature_list = [f\"feature_{idx:02d}\" for idx in range(79) if idx != 61]\ncol_feature_cat = [\"feature_09\", \"feature_10\", \"feature_11\", \"symbol_id\", \"time_id\"]\ncol_feature_cont = [item for item in col_feature_list if item not in col_feature_cat]\ncol_feature = col_feature_cat + col_feature_cont","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:27.855936Z","iopub.execute_input":"2025-01-16T02:50:27.856202Z","iopub.status.idle":"2025-01-16T02:50:27.864338Z","shell.execute_reply.started":"2025-01-16T02:50:27.856176Z","shell.execute_reply":"2025-01-16T02:50:27.863463Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_train = []\nfor i in [9]:\n    df = pl.read_parquet(f'{input_path}/train.parquet/partition_id={i}/part-0.parquet')\n    df_train.append(df)\n    del df\n    gc.collect()\ndf_train = pl.concat(df_train)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:27.865328Z","iopub.execute_input":"2025-01-16T02:50:27.865607Z","iopub.status.idle":"2025-01-16T02:50:30.334638Z","shell.execute_reply.started":"2025-01-16T02:50:27.865582Z","shell.execute_reply":"2025-01-16T02:50:30.333650Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# df_train = df_train.filter(pl.col(\"date_id\") > 1600)\ndf_train = df_train.fill_null(0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:30.335950Z","iopub.execute_input":"2025-01-16T02:50:30.336227Z","iopub.status.idle":"2025-01-16T02:50:30.540356Z","shell.execute_reply.started":"2025-01-16T02:50:30.336200Z","shell.execute_reply":"2025-01-16T02:50:30.539286Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# checking device\ndevice = torch.device(f'cuda:0' if torch.cuda.is_available() else 'cpu')\naccelerator = 'gpu' if torch.cuda.is_available() else 'cpu'\nloader_device = 'cpu'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:30.541644Z","iopub.execute_input":"2025-01-16T02:50:30.542000Z","iopub.status.idle":"2025-01-16T02:50:30.546642Z","shell.execute_reply.started":"2025-01-16T02:50:30.541962Z","shell.execute_reply":"2025-01-16T02:50:30.545771Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def encode_column(df, column, mapping):\n    max_value = max(mapping.values())  \n\n    def encode_category(category):\n        return mapping.get(category, max_value + 1)  \n    \n    return df.with_columns(\n        pl.col(column).map_elements(encode_category, return_dtype = int).alias(column)\n    )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:30.547874Z","iopub.execute_input":"2025-01-16T02:50:30.548160Z","iopub.status.idle":"2025-01-16T02:50:30.559342Z","shell.execute_reply.started":"2025-01-16T02:50:30.548132Z","shell.execute_reply":"2025-01-16T02:50:30.558535Z"},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class custom_args():\n    def __init__(self):\n        self.usegpu = True\n        self.gpuid = 0\n        self.seed = 42\n        self.loader_workers = 10\n        self.k = 16\n        self.n_blocks = 3\n        self.d_block = 512\n        self.dropout = 0.25\n        self.bs = 2048\n        self.lr = 1e-3\n        self.weight_decay = 8e-4\n        self.n_cont_features = len(col_feature_cont)\n        self.n_cat_features = len(col_feature_cat)\n        self.n_classes = None\n        self.cat_cardinalities = [23, 10, 32, 40, 969]\n        self.patience = 15\n        self.max_epochs = 2\n\n\nargs = custom_args()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:30.560394Z","iopub.execute_input":"2025-01-16T02:50:30.560699Z","iopub.status.idle":"2025-01-16T02:50:30.570321Z","shell.execute_reply.started":"2025-01-16T02:50:30.560673Z","shell.execute_reply":"2025-01-16T02:50:30.569545Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CustomDataset(Dataset):\n    def __init__(self, df, accelerator):\n        self.cont = torch.FloatTensor(df[col_feature_cont].to_numpy()).to(accelerator)\n        self.cat = torch.LongTensor(df[col_feature_cat].to_numpy()).to(accelerator)\n        self.labels = torch.FloatTensor(df[col_target].to_numpy()).to(accelerator)\n        self.weights = torch.FloatTensor(df[col_weight].to_numpy()).to(accelerator)\n\n    def __len__(self):\n        return len(self.labels)\n\n    def __getitem__(self, idx):\n        x_cont = self.cont[idx]\n        x_cat = self.cat[idx]\n        y = self.labels[idx]\n        w = self.weights[idx]\n        return x_cont, x_cat, y, w, y*w\n\n\nclass DataModule(LightningDataModule):\n    def __init__(self, train_df, batch_size, valid_df=None, accelerator='cpu'):\n        super().__init__()\n        self.df = train_df\n        self.batch_size = batch_size\n        self.accelerator = accelerator\n        self.train_df = train_df\n        self.train_dataset = None\n        self.valid_df = None\n        if valid_df is not None:\n            self.valid_df = valid_df\n        self.val_dataset = None\n\n    def setup(self):\n        self.train_dataset = CustomDataset(self.train_df, self.accelerator)\n        if self.valid_df is not None:\n            df_valid = self.valid_df\n            self.val_dataset = CustomDataset(df_valid, self.accelerator)\n\n    def train_dataloader(self, n_workers=0):\n        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=n_workers)\n\n    def val_dataloader(self, n_workers=0):\n        return DataLoader(self.val_dataset, batch_size=self.batch_size, shuffle=False, num_workers=n_workers)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:30.571394Z","iopub.execute_input":"2025-01-16T02:50:30.571681Z","iopub.status.idle":"2025-01-16T02:50:30.582939Z","shell.execute_reply.started":"2025-01-16T02:50:30.571656Z","shell.execute_reply":"2025-01-16T02:50:30.582159Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Custom R2 metric for validation\ndef r2_val(y_true, y_pred, sample_weight):\n    r2 = 1 - np.average((y_pred - y_true) ** 2, weights=sample_weight) / (np.average((y_true) ** 2, weights=sample_weight) + 1e-38)\n    return r2","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:30.583984Z","iopub.execute_input":"2025-01-16T02:50:30.584262Z","iopub.status.idle":"2025-01-16T02:50:30.601269Z","shell.execute_reply.started":"2025-01-16T02:50:30.584238Z","shell.execute_reply":"2025-01-16T02:50:30.600396Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class R2Loss(nn.Module):\n    def __init__(self):\n        super(R2Loss, self).__init__()\n\n    def forward(self, y_pred, y_true):\n        mse_loss = torch.sum((y_pred - y_true) ** 2)\n        var_y = torch.sum(y_true ** 2)\n        loss = mse_loss / (var_y + 1e-38)\n\n        return loss\n\n\nclass NN(LightningModule):\n    def __init__(self, n_cont_features, cat_cardinalities, n_classes, k, n_blocks, d_block, dropout, lr, weight_decay):\n        super().__init__()\n        self.save_hyperparameters()\n        self.k = k\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': n_blocks ,\n                    'd_block': d_block,\n                    'dropout': dropout,\n                },\n                bins=None,\n                num_embeddings= None,\n                arch_type='tabm',\n                k=self.k,\n            )\n        self.lr = lr\n        self.weight_decay = weight_decay\n        self.training_step_outputs = []\n        self.validation_step_outputs = []\n        self.loss_fn = R2Loss()\n\n    def forward(self, x_cont, x_cat):\n        return self.model(x_cont, x_cat).squeeze(-1)\n\n    def training_step(self, batch):\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        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True, batch_size=x_cont.size(0))\n        self.training_step_outputs.append((y_hat.mean(1), y, w))\n        return loss\n\n    def validation_step(self, batch):\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        loss = self.loss_fn(y_hat.flatten(0, 1), y.repeat_interleave(self.k))\n        self.log('val_loss', loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, batch_size=x_cont.size(0))\n        self.validation_step_outputs.append((y_hat.mean(1), y, w))\n        return loss\n\n    def on_validation_epoch_end(self):\n        \"\"\"Calculate validation WRMSE at the end of the epoch.\"\"\"\n        y = torch.cat([x[1] for x in self.validation_step_outputs]).cpu().numpy()\n        if self.trainer.sanity_checking:\n            prob = torch.cat([x[0] for x in self.validation_step_outputs]).cpu().numpy()\n        else:\n            prob = torch.cat([x[0] for x in self.validation_step_outputs]).cpu().numpy()\n            weights = torch.cat([x[2] for x in self.validation_step_outputs]).cpu().numpy()\n            val_r_square = r2_val(y, prob, weights)\n            self.log(\"val_r_square\", val_r_square, prog_bar=True, on_step=False, on_epoch=True)\n        self.validation_step_outputs.clear()\n\n    def configure_optimizers(self):\n        optimizer = torch.optim.AdamW(make_parameter_groups(self.model), lr=self.lr, weight_decay=self.weight_decay)\n        # scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=5,\n        #                                                        verbose=True)\n        return {\n            'optimizer': optimizer,\n            # 'lr_scheduler': {\n            #     'scheduler': scheduler,\n            #     'monitor': 'val_r_square',\n            # }\n        }\n\n    def on_train_epoch_end(self):\n        if self.trainer.sanity_checking:\n            return\n\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        # r2_training\n        train_r_square = r2_val(y, prob, weights)\n        self.log(\"train_r_square\", train_r_square, prog_bar=True, on_step=False, on_epoch=True)\n        self.training_step_outputs.clear()\n\n        epoch = self.trainer.current_epoch\n        metrics = {k: v.item() if isinstance(v, torch.Tensor) else v 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}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:30.602734Z","iopub.execute_input":"2025-01-16T02:50:30.603216Z","iopub.status.idle":"2025-01-16T02:50:30.622567Z","shell.execute_reply.started":"2025-01-16T02:50:30.603167Z","shell.execute_reply":"2025-01-16T02:50:30.621825Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_train","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:30.623716Z","iopub.execute_input":"2025-01-16T02:50:30.623959Z","iopub.status.idle":"2025-01-16T02:50:30.647556Z","shell.execute_reply.started":"2025-01-16T02:50:30.623935Z","shell.execute_reply":"2025-01-16T02:50:30.646747Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for col in col_feature_cat:\n    df_train = encode_column(df_train, col, cat_map[col])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:30.648539Z","iopub.execute_input":"2025-01-16T02:50:30.648785Z","iopub.status.idle":"2025-01-16T02:50:35.693335Z","shell.execute_reply.started":"2025-01-16T02:50:30.648760Z","shell.execute_reply":"2025-01-16T02:50:35.692640Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_tr = df_train.filter(pl.col(\"date_id\") <= 1600)\ndf_va = df_train.filter(pl.col(\"date_id\") > 1600)\ndel df_train\ngc.collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:35.694611Z","iopub.execute_input":"2025-01-16T02:50:35.695233Z","iopub.status.idle":"2025-01-16T02:50:36.063068Z","shell.execute_reply.started":"2025-01-16T02:50:35.695191Z","shell.execute_reply":"2025-01-16T02:50:36.062164Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"RS = RobustScaler()\ndf_tr[col_feature_cont] = RS.fit_transform(df_tr[col_feature_cont])\ndf_va[col_feature_cont] = RS.transform(df_va[col_feature_cont])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:50:36.064289Z","iopub.execute_input":"2025-01-16T02:50:36.065167Z","iopub.status.idle":"2025-01-16T02:51:08.042564Z","shell.execute_reply.started":"2025-01-16T02:50:36.065136Z","shell.execute_reply":"2025-01-16T02:51:08.041553Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = NN(\n    n_cont_features = args.n_cont_features,\n    cat_cardinalities = args.cat_cardinalities,\n    n_classes = args.n_classes,\n    k = args.k,\n    n_blocks = args.n_blocks,\n    d_block = args.d_block,\n    dropout = args.dropout,\n    lr=args.lr,\n    weight_decay=args.weight_decay\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:51:08.043862Z","iopub.execute_input":"2025-01-16T02:51:08.044155Z","iopub.status.idle":"2025-01-16T02:51:08.073247Z","shell.execute_reply.started":"2025-01-16T02:51:08.044127Z","shell.execute_reply":"2025-01-16T02:51:08.072294Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Initialize Logger\nlogger = None\n# Initialize Callbacks\n# early_stopping = EarlyStopping('val_loss', patience=args.patience, mode='min', verbose=False)\ncheckpoint_callback = ModelCheckpoint(monitor='val_loss', mode='min', save_top_k=1, verbose=False, filename=None)\ntimer = Timer()\n# Initialize Trainer\ntrainer = Trainer(\n    max_epochs=args.max_epochs,\n    accelerator=accelerator,\n    devices=[args.gpuid] if args.usegpu else None,\n    logger=logger,\n    callbacks=[timer],\n    enable_progress_bar=True\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:51:08.074207Z","iopub.execute_input":"2025-01-16T02:51:08.074469Z","iopub.status.idle":"2025-01-16T02:51:08.116965Z","shell.execute_reply.started":"2025-01-16T02:51:08.074436Z","shell.execute_reply":"2025-01-16T02:51:08.116129Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"data_module = DataModule(df_tr, batch_size=args.bs, valid_df=df_va, accelerator=loader_device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:51:08.118011Z","iopub.execute_input":"2025-01-16T02:51:08.118265Z","iopub.status.idle":"2025-01-16T02:51:08.122342Z","shell.execute_reply.started":"2025-01-16T02:51:08.118239Z","shell.execute_reply":"2025-01-16T02:51:08.121506Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"plt.seed_everything(args.seed)\ndata_module.setup()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:51:08.123634Z","iopub.execute_input":"2025-01-16T02:51:08.123992Z","iopub.status.idle":"2025-01-16T02:51:10.699791Z","shell.execute_reply.started":"2025-01-16T02:51:08.123952Z","shell.execute_reply":"2025-01-16T02:51:10.698776Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Start Training\ntrainer.fit(model, data_module.train_dataloader(args.loader_workers), data_module.val_dataloader(args.loader_workers))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:51:10.701007Z","iopub.execute_input":"2025-01-16T02:51:10.701300Z","iopub.status.idle":"2025-01-16T02:56:07.484641Z","shell.execute_reply.started":"2025-01-16T02:51:10.701273Z","shell.execute_reply":"2025-01-16T02:56:07.483715Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"del df_tr\ngc.collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:56:07.490481Z","iopub.execute_input":"2025-01-16T02:56:07.490812Z","iopub.status.idle":"2025-01-16T02:56:08.107990Z","shell.execute_reply.started":"2025-01-16T02:56:07.490779Z","shell.execute_reply":"2025-01-16T02:56:08.107065Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_va = df_va.sample(100000)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:56:08.109156Z","iopub.execute_input":"2025-01-16T02:56:08.109551Z","iopub.status.idle":"2025-01-16T02:56:08.123469Z","shell.execute_reply.started":"2025-01-16T02:56:08.109510Z","shell.execute_reply":"2025-01-16T02:56:08.122474Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model.eval()\nwith torch.no_grad():\n    preds = model(\n        torch.FloatTensor(df_va[col_feature_cont].to_numpy()),\n        torch.LongTensor(df_va[col_feature_cat].to_numpy())\n    ).cpu().numpy().mean(1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:56:08.124565Z","iopub.execute_input":"2025-01-16T02:56:08.124832Z","iopub.status.idle":"2025-01-16T02:56:12.238145Z","shell.execute_reply.started":"2025-01-16T02:56:08.124806Z","shell.execute_reply":"2025-01-16T02:56:12.237120Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"r2_val(df_va[col_target].to_numpy(), preds, df_va[col_weight].to_numpy())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-16T02:56:12.239498Z","iopub.execute_input":"2025-01-16T02:56:12.239871Z","iopub.status.idle":"2025-01-16T02:56:12.246859Z","shell.execute_reply.started":"2025-01-16T02:56:12.239841Z","shell.execute_reply":"2025-01-16T02:56:12.245849Z"}},"outputs":[],"execution_count":null}]}