{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":84493,"databundleVersionId":9871156,"sourceType":"competition"},{"sourceId":145414,"sourceType":"modelInstanceVersion","modelInstanceId":123281,"modelId":146350}],"dockerImageVersionId":30787,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Import necessary libraries\nimport numpy as np\nimport polars as pl\nimport pandas as pd\nimport lightgbm as lgb\nimport xgboost as xgb\nimport os\nimport joblib\nimport kaggle_evaluation.jane_street_inference_server\n\n# Set up constants\nTARGET = 'responder_6'\nFEAT_COLS = [f\"feature_{i:02d}\" for i in range(79)]\n\n# Function to load data with optional filtering\ndef load_data(date_id_range=None, time_id_range=None, columns=None, return_type='pl'):\n    data_dir = '../input/jane-street-real-time-market-data-forecasting'\n    data = pl.scan_parquet(f\"{data_dir}/train.parquet\")\n\n    if date_id_range is not None:\n        start_date, end_date = date_id_range\n        data = data.filter((pl.col(\"date_id\") >= start_date) & (pl.col(\"date_id\") <= end_date))\n\n    if time_id_range is not None:\n        start_time, end_time = time_id_range\n        data = data.filter((pl.col(\"time_id\") >= start_time) & (pl.col(\"time_id\") <= end_time))\n\n    if columns is not None:\n        data = data.select(columns)\n\n    if return_type == 'pd':\n        return data.collect().to_pandas()\n    else:\n        return data.collect()\n\n# Function to calculate R² score\ndef calculate_r2(y_true, y_pred, weights):\n    numerator = np.sum(weights * (y_true - y_pred) ** 2)\n    denominator = np.sum(weights * (y_true ** 2))\n    r2_score = 1 - (numerator / denominator)\n    return r2_score\n\n# Function to evaluate the model\ndef evaluate_model(model, test_data):\n    y_pred = model.predict(test_data[FEAT_COLS])\n    y_true = test_data[TARGET].to_numpy() \n    weights = test_data['weight'].to_numpy()  \n    r2_score = calculate_r2(y_true, y_pred, weights)\n    print(f\"Sample weighted zero-mean R-squared score (R2) on test data: {r2_score}\")\n\n# Class to manage a group of models\nclass ModelGroup:\n    def __init__(self):\n        self.models = []\n\n    def add_model(self, model):\n        self.models.append(model)\n\n    def predict(self, test_data):\n        preds = []\n        for model in self.models:\n            if isinstance(model, lgb.Booster):\n                pred = model.predict(test_data[FEAT_COLS])\n            elif isinstance(model, xgb.Booster):\n                pred = model.predict(xgb.DMatrix(test_data[FEAT_COLS]))\n            elif hasattr(model, 'predict'):\n                pred = model.predict(test_data[FEAT_COLS])\n            else:\n                raise ValueError(\"Unsupported model type\")\n            preds.append(pred)\n\n        avg_pred = np.mean(preds, axis=0)\n        return avg_pred\n    \n    @classmethod\n    def load(cls, file_path):\n        model_group = joblib.load(file_path)\n        return model_group\n\n# Function to train XGBoost with K-Folds\ndef train_xgb_kfold(total_days=1498, n_splits=5, save_models=False):\n    model_group = ModelGroup()\n    fold_size = total_days // n_splits\n    folds = [(i * fold_size, min((i + 1) * fold_size - 1, total_days - 1)) for i in range(n_splits)]\n\n    for fold_idx in range(n_splits):\n        valid_range = folds[fold_idx]\n        train_ranges = [folds[i] for i in range(n_splits) if i != fold_idx]\n\n        print(f\"Fold {fold_idx}: validation range {valid_range}, train parts: {train_ranges}\")\n\n        valid_data = load_data(date_id_range=valid_range, columns=[\"date_id\", \"weight\"] + FEAT_COLS + [TARGET], return_type='pl')\n        valid_weight = valid_data['weight'].to_pandas()\n\n        train_data = None\n        for train_range in train_ranges:\n            partial_train_data = load_data(date_id_range=train_range, columns=[\"date_id\", \"weight\"] + FEAT_COLS + [TARGET], return_type='pl')\n            if train_data is None:\n                train_data = partial_train_data\n            else:\n                train_data = train_data.vstack(partial_train_data)\n\n        train_weight = train_data['weight'].to_pandas()\n\n        dtrain = xgb.DMatrix(train_data.select(FEAT_COLS).to_pandas(), label=train_data[TARGET].to_pandas(), weight=train_weight)\n        dvalid = xgb.DMatrix(valid_data.select(FEAT_COLS).to_pandas(), label=valid_data[TARGET].to_pandas(), weight=valid_weight)\n\n        XGB_PARAMS = {\n            'eval_metric': 'rmse',\n            'learning_rate': 0.5,\n            'max_depth': 12,\n            'min_child_weight': 1.5,\n            'subsample': 0.8555,\n            'colsample_bytree': 0.85555555,\n            'random_state': 42,\n            'tree_method': 'gpu_hist',\n        }\n\n        model = xgb.train(XGB_PARAMS, dtrain, num_boost_round=1000, evals=[(dtrain, 'train'), (dvalid, 'valid')], early_stopping_rounds=100, verbose_eval=50)\n\n        y_valid_pred = model.predict(dvalid)\n        r2_score = calculate_r2(valid_data[TARGET].to_pandas(), y_valid_pred, valid_weight)\n        print(f\"Fold {fold_idx} validation R2 score: {r2_score}\")\n\n        model_group.add_model(model)\n\n    if save_models:\n        joblib.dump(model_group, \"xgb_model_group.pkl\")\n        print(\"Saved the model group to xgb_model_group.pkl\")\n    \n    return model_group\n\n# Uncomment to train a new model\n# total_days = 1699\n# xgb_models = train_xgb_kfold(total_days=total_days, n_splits=5, save_models=False)\n\n# Load pre-trained model group\nxgb_models = ModelGroup.load(\"/kaggle/input/xgb_model/other/default/1/xgb_model_group.pkl\")\nlags_ = None\n\n# Prediction function for the inference server\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    feat = test[FEAT_COLS].to_pandas()\n    pred = xgb_models.predict(feat)\n\n    predictions = test.select('row_id').with_columns(pl.Series('responder_6', pred.ravel()))\n\n    assert isinstance(predictions, (pl.DataFrame, pd.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        (\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    )\n","metadata":{"execution":{"iopub.status.busy":"2024-10-25T17:54:05.124233Z","iopub.execute_input":"2024-10-25T17:54:05.125163Z","iopub.status.idle":"2024-10-25T17:54:12.308784Z","shell.execute_reply.started":"2024-10-25T17:54:05.125112Z","shell.execute_reply":"2024-10-25T17:54:12.307725Z"},"trusted":true},"outputs":[],"execution_count":null}]}