js_xbg.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import polars as pl\n",
    "import xgboost as xgb\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import sys\n",
    "\n",
    "sys.path.append(\"/home/pqian/data/kaggle/js-2024\")\n",
    "import kaggle_evaluation.jane_street_inference_server"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# data process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "PATH = \"/home/pqian/data/kaggle/js-2024/\"\n",
    "\n",
    "\n",
    "def read_selected_data(input_path):\n",
    "    # Define the directory containing your data files\n",
    "\n",
    "    # List three specific Parquet files you want to read\n",
    "    selected_files = [f\"partition_id={i}/part-0.parquet\" for i in range(1)]\n",
    "    # Load and filter the data from only the selected Parquet files\n",
    "    dfs = []\n",
    "    for file_name in selected_files:\n",
    "        file_path = f\"{input_path}/train.parquet/{file_name}\"\n",
    "        lazy_df = pl.scan_parquet(file_path)\n",
    "        df = lazy_df.collect()\n",
    "        dfs.append(df)\n",
    "\n",
    "    # Concatenate all dataframes into a single dataframe\n",
    "    full_df = pl.concat(dfs)\n",
    "\n",
    "    return full_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = read_selected_data(PATH)\n",
    "df = df.fill_null(strategy=\"forward\")\n",
    "\n",
    "# Prepare feature names\n",
    "feature_names = [f\"feature_{i:02d}\" for i in range(79)]\n",
    "\n",
    "# Prepare training and validation data\n",
    "num_valid_dates = 100\n",
    "dates = df[\"date_id\"].unique().to_numpy()\n",
    "valid_dates = dates[-num_valid_dates:]\n",
    "train_dates = dates[:-num_valid_dates]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract features, target, and weights for validation and training sets\n",
    "X_valid = (\n",
    "    df.filter(pl.col(\"date_id\").is_in(valid_dates)).select(feature_names).to_numpy()\n",
    ")\n",
    "y_valid = (\n",
    "    df.filter(pl.col(\"date_id\").is_in(valid_dates))\n",
    "    .select(\"responder_6\")\n",
    "    .to_numpy()\n",
    "    .ravel()\n",
    ")\n",
    "w_valid = (\n",
    "    df.filter(pl.col(\"date_id\").is_in(valid_dates)).select(\"weight\").to_numpy().ravel()\n",
    ")\n",
    "\n",
    "X_train = (\n",
    "    df.filter(pl.col(\"date_id\").is_in(train_dates)).select(feature_names).to_numpy()\n",
    ")\n",
    "y_train = (\n",
    "    df.filter(pl.col(\"date_id\").is_in(train_dates))\n",
    "    .select(\"responder_6\")\n",
    "    .to_numpy()\n",
    "    .ravel()\n",
    ")\n",
    "w_train = (\n",
    "    df.filter(pl.col(\"date_id\").is_in(train_dates)).select(\"weight\").to_numpy().ravel()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def r2_xgb(y_true, y_pred, sample_weight=None):\n",
    "    if sample_weight is None:\n",
    "        sample_weight = np.ones_like(y_true)\n",
    "    r2 = 1 - np.average((y_pred - y_true) ** 2, weights=sample_weight) / (\n",
    "        np.average((y_true) ** 2, weights=sample_weight) + 1e-38\n",
    "    )\n",
    "    return -r2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the XGBoost model\n",
    "model = xgb.XGBRegressor(\n",
    "    n_estimators=2000,\n",
    "    learning_rate=0.1,\n",
    "    max_depth=6,\n",
    "    tree_method=\"hist\",\n",
    "    device=\"cuda\",\n",
    "    objective=\"reg:squarederror\",\n",
    "    eval_metric=r2_xgb,\n",
    "    disable_default_eval_metric=True,\n",
    "    early_stopping_rounds=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0]\tvalidation_0-r2_xgb:-0.00187\n",
      "[2]\tvalidation_0-r2_xgb:-0.00485\n",
      "[4]\tvalidation_0-r2_xgb:-0.00651\n",
      "[6]\tvalidation_0-r2_xgb:-0.00752\n",
      "[8]\tvalidation_0-r2_xgb:-0.00884\n",
      "[10]\tvalidation_0-r2_xgb:-0.00957\n",
      "[12]\tvalidation_0-r2_xgb:-0.00988\n",
      "[14]\tvalidation_0-r2_xgb:-0.01034\n",
      "[16]\tvalidation_0-r2_xgb:-0.01059\n",
      "[18]\tvalidation_0-r2_xgb:-0.01105\n",
      "[20]\tvalidation_0-r2_xgb:-0.01097\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<style>#sk-container-id-1 {\n",
       "  /* Definition of color scheme common for light and dark mode */\n",
       "  --sklearn-color-text: black;\n",
       "  --sklearn-color-line: gray;\n",
       "  /* Definition of color scheme for unfitted estimators */\n",
       "  --sklearn-color-unfitted-level-0: #fff5e6;\n",
       "  --sklearn-color-unfitted-level-1: #f6e4d2;\n",
       "  --sklearn-color-unfitted-level-2: #ffe0b3;\n",
       "  --sklearn-color-unfitted-level-3: chocolate;\n",
       "  /* Definition of color scheme for fitted estimators */\n",
       "  --sklearn-color-fitted-level-0: #f0f8ff;\n",
       "  --sklearn-color-fitted-level-1: #d4ebff;\n",
       "  --sklearn-color-fitted-level-2: #b3dbfd;\n",
       "  --sklearn-color-fitted-level-3: cornflowerblue;\n",
       "\n",
       "  /* Specific color for light theme */\n",
       "  --sklearn-color-text-on-default-background: var(--sg-text-color, var(--theme-code-foreground, var(--jp-content-font-color1, black)));\n",
       "  --sklearn-color-background: var(--sg-background-color, var(--theme-background, var(--jp-layout-color0, white)));\n",
       "  --sklearn-color-border-box: var(--sg-text-color, var(--theme-code-foreground, var(--jp-content-font-color1, black)));\n",
       "  --sklearn-color-icon: #696969;\n",
       "\n",
       "  @media (prefers-color-scheme: dark) {\n",
       "    /* Redefinition of color scheme for dark theme */\n",
       "    --sklearn-color-text-on-default-background: var(--sg-text-color, var(--theme-code-foreground, var(--jp-content-font-color1, white)));\n",
       "    --sklearn-color-background: var(--sg-background-color, var(--theme-background, var(--jp-layout-color0, #111)));\n",
       "    --sklearn-color-border-box: var(--sg-text-color, var(--theme-code-foreground, var(--jp-content-font-color1, white)));\n",
       "    --sklearn-color-icon: #878787;\n",
       "  }\n",
       "}\n",
       "\n",
       "#sk-container-id-1 {\n",
       "  color: var(--sklearn-color-text);\n",
       "}\n",
       "\n",
       "#sk-container-id-1 pre {\n",
       "  padding: 0;\n",
       "}\n",
       "\n",
       "#sk-container-id-1 input.sk-hidden--visually {\n",
       "  border: 0;\n",
       "  clip: rect(1px 1px 1px 1px);\n",
       "  clip: rect(1px, 1px, 1px, 1px);\n",
       "  height: 1px;\n",
       "  margin: -1px;\n",
       "  overflow: hidden;\n",
       "  padding: 0;\n",
       "  position: absolute;\n",
       "  width: 1px;\n",
       "}\n",
       "\n",
       "#sk-container-id-1 div.sk-dashed-wrapped {\n",
       "  border: 1px dashed var(--sklearn-color-line);\n",
       "  margin: 0 0.4em 0.5em 0.4em;\n",
       "  box-sizing: border-box;\n",
       "  padding-bottom: 0.4em;\n",
       "  background-color: var(--sklearn-color-background);\n",
       "}\n",
       "\n",
       "#sk-container-id-1 div.sk-container {\n",
       "  /* jupyter's `normalize.less` sets `[hidden] { display: none; }`\n",
       "     but bootstrap.min.css set `[hidden] { display: none !important; }`\n",
       "     so we also need the `!important` here to be able to override the\n",
       "     default hidden behavior on the sphinx rendered scikit-learn.org.\n",
       "     See: https://github.com/scikit-learn/scikit-learn/issues/21755 */\n",
       "  display: inline-block !important;\n",
       "  position: relative;\n",
       "}\n",
       "\n",
       "#sk-container-id-1 div.sk-text-repr-fallback {\n",
       "  display: none;\n",
       "}\n",
       "\n",
       "div.sk-parallel-item,\n",
       "div.sk-serial,\n",
       "div.sk-item {\n",
       "  /* draw centered vertical line to link estimators */\n",
       "  background-image: linear-gradient(var(--sklearn-color-text-on-default-background), var(--sklearn-color-text-on-default-background));\n",
       "  background-size: 2px 100%;\n",
       "  background-repeat: no-repeat;\n",
       "  background-position: center center;\n",
       "}\n",
       "\n",
       "/* Parallel-specific style estimator block */\n",
       "\n",
       "#sk-container-id-1 div.sk-parallel-item::after {\n",
       "  content: \"\";\n",
       "  width: 100%;\n",
       "  border-bottom: 2px solid var(--sklearn-color-text-on-default-background);\n",
       "  flex-grow: 1;\n",
       "}\n",
       "\n",
       "#sk-container-id-1 div.sk-parallel {\n",
       "  display: flex;\n",
       "  align-items: stretch;\n",
       "  justify-content: center;\n",
       "  background-color: var(--sklearn-color-background);\n",
       "  position: relative;\n",
       "}\n",
       "\n",
       "#sk-container-id-1 div.sk-parallel-item {\n",
       "  display: flex;\n",
       "  flex-direction: column;\n",
       "}\n",
       "\n",
       "#sk-container-id-1 div.sk-parallel-item:first-child::after {\n",
       "  align-self: flex-end;\n",
       "  width: 50%;\n",
       "}\n",
       "\n",
       "#sk-container-id-1 div.sk-parallel-item:last-child::after {\n",
       "  align-self: flex-start;\n",
       "  width: 50%;\n",
       "}\n",
       "\n",
       "#sk-container-id-1 div.sk-parallel-item:only-child::after {\n",
       "  width: 0;\n",
       "}\n",
       "\n",
       "/* Serial-specific style estimator block */\n",
       "\n",
       "#sk-container-id-1 div.sk-serial {\n",
       "  display: flex;\n",
       "  flex-direction: column;\n",
       "  align-items: center;\n",
       "  background-color: var(--sklearn-color-background);\n",
       "  padding-right: 1em;\n",
       "  padding-left: 1em;\n",
       "}\n",
       "\n",
       "\n",
       "/* Toggleable style: style used for estimator/Pipeline/ColumnTransformer box that is\n",
       "clickable and can be expanded/collapsed.\n",
       "- Pipeline and ColumnTransformer use this feature and define the default style\n",
       "- Estimators will overwrite some part of the style using the `sk-estimator` class\n",
       "*/\n",
       "\n",
       "/* Pipeline and ColumnTransformer style (default) */\n",
       "\n",
       "#sk-container-id-1 div.sk-toggleable {\n",
       "  /* Default theme specific background. It is overwritten whether we have a\n",
       "  specific estimator or a Pipeline/ColumnTransformer */\n",
       "  background-color: var(--sklearn-color-background);\n",
       "}\n",
       "\n",
       "/* Toggleable label */\n",
       "#sk-container-id-1 label.sk-toggleable__label {\n",
       "  cursor: pointer;\n",
       "  display: block;\n",
       "  width: 100%;\n",
       "  margin-bottom: 0;\n",
       "  padding: 0.5em;\n",
       "  box-sizing: border-box;\n",
       "  text-align: center;\n",
       "}\n",
       "\n",
       "#sk-container-id-1 label.sk-toggleable__label-arrow:before {\n",
       "  /* Arrow on the left of the label */\n",
       "  content: \"▸\";\n",
       "  float: left;\n",
       "  margin-right: 0.25em;\n",
       "  color: var(--sklearn-color-icon);\n",
       "}\n",
       "\n",
       "#sk-container-id-1 label.sk-toggleable__label-arrow:hover:before {\n",
       "  color: var(--sklearn-color-text);\n",
       "}\n",
       "\n",
       "/* Toggleable content - dropdown */\n",
       "\n",
       "#sk-container-id-1 div.sk-toggleable__content {\n",
       "  max-height: 0;\n",
       "  max-width: 0;\n",
       "  overflow: hidden;\n",
       "  text-align: left;\n",
       "  /* unfitted */\n",
       "  background-color: var(--sklearn-color-unfitted-level-0);\n",
       "}\n",
       "\n",
       "#sk-container-id-1 div.sk-toggleable__content.fitted {\n",
       "  /* fitted */\n",
       "  background-color: var(--sklearn-color-fitted-level-0);\n",
       "}\n",
       "\n",
       "#sk-container-id-1 div.sk-toggleable__content pre {\n",
       "  margin: 0.2em;\n",
       "  border-radius: 0.25em;\n",
       "  color: var(--sklearn-color-text);\n",
       "  /* unfitted */\n",
       "  background-color: var(--sklearn-color-unfitted-level-0);\n",
       "}\n",
       "\n",
       "#sk-container-id-1 div.sk-toggleable__content.fitted pre {\n",
       "  /* unfitted */\n",
       "  background-color: var(--sklearn-color-fitted-level-0);\n",
       "}\n",
       "\n",
       "#sk-container-id-1 input.sk-toggleable__control:checked~div.sk-toggleable__content {\n",
       "  /* Expand drop-down */\n",
       "  max-height: 200px;\n",
       "  max-width: 100%;\n",
       "  overflow: auto;\n",
       "}\n",
       "\n",
       "#sk-container-id-1 input.sk-toggleable__control:checked~label.sk-toggleable__label-arrow:before {\n",
       "  content: \"▾\";\n",
       "}\n",
       "\n",
       "/* Pipeline/ColumnTransformer-specific style */\n",
       "\n",
       "#sk-container-id-1 div.sk-label input.sk-toggleable__control:checked~label.sk-toggleable__label {\n",
       "  color: var(--sklearn-color-text);\n",
       "  background-color: var(--sklearn-color-unfitted-level-2);\n",
       "}\n",
       "\n",
       "#sk-container-id-1 div.sk-label.fitted input.sk-toggleable__control:checked~label.sk-toggleable__label {\n",
       "  background-color: var(--sklearn-color-fitted-level-2);\n",
       "}\n",
       "\n",
       "/* Estimator-specific style */\n",
       "\n",
       "/* Colorize estimator box */\n",
       "#sk-container-id-1 div.sk-estimator input.sk-toggleable__control:checked~label.sk-toggleable__label {\n",
       "  /* unfitted */\n",
       "  background-color: var(--sklearn-color-unfitted-level-2);\n",
       "}\n",
       "\n",
       "#sk-container-id-1 div.sk-estimator.fitted input.sk-toggleable__control:checked~label.sk-toggleable__label {\n",
       "  /* fitted */\n",
       "  background-color: var(--sklearn-color-fitted-level-2);\n",
       "}\n",
       "\n",
       "#sk-container-id-1 div.sk-label label.sk-toggleable__label,\n",
       "#sk-container-id-1 div.sk-label label {\n",
       "  /* The background is the default theme color */\n",
       "  color: var(--sklearn-color-text-on-default-background);\n",
       "}\n",
       "\n",
       "/* On hover, darken the color of the background */\n",
       "#sk-container-id-1 div.sk-label:hover label.sk-toggleable__label {\n",
       "  color: var(--sklearn-color-text);\n",
       "  background-color: var(--sklearn-color-unfitted-level-2);\n",
       "}\n",
       "\n",
       "/* Label box, darken color on hover, fitted */\n",
       "#sk-container-id-1 div.sk-label.fitted:hover label.sk-toggleable__label.fitted {\n",
       "  color: var(--sklearn-color-text);\n",
       "  background-color: var(--sklearn-color-fitted-level-2);\n",
       "}\n",
       "\n",
       "/* Estimator label */\n",
       "\n",
       "#sk-container-id-1 div.sk-label label {\n",
       "  font-family: monospace;\n",
       "  font-weight: bold;\n",
       "  display: inline-block;\n",
       "  line-height: 1.2em;\n",
       "}\n",
       "\n",
       "#sk-container-id-1 div.sk-label-container {\n",
       "  text-align: center;\n",
       "}\n",
       "\n",
       "/* Estimator-specific */\n",
       "#sk-container-id-1 div.sk-estimator {\n",
       "  font-family: monospace;\n",
       "  border: 1px dotted var(--sklearn-color-border-box);\n",
       "  border-radius: 0.25em;\n",
       "  box-sizing: border-box;\n",
       "  margin-bottom: 0.5em;\n",
       "  /* unfitted */\n",
       "  background-color: var(--sklearn-color-unfitted-level-0);\n",
       "}\n",
       "\n",
       "#sk-container-id-1 div.sk-estimator.fitted {\n",
       "  /* fitted */\n",
       "  background-color: var(--sklearn-color-fitted-level-0);\n",
       "}\n",
       "\n",
       "/* on hover */\n",
       "#sk-container-id-1 div.sk-estimator:hover {\n",
       "  /* unfitted */\n",
       "  background-color: var(--sklearn-color-unfitted-level-2);\n",
       "}\n",
       "\n",
       "#sk-container-id-1 div.sk-estimator.fitted:hover {\n",
       "  /* fitted */\n",
       "  background-color: var(--sklearn-color-fitted-level-2);\n",
       "}\n",
       "\n",
       "/* Specification for estimator info (e.g. \"i\" and \"?\") */\n",
       "\n",
       "/* Common style for \"i\" and \"?\" */\n",
       "\n",
       ".sk-estimator-doc-link,\n",
       "a:link.sk-estimator-doc-link,\n",
       "a:visited.sk-estimator-doc-link {\n",
       "  float: right;\n",
       "  font-size: smaller;\n",
       "  line-height: 1em;\n",
       "  font-family: monospace;\n",
       "  background-color: var(--sklearn-color-background);\n",
       "  border-radius: 1em;\n",
       "  height: 1em;\n",
       "  width: 1em;\n",
       "  text-decoration: none !important;\n",
       "  margin-left: 1ex;\n",
       "  /* unfitted */\n",
       "  border: var(--sklearn-color-unfitted-level-1) 1pt solid;\n",
       "  color: var(--sklearn-color-unfitted-level-1);\n",
       "}\n",
       "\n",
       ".sk-estimator-doc-link.fitted,\n",
       "a:link.sk-estimator-doc-link.fitted,\n",
       "a:visited.sk-estimator-doc-link.fitted {\n",
       "  /* fitted */\n",
       "  border: var(--sklearn-color-fitted-level-1) 1pt solid;\n",
       "  color: var(--sklearn-color-fitted-level-1);\n",
       "}\n",
       "\n",
       "/* On hover */\n",
       "div.sk-estimator:hover .sk-estimator-doc-link:hover,\n",
       ".sk-estimator-doc-link:hover,\n",
       "div.sk-label-container:hover .sk-estimator-doc-link:hover,\n",
       ".sk-estimator-doc-link:hover {\n",
       "  /* unfitted */\n",
       "  background-color: var(--sklearn-color-unfitted-level-3);\n",
       "  color: var(--sklearn-color-background);\n",
       "  text-decoration: none;\n",
       "}\n",
       "\n",
       "div.sk-estimator.fitted:hover .sk-estimator-doc-link.fitted:hover,\n",
       ".sk-estimator-doc-link.fitted:hover,\n",
       "div.sk-label-container:hover .sk-estimator-doc-link.fitted:hover,\n",
       ".sk-estimator-doc-link.fitted:hover {\n",
       "  /* fitted */\n",
       "  background-color: var(--sklearn-color-fitted-level-3);\n",
       "  color: var(--sklearn-color-background);\n",
       "  text-decoration: none;\n",
       "}\n",
       "\n",
       "/* Span, style for the box shown on hovering the info icon */\n",
       ".sk-estimator-doc-link span {\n",
       "  display: none;\n",
       "  z-index: 9999;\n",
       "  position: relative;\n",
       "  font-weight: normal;\n",
       "  right: .2ex;\n",
       "  padding: .5ex;\n",
       "  margin: .5ex;\n",
       "  width: min-content;\n",
       "  min-width: 20ex;\n",
       "  max-width: 50ex;\n",
       "  color: var(--sklearn-color-text);\n",
       "  box-shadow: 2pt 2pt 4pt #999;\n",
       "  /* unfitted */\n",
       "  background: var(--sklearn-color-unfitted-level-0);\n",
       "  border: .5pt solid var(--sklearn-color-unfitted-level-3);\n",
       "}\n",
       "\n",
       ".sk-estimator-doc-link.fitted span {\n",
       "  /* fitted */\n",
       "  background: var(--sklearn-color-fitted-level-0);\n",
       "  border: var(--sklearn-color-fitted-level-3);\n",
       "}\n",
       "\n",
       ".sk-estimator-doc-link:hover span {\n",
       "  display: block;\n",
       "}\n",
       "\n",
       "/* \"?\"-specific style due to the `<a>` HTML tag */\n",
       "\n",
       "#sk-container-id-1 a.estimator_doc_link {\n",
       "  float: right;\n",
       "  font-size: 1rem;\n",
       "  line-height: 1em;\n",
       "  font-family: monospace;\n",
       "  background-color: var(--sklearn-color-background);\n",
       "  border-radius: 1rem;\n",
       "  height: 1rem;\n",
       "  width: 1rem;\n",
       "  text-decoration: none;\n",
       "  /* unfitted */\n",
       "  color: var(--sklearn-color-unfitted-level-1);\n",
       "  border: var(--sklearn-color-unfitted-level-1) 1pt solid;\n",
       "}\n",
       "\n",
       "#sk-container-id-1 a.estimator_doc_link.fitted {\n",
       "  /* fitted */\n",
       "  border: var(--sklearn-color-fitted-level-1) 1pt solid;\n",
       "  color: var(--sklearn-color-fitted-level-1);\n",
       "}\n",
       "\n",
       "/* On hover */\n",
       "#sk-container-id-1 a.estimator_doc_link:hover {\n",
       "  /* unfitted */\n",
       "  background-color: var(--sklearn-color-unfitted-level-3);\n",
       "  color: var(--sklearn-color-background);\n",
       "  text-decoration: none;\n",
       "}\n",
       "\n",
       "#sk-container-id-1 a.estimator_doc_link.fitted:hover {\n",
       "  /* fitted */\n",
       "  background-color: var(--sklearn-color-fitted-level-3);\n",
       "}\n",
       "</style><div id=\"sk-container-id-1\" class=\"sk-top-container\"><div class=\"sk-text-repr-fallback\"><pre>XGBRegressor(base_score=None, booster=None, callbacks=None,\n",
       "             colsample_bylevel=None, colsample_bynode=None,\n",
       "             colsample_bytree=None, device=&#x27;cuda&#x27;,\n",
       "             disable_default_eval_metric=True, early_stopping_rounds=2,\n",
       "             enable_categorical=False,\n",
       "             eval_metric=&lt;function r2_xgb at 0x7edc8c5b2ac0&gt;,\n",
       "             feature_types=None, gamma=None, grow_policy=None,\n",
       "             importance_type=None, interaction_constraints=None,\n",
       "             learning_rate=0.1, max_bin=None, max_cat_threshold=None,\n",
       "             max_cat_to_onehot=None, max_delta_step=None, max_depth=6,\n",
       "             max_leaves=None, min_child_weight=None, missing=nan,\n",
       "             monotone_constraints=None, multi_strategy=None, n_estimators=2000,\n",
       "             n_jobs=None, num_parallel_tree=None, ...)</pre><b>In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook. <br />On GitHub, the HTML representation is unable to render, please try loading this page with nbviewer.org.</b></div><div class=\"sk-container\" hidden><div class=\"sk-item\"><div class=\"sk-estimator fitted sk-toggleable\"><input class=\"sk-toggleable__control sk-hidden--visually\" id=\"sk-estimator-id-1\" type=\"checkbox\" checked><label for=\"sk-estimator-id-1\" class=\"sk-toggleable__label fitted sk-toggleable__label-arrow fitted\">&nbsp;XGBRegressor<span class=\"sk-estimator-doc-link fitted\">i<span>Fitted</span></span></label><div class=\"sk-toggleable__content fitted\"><pre>XGBRegressor(base_score=None, booster=None, callbacks=None,\n",
       "             colsample_bylevel=None, colsample_bynode=None,\n",
       "             colsample_bytree=None, device=&#x27;cuda&#x27;,\n",
       "             disable_default_eval_metric=True, early_stopping_rounds=2,\n",
       "             enable_categorical=False,\n",
       "             eval_metric=&lt;function r2_xgb at 0x7edc8c5b2ac0&gt;,\n",
       "             feature_types=None, gamma=None, grow_policy=None,\n",
       "             importance_type=None, interaction_constraints=None,\n",
       "             learning_rate=0.1, max_bin=None, max_cat_threshold=None,\n",
       "             max_cat_to_onehot=None, max_delta_step=None, max_depth=6,\n",
       "             max_leaves=None, min_child_weight=None, missing=nan,\n",
       "             monotone_constraints=None, multi_strategy=None, n_estimators=2000,\n",
       "             n_jobs=None, num_parallel_tree=None, ...)</pre></div> </div></div></div></div>"
      ],
      "text/plain": [
       "XGBRegressor(base_score=None, booster=None, callbacks=None,\n",
       "             colsample_bylevel=None, colsample_bynode=None,\n",
       "             colsample_bytree=None, device='cuda',\n",
       "             disable_default_eval_metric=True, early_stopping_rounds=2,\n",
       "             enable_categorical=False,\n",
       "             eval_metric=<function r2_xgb at 0x7edc8c5b2ac0>,\n",
       "             feature_types=None, gamma=None, grow_policy=None,\n",
       "             importance_type=None, interaction_constraints=None,\n",
       "             learning_rate=0.1, max_bin=None, max_cat_threshold=None,\n",
       "             max_cat_to_onehot=None, max_delta_step=None, max_depth=6,\n",
       "             max_leaves=None, min_child_weight=None, missing=nan,\n",
       "             monotone_constraints=None, multi_strategy=None, n_estimators=2000,\n",
       "             n_jobs=None, num_parallel_tree=None, ...)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.fit(\n",
    "    X_train,\n",
    "    y_train,\n",
    "    sample_weight=w_train,\n",
    "    eval_set=[(X_valid, y_valid)],\n",
    "    sample_weight_eval_set=[w_valid],\n",
    "    verbose=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# save model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(\"./model_save\"):\n",
    "    # Create the directory if it does not exist\n",
    "    os.mkdir(\"./model_save\")\n",
    "model.save_model(\"./model_save/js_xgboost_model_baseline.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# load model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_loaded = xgb.XGBRegressor()\n",
    "model_loaded.load_model(\"./model_save/js_xgboost_model_baseline.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# load test data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "test = pl.scan_parquet(\n",
    "    \"/home/pqian/data/kaggle/js-2024/test.parquet/date_id=0/part-0.parquet\"\n",
    ")\n",
    "test = test.collect()\n",
    "test = test.to_pandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_df = test[feature_names].values\n",
    "predictions = model_loaded.predict(test_df)\n",
    "\n",
    "output_df = pd.DataFrame({\"row_id\": test[\"row_id\"], \"responder_6\": predictions})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# submission API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "lags_: pl.DataFrame | None = None\n",
    "\n",
    "\n",
    "# Replace this function with your inference code.\n",
    "# You can return either a Pandas or Polars dataframe, though Polars is recommended.\n",
    "# Each batch of predictions (except the very first) must be returned within 1 minute of the batch features being provided.\n",
    "def predict(\n",
    "    test: pl.DataFrame, lags: pl.DataFrame | None\n",
    ") -> pl.DataFrame | pd.DataFrame:\n",
    "    \"\"\"Make a prediction.\"\"\"\n",
    "    # All the responders from the previous day are passed in at time_id == 0. We save them in a global variable for access at every time_id.\n",
    "    # Use them as extra features, if you like.\n",
    "    global lags_\n",
    "    if lags is not None:\n",
    "        lags_ = lags\n",
    "    # Replace this section with your own predictions\n",
    "    predictions = test.select(\n",
    "        \"row_id\",\n",
    "        pl.lit(0.0).alias(\"responder_6\"),\n",
    "    )\n",
    "\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",
    "    # Confirm has as many rows as the test data.\n",
    "    assert len(predictions) == len(test)\n",
    "\n",
    "    return predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "inference_server = kaggle_evaluation.jane_street_inference_server.JSInferenceServer(\n",
    "    predict\n",
    ")\n",
    "\n",
    "if os.getenv(\"KAGGLE_IS_COMPETITION_RERUN\"):\n",
    "    inference_server.serve()\n",
    "else:\n",
    "    inference_server.run_local_gateway(\n",
    "        (\n",
    "            \"/home/pqian/data/kaggle/js-2024/test.parquet\",\n",
    "            \"/home/pqian/data/kaggle/js-2024/lags.parquet\",\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "kaggle",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}