{"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.10.14"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":84493,"databundleVersionId":9871156,"sourceType":"competition"},{"sourceId":9806342,"sourceType":"datasetVersion","datasetId":6010899},{"sourceId":203900450,"sourceType":"kernelVersion"}],"dockerImageVersionId":30787,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":7.594014,"end_time":"2024-10-10T11:58:36.355301","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2024-10-10T11:58:28.761287","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Useful notebooks:\n\n- Preprocessing : https://www.kaggle.com/code/motono0223/js24-preprocessing-create-lags\n- Training (XGB) : https://www.kaggle.com/code/motono0223/js24-train-gbdt-model-with-lags-singlemodel\n  - trained XGB model : https://www.kaggle.com/datasets/motono0223/js24-trained-gbdt-model\n- Training (NN): https://www.kaggle.com/code/voix97/jane-street-rmf-training-nn\n  - trained NN model : https://www.kaggle.com/datasets/voix97/js-xs-nn-trained-model\n- Inference of NN : **this notebook** https://www.kaggle.com/code/voix97/jane-street-rmf-nn-with-pytorch-lightning\n- Inference of NN+XGB:  https://www.kaggle.com/code/voix97/jane-street-rmf-nn-xgb\n- EDA(1) : https://www.kaggle.com/code/motono0223/eda-jane-street-real-time-market-data-forecasting\n- EDA(2) : https://www.kaggle.com/code/motono0223/eda-v2-jane-street-real-time-market-forecasting","metadata":{}},{"cell_type":"markdown","source":"# Pytorch Lightning Neural Networks Inference ","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport polars as pl\nimport numpy as np\nimport os, gc\nfrom tqdm.auto import tqdm\nfrom matplotlib import pyplot as plt\nimport pickle\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pytorch_lightning import (LightningDataModule, LightningModule, Trainer)\nfrom pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint, Timer\n\nimport pandas as pd\nimport numpy as np\nfrom sklearn.metrics import r2_score\nfrom sklearn.model_selection import train_test_split\nfrom torch.utils.data import Dataset, DataLoader\n\n\nfrom sklearn.metrics import r2_score\nfrom lightgbm import LGBMRegressor\nimport lightgbm as lgb\nfrom xgboost import XGBRegressor\nfrom catboost import CatBoostRegressor\nfrom sklearn.ensemble import VotingRegressor\n\nimport warnings\nwarnings.filterwarnings('ignore')\npd.options.display.max_columns = None\n\nimport kaggle_evaluation.jane_street_inference_server","metadata":{"execution":{"iopub.status.busy":"2024-11-13T15:54:37.642612Z","iopub.execute_input":"2024-11-13T15:54:37.642997Z","iopub.status.idle":"2024-11-13T15:54:51.996933Z","shell.execute_reply.started":"2024-11-13T15:54:37.642952Z","shell.execute_reply":"2024-11-13T15:54:51.995974Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Configurations","metadata":{}},{"cell_type":"code","source":"class CONFIG:\n    seed = 42\n    target_col = \"responder_6\"\n    # feature_cols = [\"symbol_id\", \"time_id\"] + [f\"feature_{idx:02d}\" for idx in range(79)]+ [f\"responder_{idx}_lag_1\" for idx in range(9)]\n    feature_cols = [f\"feature_{idx:02d}\" for idx in range(79)]+ [f\"responder_{idx}_lag_1\" for idx in range(9)]\n    \n    model_paths = [\n        #\"/kaggle/input/js24-train-gbdt-model-with-lags-singlemodel/result.pkl\",\n        #\"/kaggle/input/js24-trained-gbdt-model/result.pkl\",\n        \"/kaggle/input/js-xs-nn-trained-model\",\n    ]","metadata":{"execution":{"iopub.status.busy":"2024-11-13T15:54:51.998173Z","iopub.execute_input":"2024-11-13T15:54:51.998985Z","iopub.status.idle":"2024-11-13T15:54:52.004625Z","shell.execute_reply.started":"2024-11-13T15:54:51.998947Z","shell.execute_reply":"2024-11-13T15:54:52.003646Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Load preprocessed data (to calculate CV)","metadata":{}},{"cell_type":"code","source":"valid = pl.scan_parquet(\n    f\"/kaggle/input/js24-preprocessing-create-lags/validation.parquet/\"\n).collect().to_pandas()","metadata":{"execution":{"iopub.status.busy":"2024-11-13T15:54:52.005839Z","iopub.execute_input":"2024-11-13T15:54:52.006169Z","iopub.status.idle":"2024-11-13T15:54:53.896950Z","shell.execute_reply.started":"2024-11-13T15:54:52.006136Z","shell.execute_reply":"2024-11-13T15:54:53.895840Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Load model","metadata":{}},{"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\n\n\nclass NN(LightningModule):\n    def __init__(self, input_dim, hidden_dims, dropouts, lr, weight_decay):\n        super().__init__()\n        self.save_hyperparameters()\n        layers = []\n        in_dim = input_dim\n        for i, hidden_dim in enumerate(hidden_dims):\n            layers.append(nn.BatchNorm1d(in_dim))\n            if i > 0:\n                layers.append(nn.SiLU())\n            if i < len(dropouts):\n                layers.append(nn.Dropout(dropouts[i]))\n            layers.append(nn.Linear(in_dim, hidden_dim))\n            # layers.append(nn.ReLU())\n            in_dim = hidden_dim\n        layers.append(nn.Linear(in_dim, 1))  # 输出层\n        layers.append(nn.Tanh())\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)  # 输出为一维张量\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  # 考虑样本权重\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        \"\"\"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            # r2_val\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.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)\n        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5,\n                                                               verbose=True)\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 self.trainer.sanity_checking:\n            return\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":"2024-11-13T15:54:53.899762Z","iopub.execute_input":"2024-11-13T15:54:53.900189Z","iopub.status.idle":"2024-11-13T15:54:53.921502Z","shell.execute_reply.started":"2024-11-13T15:54:53.900144Z","shell.execute_reply":"2024-11-13T15:54:53.920566Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"N_folds = 5\n# 加载最佳模型\nmodels = []\nfor fold in range(N_folds):\n    checkpoint_path = f\"{CONFIG.model_paths[0]}/nn_{fold}.model\"\n    model = NN.load_from_checkpoint(checkpoint_path)\n    models.append(model.to(\"cuda:0\"))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T15:54:53.922882Z","iopub.execute_input":"2024-11-13T15:54:53.923236Z","iopub.status.idle":"2024-11-13T15:54:55.021580Z","shell.execute_reply.started":"2024-11-13T15:54:53.923202Z","shell.execute_reply":"2024-11-13T15:54:55.020737Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"models[0]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T15:55:04.368460Z","iopub.execute_input":"2024-11-13T15:55:04.368844Z","iopub.status.idle":"2024-11-13T15:55:04.377124Z","shell.execute_reply.started":"2024-11-13T15:55:04.368806Z","shell.execute_reply":"2024-11-13T15:55:04.376083Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# CV Score","metadata":{}},{"cell_type":"code","source":"X_valid = valid[ CONFIG.feature_cols ]\ny_valid = valid[ CONFIG.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":{"execution":{"iopub.status.busy":"2024-11-07T17:39:37.640886Z","iopub.execute_input":"2024-11-07T17:39:37.641226Z","iopub.status.idle":"2024-11-07T17:39:39.747550Z","shell.execute_reply.started":"2024-11-07T17:39:37.641191Z","shell.execute_reply":"2024-11-07T17:39:39.746562Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"y_pred_valid = np.zeros(y_valid.shape)\nwith torch.no_grad():\n    for model in models:\n        model.eval()\n        y_pred_valid += model(torch.FloatTensor(X_valid.values).to(\"cuda:0\")).cpu().numpy() / len(models)\nvalid_score = r2_score( y_valid, y_pred_valid, sample_weight=w_valid )\nvalid_score","metadata":{"execution":{"iopub.status.busy":"2024-11-07T17:39:39.748584Z","iopub.execute_input":"2024-11-07T17:39:39.748854Z","iopub.status.idle":"2024-11-07T17:39:46.530347Z","shell.execute_reply.started":"2024-11-07T17:39:39.748824Z","shell.execute_reply":"2024-11-07T17:39:46.529359Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"del valid, X_valid, y_valid, w_valid\ngc.collect()","metadata":{"execution":{"iopub.status.busy":"2024-11-07T17:39:46.531475Z","iopub.execute_input":"2024-11-07T17:39:46.531772Z","iopub.status.idle":"2024-11-07T17:39:46.740179Z","shell.execute_reply.started":"2024-11-07T17:39:46.531739Z","shell.execute_reply":"2024-11-07T17:39:46.739143Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"#### There seems to be bug in official code, can only submit polars dataframe","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    if not lags is None:\n        lags = lags.group_by([\"date_id\", \"symbol_id\"], maintain_order=True).last() # pick up last record of previous date\n        test = test.join(lags, on=[\"date_id\", \"symbol_id\"],  how=\"left\")\n        # print(a)\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    preds = np.zeros((test.shape[0],))\n    # for i, model in enumerate(tqdm(models)):\n    #     preds += model.predict(test[CONFIG.feature_cols].to_pandas()) / len(models)\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).to(\"cuda:0\")\n    with torch.no_grad():\n        for i, nn_model in enumerate(tqdm(models)):\n            nn_model.eval()\n            preds += nn_model(test_input).cpu().numpy() / len(models)\n    print(f\"predict> preds.shape =\", preds.shape)\n    \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    # The predict function must return a DataFrame\n    assert isinstance(predictions, pl.DataFrame | pd.DataFrame)\n    # with columns 'row_id', 'responer_6'\n    assert list(predictions.columns) == ['row_id', 'responder_6']\n    # and as many rows as the test data.\n    assert len(predictions) == len(test)\n\n    return predictions","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","papermill":{"duration":0.018344,"end_time":"2024-10-10T11:58:33.59684","exception":false,"start_time":"2024-10-10T11:58:33.578496","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-11-07T17:51:16.499233Z","iopub.execute_input":"2024-11-07T17:51:16.499648Z","iopub.status.idle":"2024-11-07T17:51:16.512183Z","shell.execute_reply.started":"2024-11-07T17:51:16.499612Z","shell.execute_reply":"2024-11-07T17:51:16.511262Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"When your notebook is run on the hidden test set, inference_server.serve must be called within 15 minutes of the notebook starting or the gateway will throw an error. If you need more than 15 minutes to load your model you can do so during the very first `predict` call, which does not have the usual 10 minute response deadline.","metadata":{"papermill":{"duration":0.002521,"end_time":"2024-10-10T11:58:33.6023","exception":false,"start_time":"2024-10-10T11:58:33.599779","status":"completed"},"tags":[]}},{"cell_type":"code","source":"inference_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            '/kaggle/input/jane-street-realtime-marketdata-forecasting/test.parquet',\n            '/kaggle/input/jane-street-realtime-marketdata-forecasting/lags.parquet',\n        )\n    )","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","papermill":{"duration":2.225871,"end_time":"2024-10-10T11:58:35.830964","exception":false,"start_time":"2024-10-10T11:58:33.605093","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-11-07T17:51:18.438863Z","iopub.execute_input":"2024-11-07T17:51:18.439244Z","iopub.status.idle":"2024-11-07T17:51:18.514114Z","shell.execute_reply.started":"2024-11-07T17:51:18.439200Z","shell.execute_reply":"2024-11-07T17:51:18.513306Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}