{"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":9801075,"sourceType":"datasetVersion","datasetId":6006872},{"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":"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\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":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CONFIG:\n   \n    seed = 42\n\n    target_col = \"responder_6\"\n\n   \n    feature_cols = [f\"feature_{idx:02d}\" for idx in range(79)] + [f\"responder_{idx}_lag_1\" for idx in range(9)]\n    \n    \n    model_paths = [\n       \n        \"/kaggle/input/js-xs-nn-trained-model\",  \n        \"/kaggle/input/js-with-lags-trained-xgb/result.pkl\",  \n    ]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"valid = pl.scan_parquet(\n    f\"/kaggle/input/js24-preprocessing-create-lags/validation.parquet/\"\n).collect().to_pandas()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"xgb_model = None\nmodel_path = CONFIG.model_paths[1]\nwith open(model_path, \"rb\") as fp:\n    result = pickle.load(fp)\n    xgb_model = result[\"model\"]\nxgb_feature_cols = [\"symbol_id\", \"time_id\"] + CONFIG.feature_cols\ndisplay(xgb_model)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def r2_val(y_true, y_pred, sample_weight):\n   \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   \n    def __init__(self, input_dim, hidden_dims, dropouts, lr, weight_decay):\n        super().__init__()\n       \n        self.save_hyperparameters()\n\n        layers = []\n        in_dim = input_dim\n        \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            in_dim = hidden_dim  \n            \n        layers.append(nn.Linear(in_dim, 1))  \n        layers.append(nn.Tanh())  \n        self.model = nn.Sequential(*layers)  \n        \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\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            \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\n    def configure_optimizers(self):\n        \n        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)\n       \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\n    def on_train_epoch_end(self):\n        if self.trainer.sanity_checking:\n            return  \n        epoch = self.trainer.current_epoch  \n        \n        metrics = {k: v.item() if isinstance(v, torch.Tensor) else v for k, v in self.trainer.logged_metrics.items()}\n        \n        formatted_metrics = {k: f\"{v:.5f}\" for k, v in metrics.items()}\n        \n        print(f\"Epoch {epoch}: {formatted_metrics}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"N_folds = 5\n\nmodels = []\n\nfor fold in range(N_folds):\n\n    checkpoint_path = f\"{CONFIG.model_paths[0]}/nn_{fold}.model\"\n    \n    model = NN.load_from_checkpoint(checkpoint_path)\n\n    models.append(model.to(\"cuda:0\"))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X_valid = valid[xgb_feature_cols]  \ny_valid = valid[CONFIG.target_col]  \n\nw_valid = valid[\"weight\"]\n\n\ny_pred_valid_xgb = xgb_model.predict(X_valid)\n\n\nvalid_score = r2_score(y_valid, y_pred_valid_xgb, sample_weight=w_valid)\n\nvalid_score","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X_valid = valid[CONFIG.feature_cols]  \ny_valid = valid[CONFIG.target_col]  \n\nw_valid = valid[\"weight\"]\n\n\nX_valid = X_valid.fillna(method='ffill').fillna(0)\n\nX_valid.shape, y_valid.shape, w_valid.shape","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"y_pred_valid_nn = np.zeros(y_valid.shape)\n\nwith torch.no_grad():\n    for model in models:\n        model.eval()  \n\n        y_pred_valid_nn += model(torch.FloatTensor(X_valid.values).to(\"cuda:0\")).cpu().numpy() / len(models)\n\nvalid_score = r2_score(y_valid, y_pred_valid_nn, sample_weight=w_valid)\n\nvalid_score","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"y_pred_valid_ensemble = 0.5 * (y_pred_valid_xgb + y_pred_valid_nn)\n\n\nvalid_score = r2_score(y_valid, y_pred_valid_ensemble, sample_weight=w_valid)\n\n\nvalid_score","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"del valid, X_valid, y_valid, w_valid\n\ngc.collect()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"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    \n\n    if lags is not None:\n        lags_ = lags\n\n\n    predictions = test.select(\n        'row_id',  \n        pl.lit(0.0).alias('responder_6'),  \n    )\n    \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()  \n        test = test.join(lags, on=[\"date_id\", \"symbol_id\"], how=\"left\")  \n    else:\n\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    \n    preds += xgb_model.predict(test[xgb_feature_cols].to_pandas()) / 2  \n\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\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() / 10  \n\n    print(f\"predict> preds.shape =\", preds.shape)\n\n    predictions = 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\n    assert isinstance(predictions, pl.DataFrame | pd.DataFrame)\n\n    assert list(predictions.columns) == ['row_id', 'responder_6']\n\n\n    assert len(predictions) == len(test)\n\n    return predictions  ","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"inference_server = kaggle_evaluation.jane_street_inference_server.JSInferenceServer(predict)\n\n\nif os.getenv('KAGGLE_IS_COMPETITION_RERUN'):\n    inference_server.serve()  \nelse:\n\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},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}