{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":84493,"databundleVersionId":9871156,"sourceType":"competition"},{"sourceId":10422214,"sourceType":"datasetVersion","datasetId":6344006},{"sourceId":214892435,"sourceType":"kernelVersion"},{"sourceId":215616115,"sourceType":"kernelVersion"},{"sourceId":215844687,"sourceType":"kernelVersion"},{"sourceId":216017958,"sourceType":"kernelVersion"},{"sourceId":216986694,"sourceType":"kernelVersion"}],"dockerImageVersionId":30823,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":" !pip install rtdl_num_embeddings -q --no-index --find-links=/kaggle/input/jane-street-import/rtdl_num_embeddings","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-01-11T09:38:27.736632Z","iopub.execute_input":"2025-01-11T09:38:27.737012Z","iopub.status.idle":"2025-01-11T09:38:32.157309Z","shell.execute_reply.started":"2025-01-11T09:38:27.736983Z","shell.execute_reply":"2025-01-11T09:38:32.156237Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim\nfrom torch.utils.data import Dataset, DataLoader, TensorDataset\n\nfrom sklearn.model_selection import train_test_split\n\nfrom sklearn.metrics import r2_score\nimport pandas as pd\nimport math\nimport numpy as np\nfrom tqdm import tqdm\nimport polars as pl\nfrom collections import OrderedDict\nimport sys\nfrom tabm_reference import Model, make_parameter_groups\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\nimport kaggle_evaluation.jane_street_inference_server\n\nimport os\n\nimport joblib\n\nfrom pytorch_lightning import (LightningDataModule, LightningModule, Trainer)\nfrom pytorch_lightning.callbacks import Callback\nimport gc","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-11T09:38:32.158586Z","iopub.execute_input":"2025-01-11T09:38:32.158873Z","iopub.status.idle":"2025-01-11T09:38:39.179927Z","shell.execute_reply.started":"2025-01-11T09:38:32.158843Z","shell.execute_reply":"2025-01-11T09:38:39.179276Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"feature_list = [f\"feature_{idx:02d}\" for idx in range(79) if idx != 61]\n\ntarget_col = \"responder_6\" \ntarget_col2 = \"responder_3\" \n\nfeature_test = feature_list + [f\"responder_{idx}_lag_1\" for idx in range(9)]\\\n                    + ['sin_time_id', 'cos_time_id', 'sin_time_id_halfday', 'cos_time_id_halfday', 'sin_feature_61', 'cos_feature_61']\n\n\nfeature_cat = [\"feature_09\", \"feature_10\", \"feature_11\", \"symbol_id\", \"time_id\"]\nfeature_cont = [item for item in feature_test if item not in feature_cat]\n\nbatch_size = 8192\n\nstd_feature = [i for i in feature_list if i not in feature_cat] + [f\"responder_{idx}_lag_1\" for idx in range(9)]\n\nlag_statscols_rename = { f\"responder_{idx}_lag_1\" : f\"responder_{idx}\" for idx in range(9)}\n\ndata_stats = joblib.load(\"/kaggle/input/js-gbdt-tabm/data_stats.pkl\")\nmeans = data_stats['mean']\nstds = data_stats['std']\n\ndef 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    ])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-11T09:38:39.181194Z","iopub.execute_input":"2025-01-11T09:38:39.181618Z","iopub.status.idle":"2025-01-11T09:38:39.195031Z","shell.execute_reply.started":"2025-01-11T09:38:39.181594Z","shell.execute_reply":"2025-01-11T09:38:39.194356Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"category_mappings = {'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)}}\n\ndef 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=pl.Int64).alias(column)\n    )\n\nclass 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-8)\n        return loss","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-11T09:38:39.196279Z","iopub.execute_input":"2025-01-11T09:38:39.196593Z","iopub.status.idle":"2025-01-11T09:38:39.205107Z","shell.execute_reply.started":"2025-01-11T09:38:39.196571Z","shell.execute_reply":"2025-01-11T09:38:39.204485Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# TabM","metadata":{}},{"cell_type":"code","source":"\n\nclass NN(LightningModule):\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\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,\n                    'dropout': 0,\n                },\n                bins=None,\n                num_embeddings= None,\n                arch_type='tabm',\n                k=self.k,\n            )\n\n        self.lr = lr\n        self.weight_decay = weight_decay\n        self.training_step_outputs = []\n        self.validation_step_outputs = []\n        self.n_classes = n_classes\n        self.loss_fn = R2Loss()\n        # self.loss_fn = nn.MSELoss()\n\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_data, y_ol, y2_ol = batch\n        X_ol = X_data[:, :-2]\n        symbol_ol = X_data[:, -2]\n        time_ol = X_data[:, -1]\n\n        x_cont_ol = X_ol[:, [col for col in range(X_ol.shape[1]) if col not in [9, 10, 11]]]\n        x_cont_ol = x_cont_ol + torch.randn_like(x_cont_ol) * 0.02\n\n        x_cat_ol = X_ol[:, [9, 10, 11]]\n        x_cat_ol = (torch.concat([x_cat_ol, symbol_ol.unsqueeze(-1), time_ol.unsqueeze(-1)], axis=1)).to(torch.int64)\n\n        y_hat = self(x_cont_ol, x_cat_ol)\n\n\n        loss1 = self.loss_fn(y_hat[:, :, 0].flatten(0, 1), y_ol.repeat_interleave(self.k))\n        loss2 = self.loss_fn(y_hat[:, :, 1].flatten(0, 1), y2_ol.repeat_interleave(self.k))\n        loss = 0.85 * loss1 + 0.15 * loss2\n\n        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True, batch_size=x_cont_ol.size(0))\n\n        return loss\n\n\n    def configure_optimizers(self):\n        optimizer = torch.optim.AdamW(make_parameter_groups(self.model), lr=self.lr, weight_decay=self.weight_decay, eps=1e-4)\n        return {\n            'optimizer': optimizer,\n\n        }\n\n    def on_train_epoch_end(self):\n        if self.trainer.sanity_checking:\n            return\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}\")\n        \n\ndevice = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-11T09:38:39.205900Z","iopub.execute_input":"2025-01-11T09:38:39.206159Z","iopub.status.idle":"2025-01-11T09:38:39.244977Z","shell.execute_reply.started":"2025-01-11T09:38:39.206132Z","shell.execute_reply":"2025-01-11T09:38:39.244395Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# cat_cardinalities = [23, 10, 32, 40, 969]\n# device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n# pl_model = Model(\n#         n_num_features=len(feature_cont),\n#         cat_cardinalities=cat_cardinalities,\n#         n_classes=2,\n#         backbone={\n#             'type': 'MLP',\n#             'n_blocks': 3 ,\n#             'd_block': 512,\n#             'dropout': 0,\n#         },\n#         bins=None,\n#         num_embeddings= None,\n#         # cat_embeddings={\n#         #     None if cat_cardinalities is None else\n#         #         'type': 'TrainablePositionEncoding',\n#         #         'd_embedding' : [32, 32, 32, 32, 64],\n#         #         'cardinality' : cat_cardinalities,\n#         # },\n#         # cat_dmodel = [32, 32, 32, 32, 64],\n#         arch_type='tabm',\n#         k=16,\n#     )\n\n# pl_model.load_state_dict(torch.load(\"/kaggle/input/js-gbdt-tabm/tabmv2_ol.model\"))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-11T09:38:39.245602Z","iopub.execute_input":"2025-01-11T09:38:39.245818Z","iopub.status.idle":"2025-01-11T09:38:39.248891Z","shell.execute_reply.started":"2025-01-11T09:38:39.245798Z","shell.execute_reply":"2025-01-11T09:38:39.248263Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# valid_original = pl.scan_parquet(\"/kaggle/input/jane-street-data-preprocessing/validation.parquet\").sort(['date_id', 'time_id', 'symbol_id'])         \n# for col in feature_cat:\n#     valid_original = encode_column(valid_original, col, category_mappings[col])\n \n# valid_data = valid_original \\\n#              .select(feature_test + [target_col, 'weight', 'symbol_id', 'time_id'])\n\n# valid_data_tensor = torch.tensor(valid_data.collect().to_numpy(), dtype=torch.float32)\n# valid_ds = TensorDataset(valid_data_tensor)\n# valid_dl = DataLoader(valid_ds, batch_size=batch_size, num_workers=4, pin_memory=True, shuffle=False)\n# del valid_data_tensor\n# gc.collect()\n\n# print(\"Load done\")\n\n# def r2_val(y_true, y_pred, sample_weight):\n#     residuals = sample_weight * (y_true - y_pred) ** 2\n#     weighted_residual_sum = np.sum(residuals)\n\n#     # Calculate weighted sum of squared true values (denominator)\n#     weighted_true_sum = np.sum(sample_weight * (y_true) ** 2)\n\n#     # Calculate weighted R2\n#     r2 = 1 - weighted_residual_sum / weighted_true_sum\n\n#     return r2\n\n# pl_model.eval()\n# valid_pred_list = []\n# for valid_tensor in tqdm(valid_dl):\n#     X_valid = valid_tensor[0][:, :-4].to(device)\n#     y_valid = valid_tensor[0][:, -4].to(device)\n#     w_valid = valid_tensor[0][:, -3].to(device)\n#     symbol_valid = valid_tensor[0][:, -2].to(device)\n#     time_valid = valid_tensor[0][:, -1].to(device)\n\n    \n#     x_cont_valid = X_valid[:, [col for col in range(X_valid.shape[1]) if col not in [9, 10, 11]]]\n#     # x_cont_valid = x_cont_valid + torch.randn_like(x_cont_valid) * 0.02\n    \n#     x_cat_valid = X_valid[:, [9, 10, 11]]\n#     x_cat_valid = (torch.concat([x_cat_valid, symbol_valid.unsqueeze(-1),time_valid.unsqueeze(-1)], axis=1)).to(torch.int64)\n\n#     with torch.no_grad():\n#         y_pred = pl_model(x_cont_valid, x_cat_valid)\n\n    \n#     valid_pred_list.append((y_pred[:, :, 0].mean(1), y_valid, w_valid))\n# weights_eval = torch.cat([x[2] for x in valid_pred_list]).cpu().numpy()\n# y_eval = torch.cat([x[1] for x in valid_pred_list]).cpu().numpy()\n# prob_eval = torch.cat([x[0] for x in valid_pred_list]).cpu().numpy()\n# val_r2 = r2_val(y_eval, prob_eval, weights_eval)\n\n# print(val_r2)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-11T09:38:39.249647Z","iopub.execute_input":"2025-01-11T09:38:39.249903Z","iopub.status.idle":"2025-01-11T09:38:39.259987Z","shell.execute_reply.started":"2025-01-11T09:38:39.249884Z","shell.execute_reply":"2025-01-11T09:38:39.259257Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"previous_test: pl.DataFrame | None = None\nlags_: pl.DataFrame | None = None\npre_train: pl.DataFrame | None = None\nday_counts = 0\n\n# pl_model = NN.load_from_checkpoint('/kaggle/input/js-gbdt-tabm/tabmv2_epoch05.ckpt')\npl_model = NN.load_from_checkpoint('/kaggle/input/js-gbdt-tabm/tabmv2_original.ckpt').model\n# 清理GPU内存\ntorch.cuda.empty_cache()\n\n\nprint('2 gpus')\npl_model = nn.DataParallel(pl_model, device_ids=[0, 1])\npl_model.cuda(device=0)\n\ndef predict(test: pl.DataFrame, lags: pl.DataFrame | None) -> pl.DataFrame | pd.DataFrame:\n    \"\"\"Make a prediction with incremental learning using the entire dataset.\"\"\"\n    global lags_\n    global previous_test\n    global pre_train\n    global pl_model\n    global day_counts\n\n    # Update lags if provided\n    if lags is not None:\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        lags = lags.with_columns(\n            date_id = pl.col('date_id') + 1\n        )\n        lags_ = lags\n\n\n    # Encode categorical features\n    for col in feature_cat:\n        test = encode_column(test, col, category_mappings[col])\n        if (lags is not None) and (col == 'symbol_id' or col == 'time_id'):\n            lags = encode_column(lags, col, category_mappings[col])\n\n    # Add time-based features\n    test = test.with_columns([\n        (2 * np.pi * pl.col('time_id') / 967).sin().alias('sin_time_id'),\n        (2 * np.pi * pl.col('time_id') / 967).cos().alias('cos_time_id'),\n        (2 * np.pi * pl.col('time_id') / 483).sin().alias('sin_time_id_halfday'),\n        (2 * np.pi * pl.col('time_id') / 483).cos().alias('cos_time_id_halfday'),\n        (2 * np.pi * pl.col('feature_61') / 20).sin().alias('sin_feature_61'),\n        (2 * np.pi * pl.col('feature_61') / 20).cos().alias('cos_feature_61'),\n    ])\n\n    # 在开始时清理 GPU 缓存\n    torch.cuda.empty_cache()\n\n    # Incremental Learning Part\n    if previous_test is not None and lags is not None:\n        # Join previous test data with lags for training\n        train = previous_test.join(\n            lags_.select([\"time_id\", \"symbol_id\", pl.col(\"responder_6_lag_1\").alias(\"responder_6\"), pl.col(\"responder_3_lag_1\").alias(\"responder_3\")]),\n            on=[\"time_id\", \"symbol_id\"],\n            how=\"left\"\n        )\n        train = train.fill_null(0)\n\n        # Manage training data size\n        if pre_train is not None and len(pre_train) > 10000:\n            pre_train = pre_train.sample(n=10000, seed=2025)\n\n        \n        if pre_train is None:\n            pre_train = train\n        else:\n            pre_train = pl.concat([pre_train, train])\n        # 检查并保留最后五万个样本\n        if len(pre_train) > 50000:\n            pre_train = pre_train.tail(50000)\n        \n        print(len(pre_train))\n        # print(pre_train)\n\n        # Prepare training data\n        train_data = pre_train\n        \n        X_train = train_data[feature_test].to_numpy()\n        y_train = pre_train['responder_6'].to_numpy()\n        y2_train = pre_train['responder_3'].to_numpy()\n        \n        # Convert to tensors\n        X_train_tensor = torch.tensor(X_train, dtype=torch.float32).to(device)\n        symbol_ids_train = train_data.select('symbol_id').to_numpy()[:, 0]\n        time_ids_train = train_data.select('time_id').to_numpy()[:, 0]\n        \n        X_cat_train = X_train_tensor[:, [9, 10, 11]]\n        X_cont_train = X_train_tensor[:, [i for i in range(X_train_tensor.shape[1]) if i not in [9, 10, 11]]]\n        \n        symbol_tensor_train = torch.tensor(symbol_ids_train, dtype=torch.float32).to(device)\n        time_tensor_train = torch.tensor(time_ids_train, dtype=torch.float32).to(device)\n        X_cat_train = torch.concat([X_cat_train, symbol_tensor_train.unsqueeze(-1), \n                                  time_tensor_train.unsqueeze(-1)], axis=1).to(torch.int64)\n\n        y_train_tensor = torch.FloatTensor(y_train).to(device)\n        y2_train_tensor = torch.FloatTensor(y2_train).to(device)\n\n\n        print()\n        # 训练循环\n        pl_model.train()\n        optimizer = torch.optim.AdamW(make_parameter_groups(pl_model), lr=1e-4, weight_decay=8e-4)\n        # criterion = nn.MSELoss()\n        criterion = R2Loss()\n        \n        for epoch in range(5):\n            outputs = pl_model(X_cont_train, X_cat_train)\n            loss1 = criterion(outputs[:, :, 0].flatten(0, 1), y_train_tensor.repeat_interleave(16))\n            loss2 = criterion(outputs[:, :, 1].flatten(0, 1), y2_train_tensor.repeat_interleave(16))\n            loss = 0.85 * loss1 + 0.15 * loss2\n            # outputs = pl_model(X_cont_train, X_cat_train)\n            # loss = criterion(outputs[:, :, 0].flatten(0, 1), y_train_tensor.repeat_interleave(16))\n            loss.backward()\n            optimizer.step()\n            \n            print(f\"Epoch {epoch + 1}/3, Loss: {loss.item()}\")\n        \n        previous_test = pl.DataFrame()\n\n    # Prepare test data for prediction\n    time_id = test.select(\"time_id\").to_numpy()[0]\n    symbol_ids = test.select('symbol_id').to_numpy()[:, 0]\n    timie_id_array = test.select(\"time_id\").to_numpy()[:, 0]\n\n    if not lags_ is None:\n\n        # lags_feature = lags_.filter(pl.col(\"time_id\") == time_id)\n        # lags_feature = lags_feature.drop('time_id')\n        # test = test.join(lags_feature, on=[\"date_id\", \"symbol_id\"], how=\"left\") # 用来预测\n\n        # 下面是debug的时候用\n        lags_feature = lags_.group_by([\"date_id\", \"symbol_id\"], maintain_order=True).last()\n        lags_feature = lags_feature.drop('time_id')\n        test = test.join(lags_feature, on=[\"date_id\", \"symbol_id\"], how=\"left\")\n\n\n    else:\n        test = test.with_columns(\n            (pl.lit(0.0).alias(f'responder_{idx}_lag_1') for idx in range(9))\n        )\n\n    # test_ol = test_ol.with_columns([\n    #     pl.col(col).fill_null(0) for col in feature_test + ['symbol_id', 'time_id']\n    # ])\n    \n    test = test.with_columns([\n        pl.col(col).fill_null(0) for col in feature_test + ['symbol_id', 'time_id']\n    ])\n    \n    test = standardize(test, std_feature, means, stds)\n    # test_ol = standardize(test_ol, std_feature, means, stds)\n    \n    # Prepare tensors for prediction\n    X_test = test[feature_test].to_numpy()\n    X_test_tensor = torch.tensor(X_test, dtype=torch.float32).to(device)\n    \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    X_cat = X_test_tensor[:, [9, 10, 11]]\n    X_cont = X_test_tensor[:, [i for i in range(X_test_tensor.shape[1]) if i not in [9, 10, 11]]]\n    \n    X_cat = torch.concat([X_cat, symbol_tensor.unsqueeze(-1), time_tensor.unsqueeze(-1)], \n                        axis=1).to(torch.int64)\n\n    # Make predictions\n    pl_model.eval()\n    with torch.no_grad():\n        outputs = pl_model(X_cont, X_cat)  # shape: [batch_size, 16, 2]\n        outputs = outputs[:, :, 0]  # 取第一个预测值\n        preds = outputs.mean(dim=1).cpu().numpy()  # 在ensemble维度上取平均\n\n    # Store current test data for next iteration\n    if previous_test is None:\n        previous_test = test\n    else:\n        previous_test = pl.concat([previous_test, test])\n\n\n    # Format predictions\n    try:\n        predictions = test.select('row_id').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    except:\n        print(preds.shape)\n\n    # Validate output\n    if isinstance(predictions, pl.DataFrame):\n        assert predictions.columns == ['row_id', 'responder_6']\n    elif isinstance(predictions, pd.DataFrame):\n        assert (predictions.columns == ['row_id', 'responder_6']).all()\n    else:\n        raise TypeError('The predict function must return a DataFrame')\n    assert len(predictions) == len(test)\n\n    return predictions","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-11T05:26:17.008867Z","iopub.execute_input":"2025-01-11T05:26:17.009175Z","iopub.status.idle":"2025-01-11T05:26:17.090657Z","shell.execute_reply.started":"2025-01-11T05:26:17.009152Z","shell.execute_reply":"2025-01-11T05:26:17.090001Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%time\n\nEVAL =  True\nif EVAL:\n    test_dir = '/kaggle/input/k/i2nfinit3y/janestreet-updated-simulator-for-time-series-api/debug/test.parquet'\n    lags_dir = '/kaggle/input/k/i2nfinit3y/janestreet-updated-simulator-for-time-series-api/debug/lags.parquet'\n    # test_dir = '/kaggle/input/js24-rmf-submission-api-debug-with-synthetic-test/synthetic_test.parquet'\n    # lags_dir = '/kaggle/input/js24-rmf-submission-api-debug-with-synthetic-test/synthetic_lag.parquet'\nelse:\n    test_dir = '/kaggle/input/jane-street-real-time-market-data-forecasting/test.parquet'\n    lags_dir = '/kaggle/input/jane-street-real-time-market-data-forecasting/lags.parquet'\n\ninference_server = kaggle_evaluation.jane_street_inference_server.JSInferenceServer(predict)\n\nif os.getenv('KAGGLE_IS_COMPETITION_RERUN'):\n    inference_server.serve()\nelse:\n    inference_server.run_local_gateway(\n        (\n            test_dir,\n            lags_dir\n        )\n    )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-11T05:26:18.712953Z","iopub.execute_input":"2025-01-11T05:26:18.713230Z","iopub.status.idle":"2025-01-11T05:26:25.415443Z","shell.execute_reply.started":"2025-01-11T05:26:18.713207Z","shell.execute_reply":"2025-01-11T05:26:25.414765Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def weighted_zero_mean_r2(y_true, y_pred, weights):\n    \"\"\"\n    Calculate the sample weighted zero-mean R-squared score.\n\n    Parameters:\n    y_true (numpy.ndarray): Ground-truth values for responder_6.\n    y_pred (numpy.ndarray): Predicted values for responder_6.\n    weights (numpy.ndarray): Sample weight vector.\n\n    Returns:\n    float: The weighted zero-mean R-squared score.\n    \"\"\"\n    numerator = np.sum(weights * (y_true - y_pred)**2)\n    denominator = np.sum(weights * y_true**2)\n    \n    r2_score = 1 - numerator / denominator\n    return r2_score","metadata":{"trusted":true,"execution":{"execution_failed":"2025-01-11T02:32:22.694Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if EVAL:\n\n    submission_file = pd.read_parquet('/kaggle/working/submission.parquet')\n    y_pred = submission_file['responder_6']\n    \n    valid_df = pl.read_parquet(\"/kaggle/input/k/i2nfinit3y/janestreet-updated-simulator-for-time-series-api/valid_df.parquet\")\n    \n    y_true = valid_df.select(\"responder_6\").to_numpy().reshape(-1)\n    \n    weights = valid_df.select(\"weight\").to_numpy().reshape(-1)\n    \n    print(weighted_zero_mean_r2(y_true, y_pred, weights))","metadata":{"trusted":true,"execution":{"execution_failed":"2025-01-11T02:32:22.694Z"}},"outputs":[],"execution_count":null}]}