{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10.12","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":10139918,"sourceType":"datasetVersion","datasetId":6258261},{"sourceId":10139922,"sourceType":"datasetVersion","datasetId":6258265},{"sourceId":10290183,"sourceType":"datasetVersion","datasetId":6368446},{"sourceId":203900450,"sourceType":"kernelVersion"},{"sourceId":213144305,"sourceType":"kernelVersion"},{"sourceId":214286693,"sourceType":"kernelVersion"}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":80.344101,"end_time":"2024-10-26T03:27:42.952247","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2024-10-26T03:26:22.608146","version":"2.6.0"},"widgets":{"application/vnd.jupyter.widget-state+json":{"state":{"252dd2de87de42f4becd877fbbafc26b":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"HTMLModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_ff55584ae0ab4f39ae471f314eaad988","placeholder":"​","style":"IPY_MODEL_b8d338c473aa4dc3ba25f37a997a9037","value":" 1/1 [00:00&lt;00:00, 33.79it/s]"}},"48a2731fb59b4ce8ace5b53d6f0e3337":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"HTMLModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HTMLModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HTMLView","description":"","description_tooltip":null,"layout":"IPY_MODEL_a8dbc79c7c5a48318466a102c55801bb","placeholder":"​","style":"IPY_MODEL_ebd06aaaa7024a1684ba1b9fe89358cf","value":"100%"}},"56da652f3eeb42aca986e6fd815629dc":{"model_module":"@jupyter-widgets/base","model_module_version":"1.2.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"5bb4e0df01c44716afebab25aafe9f5f":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"FloatProgressModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"FloatProgressModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"ProgressView","bar_style":"success","description":"","description_tooltip":null,"layout":"IPY_MODEL_f4e41234b0cc4f3e9313d971f5aadc9f","max":1,"min":0,"orientation":"horizontal","style":"IPY_MODEL_6408698650f74dd699bcc914b850e396","value":1}},"6408698650f74dd699bcc914b850e396":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"ProgressStyleModel","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"ProgressStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","bar_color":null,"description_width":""}},"9cb493ec04fc4ed391f5ac28cc84500e":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"HBoxModel","state":{"_dom_classes":[],"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"HBoxModel","_view_count":null,"_view_module":"@jupyter-widgets/controls","_view_module_version":"1.5.0","_view_name":"HBoxView","box_style":"","children":["IPY_MODEL_48a2731fb59b4ce8ace5b53d6f0e3337","IPY_MODEL_5bb4e0df01c44716afebab25aafe9f5f","IPY_MODEL_252dd2de87de42f4becd877fbbafc26b"],"layout":"IPY_MODEL_56da652f3eeb42aca986e6fd815629dc"}},"a8dbc79c7c5a48318466a102c55801bb":{"model_module":"@jupyter-widgets/base","model_module_version":"1.2.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"b8d338c473aa4dc3ba25f37a997a9037":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"DescriptionStyleModel","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"ebd06aaaa7024a1684ba1b9fe89358cf":{"model_module":"@jupyter-widgets/controls","model_module_version":"1.5.0","model_name":"DescriptionStyleModel","state":{"_model_module":"@jupyter-widgets/controls","_model_module_version":"1.5.0","_model_name":"DescriptionStyleModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"StyleView","description_width":""}},"f4e41234b0cc4f3e9313d971f5aadc9f":{"model_module":"@jupyter-widgets/base","model_module_version":"1.2.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}},"ff55584ae0ab4f39ae471f314eaad988":{"model_module":"@jupyter-widgets/base","model_module_version":"1.2.0","model_name":"LayoutModel","state":{"_model_module":"@jupyter-widgets/base","_model_module_version":"1.2.0","_model_name":"LayoutModel","_view_count":null,"_view_module":"@jupyter-widgets/base","_view_module_version":"1.2.0","_view_name":"LayoutView","align_content":null,"align_items":null,"align_self":null,"border":null,"bottom":null,"display":null,"flex":null,"flex_flow":null,"grid_area":null,"grid_auto_columns":null,"grid_auto_flow":null,"grid_auto_rows":null,"grid_column":null,"grid_gap":null,"grid_row":null,"grid_template_areas":null,"grid_template_columns":null,"grid_template_rows":null,"height":null,"justify_content":null,"justify_items":null,"left":null,"margin":null,"max_height":null,"max_width":null,"min_height":null,"min_width":null,"object_fit":null,"object_position":null,"order":null,"overflow":null,"overflow_x":null,"overflow_y":null,"padding":null,"right":null,"top":null,"visibility":null,"width":null}}},"version_major":2,"version_minor":0}}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Libraries","metadata":{}},{"cell_type":"code","source":"import polars as pl\n!pip install rtdl_num_embeddings --no-index --find-links=/kaggle/input/jane-street-import/rtdl_num_embeddings","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-02T13:09:37.455748Z","iopub.execute_input":"2025-01-02T13:09:37.456006Z","iopub.status.idle":"2025-01-02T13:09:42.194106Z","shell.execute_reply.started":"2025-01-02T13:09:37.455984Z","shell.execute_reply":"2025-01-02T13:09:42.193031Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, sys, gc\nimport enum\nimport datetime\nimport pickle\nimport dill\nimport numpy as np\nimport pandas as pd\nimport polars as pl\n\nfrom sklearn.metrics import r2_score\n\nimport lightgbm as lgb\nfrom lightgbm import LGBMRegressor, Booster\nfrom xgboost import XGBRegressor\nfrom catboost import CatBoostRegressor\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader, Dataset\nfrom pytorch_lightning import LightningModule\n\nsys.path.append(\"/kaggle/input/jane-street-real-time-market-data-forecasting\")\nimport kaggle_evaluation.jane_street_inference_server\n\nsys.path.append(\"/kaggle/input/src/tabm_reference\")\nfrom tanm_reference import Model, make_parameter_groups\n\nimport warnings\nwarnings.filterwarnings('ignore')\npd.options.display.max_columns = None\n\n@enum.unique\nclass DataEnum(enum.IntEnum):\n    Train = 0\n    Valid = 1\n    Test = 2\n    Infer = 3\n\nis_debug = False\nis_rerun = os.environ.get('KAGGLE_IS_COMPETITION_RERUN', \"\") != \"\" \nis_local = os.environ.get(\"DOCKER_USING\", \"\") == \"LOCAL\"\nnum_workers = 4\nif is_rerun:\n    is_debug = False\n\ndef load_from_dill(model_name, model_path=None, file_ext='.dill'):\n    model_object = None\n    with open(f\"{model_path}/{model_name}{file_ext}\", \"rb\") as file_handle:\n        model_object = dill.load(file_handle)\n    return model_object\n\n# RMF-NN, Starter, TabM, Ridge\nensem_koef = dict(\n    nn = 0.55,\n    starter = 0.35,\n    tabm = 0.3,\n    ridge = 0.10,\n)\nsum_ensem_koef = 0.0\nfor value in ensem_koef.values():\n    sum_ensem_koef += value \nfor key, value in ensem_koef.items():\n    ensem_koef[key] = value / sum_ensem_koef\nprint(ensem_koef)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-02T13:09:42.195342Z","iopub.execute_input":"2025-01-02T13:09:42.196115Z","iopub.status.idle":"2025-01-02T13:09:52.506856Z","shell.execute_reply.started":"2025-01-02T13:09:42.195794Z","shell.execute_reply":"2025-01-02T13:09:52.505975Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# [yunsuxiaozi - JS2024 Starter v10 LB=0.0070 (0.0054)](https://www.kaggle.com/code/yunsuxiaozi/js2024-starter?scriptVersionId=206770572)","metadata":{"papermill":{"duration":0.022173,"end_time":"2024-10-26T03:27:39.285949","exception":false,"start_time":"2024-10-26T03:27:39.263776","status":"completed"},"tags":[]}},{"cell_type":"code","source":"model_path = \"/kaggle/input/js-2024/13-01-11\" if is_local else \"/kaggle/input/js-2024-13-01-11\"\nfinal_feature = ['symbol_id','sin_time_id','cos_time_id','sin_time_id_halfday','cos_time_id_halfday'] + [f'feature_{i:02}' for i in range(79)]\n\nlgb = Booster(model_file=f\"{model_path}/lgb.model\")\ncat = CatBoostRegressor()\ncat.load_model(f\"{model_path}/cat.model\")\nxgb = XGBRegressor()\nxgb.load_model(f\"{model_path}/xgb.model\")\n\ndef predict_starter(test,lags):\n    predictions = test.select(\n        'row_id',\n        pl.lit(0.0).alias('responder_6'),\n    )\n    test=test.to_pandas()\n    test['sin_time_id']=np.sin(2*np.pi*test['time_id']/967)\n    test['cos_time_id']=np.cos(2*np.pi*test['time_id']/967)\n    test['sin_time_id_halfday']=np.sin(2*np.pi*test['time_id']/483)\n    test['cos_time_id_halfday']=np.cos(2*np.pi*test['time_id']/483)\n    test=test.fillna(-1)\n    test=test[final_feature]\n    eps=1e-10\n    test_preds=0.55*lgb.predict(test)+0.2*cat.predict(test)+0.25*xgb.predict(test)\n    test_preds=np.clip(test_preds,-5+eps,5-eps)\n    predictions = predictions.with_columns(pl.Series('responder_6', test_preds.ravel()))\n    return predictions","metadata":{"_kg_hide-input":true,"papermill":{"duration":0.031534,"end_time":"2024-10-26T03:27:39.453592","exception":false,"start_time":"2024-10-26T03:27:39.422058","status":"completed"},"tags":[],"_kg_hide-output":true,"trusted":true,"execution":{"iopub.status.busy":"2025-01-02T13:09:52.507802Z","iopub.execute_input":"2025-01-02T13:09:52.508486Z","iopub.status.idle":"2025-01-02T13:09:52.826420Z","shell.execute_reply.started":"2025-01-02T13:09:52.508453Z","shell.execute_reply":"2025-01-02T13:09:52.825523Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# [I2nfinit3y - Jane Street | TabM/FT-Transformer inference LB=0.0064 (0.0074)](https://www.kaggle.com/code/i2nfinit3y/jane-street-tabm-ft-transformer-inference?scriptVersionId=213715783)","metadata":{}},{"cell_type":"code","source":"model_path = '/kaggle/input/js-2024' + ('/' if is_local else '-') + '19-02-1/last_tabm.pt'\nstats_path = '/kaggle/input/js-2024' + ('/' if is_local else '-') + '19-02-1/data_stats.dill'\ndevice = torch.device('cuda:0')\n\ntarget_col = \"responder_6\"\nnecessary_cols = [target_col, 'weight']\nfeat_clear_categ = [\"feature_09\", \"feature_10\", \"feature_11\"]\nfeature_categ = feat_clear_categ + ['symbol_id', 'time_id']\nfeature_cols = [f\"feature_{idx:02d}\" for idx in range(79) if idx not in [9, 10, 11, 61]]\nresponder_cols = [f\"responder_{idx}_lag_1\" for idx in range(9)] \nfeature_cont = feature_cols + responder_cols\ndataset_cols = feature_cont + necessary_cols + feature_categ\nstd_feature = [i for i in feature_cont]\n\nbatch_size = 8192\nn_cont_features = len(feature_cont)\nn_cat_features = len(feature_categ)\nn_classes = None\ncat_cardinalities = [23, 10, 32, 40, 969]\n# TabM\narch_type = 'tabm'\nbins = None\nmodel_koef = 32\n\nprint(n_cont_features, n_cat_features, len(dataset_cols))\n\ncategory_mappings = {\n    'feature_09': {2: 0, 4: 1, 9: 2, 11: 3, 12: 4, 14: 5, 15: 6, 25: 7, 26: 8, 30: 9, \n        34: 10, 42: 11, 44: 12, 46: 13, 49: 14, 50: 15, 57: 16, 64: 17, 68: 18, 70: 19, 81: 20, 82: 21},\n    'feature_10': {1: 0, 2: 1, 3: 2, 4: 3, 5: 4, 6: 5, 7: 6, 10: 7, 12: 8},\n    'feature_11': {9: 0, 11: 1, 13: 2, 16: 3, 24: 4, 25: 5, 34: 6, 40: 7, 48: 8, 50: 9, 59: 10, 62: 11, 63: 12, 66: 13,\n        76: 14, 150: 15, 158: 16, 159: 17, 171: 18, 195: 19, 214: 20, 230: 21, 261: 22, 297: 23, 336: 24, 376: 25, 388: 26, 410: 27, 522: 28, 534: 29, 539: 30},\n    'symbol_id': {i : i for i in range(39)},\n    'time_id' : {i : i for i in range(968)}\n}\n\ndef standardize(df, feature_cols, means, stds):\n    return df.with_columns([\n        ((pl.col(col) - means[col]) / stds[col]).alias(col) for col in feature_cols\n    ])\n\ndef encode_column(df, column, mapping):\n    max_value = max(mapping.values())\n    def encode_category(category):\n        return mapping.get(category, max_value + 1)\n    return df.with_columns(pl.col(column).map_elements(encode_category, return_dtype=pl.Int64).alias(column))\n\nmodel_tabm = Model(\n    n_num_features=n_cont_features,\n    cat_cardinalities=cat_cardinalities,\n    n_classes=n_classes,\n    backbone={\n        'type': 'MLP',\n        'n_blocks': 3 ,\n        'd_block': 512,\n        'dropout': 0.25,\n    },\n    bins=bins,\n    num_embeddings=(\n        None\n        # {\n        #     'type': 'PeriodicEmbeddings',\n        #     'd_embedding': 16,\n        #     'lite':True,\n        # }\n    ),\n    arch_type=arch_type,\n    k=model_koef,\n).to(device)\n\nwith open(stats_path, \"rb\") as file_handle:\n    data_stats = dill.load(file_handle)\nmeans, stds = data_stats['means'], data_stats['stds']\n\ncheckpoint = torch.load(model_path, weights_only=True)\nmodel_tabm.load_state_dict(checkpoint['model_state_dict'])\nmodel_tabm.to(device)\n\nlags_history = None\n\ndef predict_tabm(test: pl.DataFrame, lags: pl.DataFrame | None) -> pl.DataFrame:\n    global lags_history\n    for col in feature_categ:\n        test = encode_column(test, col, category_mappings[col])\n    predictions = test.select(\n        'row_id',\n        pl.lit(0.0).alias('responder_6'),\n    )\n    time_id = test.select(\"time_id\").to_numpy()[0]\n    if time_id == 0:\n        lags = lags.with_columns(pl.col('time_id').cast(pl.Int64))\n        lags = lags.with_columns(pl.col('symbol_id').cast(pl.Int64))    \n        lags_history = lags\n    lags = lags_history.clone().group_by([\"date_id\", \"symbol_id\"], maintain_order=True).last() # pick up last record of previous date\n    test = test.join(lags, on=[\"time_id\", \"symbol_id\"], how=\"left\")\n    \n    test = test.select(pl.all().forward_fill())\n    test = test.with_columns([\n        pl.col(col).fill_null(0) for col in feature_cont + feat_clear_categ\n    ])\n    test = standardize(test, std_feature, means, stds)\n\n    model_tabm.eval()\n    with torch.no_grad():\n        X_cont = torch.tensor(test[feature_cont].to_numpy(), dtype=torch.float32).to(device)\n        X_categ = torch.tensor(test[feature_categ].to_numpy(), dtype=torch.int64).to(device)\n        outputs = model_tabm(X_cont, X_categ)\n        # Assuming the model outputs a tensor of shape (batch_size, 1)\n        preds = outputs.squeeze(-1).cpu().numpy()\n        preds = preds.mean(1)\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)\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    return predictions","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-02T13:09:52.828009Z","iopub.execute_input":"2025-01-02T13:09:52.828282Z","iopub.status.idle":"2025-01-02T13:09:53.250753Z","shell.execute_reply.started":"2025-01-02T13:09:52.828260Z","shell.execute_reply":"2025-01-02T13:09:53.250004Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# [yunsuxiaozi - JS Ridge baseline v03 LB=0.0033](https://www.kaggle.com/code/yunsuxiaozi/js-ridge-baseline?scriptVersionId=202739388)","metadata":{}},{"cell_type":"code","source":"rdg = load_from_dill(model_name='ridge', model_path=\"/kaggle/input/js-2024/16-01\" if is_local else \"/kaggle/input/js-2024-16-01\")\n\ndef predict_ridge(test, lags):\n    cols = [f'feature_{i:02}' for i in range(79)]\n    predictions = test.select(\n        'row_id',\n        pl.lit(0.0).alias('responder_6'),\n    )\n    test_preds = rdg.predict(test[cols].to_pandas().fillna(3).values)\n    predictions = predictions.with_columns(pl.Series('responder_6', test_preds.ravel()))\n    return predictions","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-02T13:09:53.251719Z","iopub.execute_input":"2025-01-02T13:09:53.252023Z","iopub.status.idle":"2025-01-02T13:09:53.259448Z","shell.execute_reply.started":"2025-01-02T13:09:53.251991Z","shell.execute_reply":"2025-01-02T13:09:53.258791Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# [Xiang Sheng - Jane Street RMF NN + XGB Lb=0.0076 (0.0056)](https://www.kaggle.com/code/voix97/jane-street-rmf-nn-xgb)","metadata":{}},{"cell_type":"code","source":"class 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\nvalid = pl.scan_parquet(f\"/kaggle/input/js24-preprocessing-create-lags/validation.parquet/\").collect().to_pandas()\n\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\"]\nxgb_feature_cols = [\"symbol_id\", \"time_id\"] + CONFIG.feature_cols\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            # 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, 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}\")\n\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\nX_valid = valid[ xgb_feature_cols ]\ny_valid = valid[ CONFIG.target_col ]\nw_valid = valid[ \"weight\" ]\ny_pred_valid_xgb = xgb_model.predict(X_valid)\nvalid_score = r2_score( y_valid, y_pred_valid_xgb, sample_weight=w_valid )\nvalid_score\n\nX_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\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(torch.FloatTensor(X_valid.values).to(\"cuda:0\")).cpu().numpy() / len(models)\nvalid_score = r2_score( y_valid, y_pred_valid_nn, sample_weight=w_valid )\nvalid_score\n\ny_pred_valid_ensemble = 0.5 * (y_pred_valid_xgb + y_pred_valid_nn)\nvalid_score = r2_score( y_valid, y_pred_valid_ensemble, sample_weight=w_valid )\nvalid_score\n\ndel valid, X_valid, y_valid, w_valid\ngc.collect()\n\nlags_ : pl.DataFrame | None = None\n\ndef predict_nn(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_14 = test.select('row_id', pl.lit(0.0).alias('responder_6'),)\n\n    # add this part to reuse lags of previous date ids when rows have more than 0 time_ids.\n    lags = lags_.clone().group_by([\"date_id\", \"symbol_id\"], maintain_order=True).last()\n    test = test.join(lags, on=[\"date_id\", \"symbol_id\"],  how=\"left\")\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    #else:\n    #    test = test.with_columns(( pl.lit(0.0).alias(f'responder_{idx}_lag_1') for idx in range(9) ))\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    with torch.no_grad():\n        for i, nn_model in enumerate(models):\n            nn_model.eval()\n            preds += nn_model(test_input).cpu().numpy() / 10\n    #print(f\"predict> preds.shape =\", preds.shape)\n\n    predictions_14 = 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    return predictions_14\n","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true,"execution":{"iopub.status.busy":"2025-01-02T13:09:53.260239Z","iopub.execute_input":"2025-01-02T13:09:53.260513Z","iopub.status.idle":"2025-01-02T13:10:06.530598Z","shell.execute_reply.started":"2025-01-02T13:09:53.260493Z","shell.execute_reply":"2025-01-02T13:10:06.529884Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Ensemble","metadata":{}},{"cell_type":"code","source":"is_first = False\n\ndef predict(test:pl.DataFrame, lags:pl.DataFrame | None) -> pl.DataFrame | pd.DataFrame:\n    global is_first\n\n    pd_nn = predict_nn(test,lags).to_pandas()\n    pd_starter = predict_starter(test,lags).to_pandas()\n    pd_ridge = predict_ridge(test,lags).to_pandas()\n    pd_tabm = predict_tabm(test,lags).to_pandas()\n\n    pd_nn = pd_nn.rename(columns={'responder_6': 'col_nn'})\n    pd_starter = pd_starter.rename(columns={'responder_6': 'col_starter'})\n    pd_ridge = pd_ridge.rename(columns={'responder_6': 'col_ridge'})\n    pd_tabm = pd_tabm.rename(columns={'responder_6': 'col_tabm'})\n\n    pds = pd.merge(pd_nn, pd_starter, on=['row_id'])\n    pds = pd.merge(pds, pd_ridge, on=['row_id'])\n    pds = pd.merge(pds, pd_tabm, on=['row_id'])\n\n    pds['responder_6'] = \\\n        pds['col_nn'] * ensem_koef['nn'] + \\\n        pds['col_starter'] * ensem_koef['starter'] + \\\n        pds['col_ridge'] * ensem_koef['ridge'] + \\\n        pds['col_tabm'] * ensem_koef['tabm'] \n\n    if not is_first:\n        display(pds)\n        is_first = True\n\n    predictions = test.select('row_id', pl.lit(0.0).alias('responder_6'))\n    pred = pds['responder_6'].to_numpy()\n    predictions = predictions.with_columns(pl.Series('responder_6', pred.ravel()))\n    return predictions\n\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    )","metadata":{"_kg_hide-input":false,"_kg_hide-output":false,"papermill":{"duration":0.351292,"end_time":"2024-10-26T03:27:42.101707","exception":false,"start_time":"2024-10-26T03:27:41.750415","status":"completed"},"tags":[],"trusted":true,"execution":{"iopub.status.busy":"2025-01-02T13:10:06.531359Z","iopub.execute_input":"2025-01-02T13:10:06.531584Z","iopub.status.idle":"2025-01-02T13:10:07.020076Z","shell.execute_reply.started":"2025-01-02T13:10:06.531556Z","shell.execute_reply":"2025-01-02T13:10:07.019418Z"}},"outputs":[],"execution_count":null}]}