{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","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":9801075,"sourceType":"datasetVersion","datasetId":6006872},{"sourceId":9806342,"sourceType":"datasetVersion","datasetId":6010899},{"sourceId":203900450,"sourceType":"kernelVersion"},{"sourceId":207107641,"sourceType":"kernelVersion"}],"dockerImageVersionId":30787,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# **JANE STREER MODEL**\n","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport polars as pl\nimport numpy as np\nimport os\nimport 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\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\n# Configurations\nclass CONFIG:\n    seed = 42\n    target_col = \"responder_6\"\n    feature_cols = [f\"feature_{idx:02d}\" for idx in range(79)] + [f\"responder_{idx}_lag_1\" for idx in range(9)]\n    model_paths = [\n        \"/kaggle/input/js-xs-nn-trained-model\",\n        \"/kaggle/input/js-with-lags-trained-xgb/result.pkl\",\n    ]\n\n# 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\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            in_dim = hidden_dim\n        layers.append(nn.Linear(in_dim, 1))  # Output layer\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)  # Output as a 1D tensor\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 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        y = torch.cat([x[1] for x in self.validation_step_outputs]).cpu().numpy()\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.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, verbose=True)\n        return {'optimizer': optimizer, 'lr_scheduler': {'scheduler': scheduler, 'monitor': 'val_loss'}}\n\n# Load XGBoost model\nxgb_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\"]\n\nxgb_feature_cols = [\"symbol_id\", \"time_id\"] + CONFIG.feature_cols\n\n# Load the best neural network models\nN_folds = 5\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\"))\n\n# Cross-validation score calculation for the validation set\nvalid = pl.scan_parquet(f\"/kaggle/input/js24-preprocessing-create-lags/validation.parquet/\").collect().to_pandas()\nX_valid = valid[xgb_feature_cols]\ny_valid = valid[CONFIG.target_col]\nw_valid = valid[\"weight\"]\n\ny_pred_valid_xgb = xgb_model.predict(X_valid)\nvalid_score_xgb = r2_score(y_valid, y_pred_valid_xgb, sample_weight=w_valid)\n\nX_valid = valid[CONFIG.feature_cols]\nX_valid = X_valid.fillna(method='ffill').fillna(0)\nX_valid_tensor = torch.FloatTensor(X_valid.values).to(\"cuda:0\")\n\ny_pred_valid_nn = np.zeros(y_valid.shape)\nwith torch.no_grad():\n    for model in models:\n        model.eval()\n        y_pred_valid_nn += model(X_valid_tensor).cpu().numpy() / N_folds\n\nvalid_score_nn = r2_score(y_valid, y_pred_valid_nn, sample_weight=w_valid)\n\n# Weighted average of XGBoost and NN based on CV scores\nensemble_weight_xgb = valid_score_xgb / (valid_score_xgb + valid_score_nn)\nensemble_weight_nn = valid_score_nn / (valid_score_xgb + valid_score_nn)\n\ny_pred_valid_ensemble = ensemble_weight_xgb * y_pred_valid_xgb + ensemble_weight_nn * y_pred_valid_nn\nvalid_score = r2_score(y_valid, y_pred_valid_ensemble, sample_weight=w_valid)\n\n# Clean up memory\ndel valid, X_valid, y_valid, w_valid\ngc.collect()\n\n# Define the inference function\ndef predict(test: pl.DataFrame, lags: pl.DataFrame | None) -> pl.DataFrame:\n    global lags_\n    if lags is not None:\n        lags_ = lags\n\n    predictions = test.select('row_id', pl.lit(0.0).alias('responder_6'))\n    symbol_ids = test.select('symbol_id').to_numpy()[:, 0]\n\n    if lags is not None:\n        lags = lags.group_by([\"date_id\", \"symbol_id\"], maintain_order=True).last()  # Pick last record of the previous date\n        test = test.join(lags, on=[\"date_id\", \"symbol_id\"], how=\"left\")\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    preds += xgb_model.predict(test[xgb_feature_cols].to_pandas()) / 2\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 nn_model in models:\n            nn_model.eval()\n            preds += nn_model(test_input).cpu().numpy() / N_folds\n    \n    # Ensure predictions are within a reasonable range\n    preds = np.clip(preds, a_min=-5, a_max=5)\n    \n    predictions = test.select('row_id').with_columns(\n        pl.Series(name='responder_6', values=preds, dtype=pl.Float64)\n    )\n\n    assert isinstance(predictions, pl.DataFrame)\n    assert list(predictions.columns) == ['row_id', 'responder_6']\n    assert len(predictions) == len(test)\n\n    return predictions\n\n# Set up the inference server\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        '/kaggle/input/jane-street-real-time-market-data-forecasting/test.parquet',\n        '/kaggle/input/jane-street-real-time-market-data-forecasting/lags.parquet',\n    ))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-14T18:59:54.537209Z","iopub.execute_input":"2024-11-14T18:59:54.537593Z","iopub.status.idle":"2024-11-14T19:00:02.986666Z","shell.execute_reply.started":"2024-11-14T18:59:54.537560Z","shell.execute_reply":"2024-11-14T19:00:02.985151Z"}},"outputs":[],"execution_count":null}]}