{"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":10004078,"sourceType":"datasetVersion","datasetId":6158005},{"sourceId":10004352,"sourceType":"datasetVersion","datasetId":6158178},{"sourceId":10004448,"sourceType":"datasetVersion","datasetId":6158247},{"sourceId":177576,"sourceType":"modelInstanceVersion","modelInstanceId":151273,"modelId":173745},{"sourceId":177591,"sourceType":"modelInstanceVersion","modelInstanceId":151288,"modelId":173760},{"sourceId":177620,"sourceType":"modelInstanceVersion","modelInstanceId":151306,"modelId":173777},{"sourceId":182165,"sourceType":"modelInstanceVersion","modelInstanceId":155278,"modelId":177752},{"sourceId":185164,"sourceType":"modelInstanceVersion","modelInstanceId":157859,"modelId":180261},{"sourceId":186257,"sourceType":"modelInstanceVersion","modelInstanceId":158797,"modelId":181180},{"sourceId":226962,"sourceType":"modelInstanceVersion","modelInstanceId":193541,"modelId":215469},{"sourceId":226966,"sourceType":"modelInstanceVersion","modelInstanceId":193544,"modelId":215472},{"sourceId":228150,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":194534,"modelId":216442},{"sourceId":228151,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":194535,"modelId":216443}],"dockerImageVersionId":30786,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport polars as pl\nimport pandas as pd\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport joblib\nimport sys\nimport kaggle_evaluation.jane_street_inference_server\nfrom typing import Optional\n\n# Globals for caching\nloaded_model = None\nloaded_scaler = None\nlags_ = None  # If you still use day-lag columns, store them here (renamed) for each day\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n\nclass TwoLayerLSTM(nn.Module):\n    def __init__(self, input_dim, hidden_dim, output_dim=1, dropout=0.2):\n        super(TwoLayerLSTM, self).__init__()\n        self.lstm = nn.LSTM(\n            input_size=input_dim,\n            hidden_size=hidden_dim,\n            num_layers=2,\n            batch_first=True,\n            dropout=dropout\n        )\n        self.fc = nn.Linear(hidden_dim, output_dim)\n\n    def forward(self, x):\n        # x shape: [batch_size, seq_length=1, input_dim]\n        out, (hn, cn) = self.lstm(x)\n        out = out[:, -1, :]  # last time step (since seq_length=1, it's out[:, 0, :]\n        out = self.fc(out)\n        return out\n\n# Model / Scaler paths in the Kaggle environment\nMODEL_PATH = '/kaggle/input/best-model-lstm-nn-lags/pytorch/default/1/best_model_LSTM_NN.pth'\nSCALER_PATH = '/kaggle/input/best-scaler-lstm-nn-lags/scikitlearn/default/1/scaler_best_model_LSTM_NN.joblib'\n\n\ndef filling_na(feature_cols, df: pl.DataFrame) -> pl.DataFrame:\n    \"\"\"\n    Fill all null values with 0, as done in training.\n    \"\"\"\n    return df.fill_null(0)\n\n\ndef predict(test: pl.DataFrame, lags: Optional[pl.DataFrame]) -> pl.DataFrame:\n    \"\"\"\n    Called by the Kaggle real-time API. \n    Must return a Polars DF with columns ['row_id', 'responder_6'].\n    Each row is predicted independently with seq_length=1.\n    \"\"\"\n\n    global loaded_model, loaded_scaler, lags_, device\n\n    # 1) Lazy-load the model/scaler once\n    if loaded_model is None:\n        input_dim = 91          # match the training dimension\n        hidden_dim = 32\n        dropout = 0.4\n        loaded_model = TwoLayerLSTM(input_dim, hidden_dim, output_dim=1, dropout=dropout)\n        loaded_model.load_state_dict(torch.load(MODEL_PATH, map_location=device))\n        loaded_model.to(device)\n        loaded_model.eval()\n\n    if loaded_scaler is None:\n        loaded_scaler = joblib.load(SCALER_PATH)\n\n    # 2) If we have new day-lag data, rename columns and store (if you do day-lag merges)\n    if lags is not None:\n        # e.g., rename from 'responder_0_lag_1' to 'lag_1_responder_0', etc.\n        lags_renamed = lags.rename({\n            'responder_0_lag_1': 'lag_1_responder_0',\n            'responder_1_lag_1': 'lag_1_responder_1',\n            'responder_2_lag_1': 'lag_1_responder_2',\n            'responder_3_lag_1': 'lag_1_responder_3',\n            'responder_4_lag_1': 'lag_1_responder_4',\n            'responder_5_lag_1': 'lag_1_responder_5',\n            'responder_6_lag_1': 'lag_1_responder_6',\n            'responder_7_lag_1': 'lag_1_responder_7',\n            'responder_8_lag_1': 'lag_1_responder_8'\n        })\n        lags_ = lags_renamed\n\n    # 3) Merge 'test' with 'lags_' if available\n    if lags_ is not None:\n        combined = test.join(lags_, on=[\"date_id\",\"time_id\",\"symbol_id\"], how=\"left\")\n    else:\n        combined = test\n\n    # 4) Exclude columns not used in training\n    excluded_cols = [c for c in combined.columns if c.startswith('responder_')]\n    excluded_cols += ['weight', 'row_id', 'is_scored']\n    feature_cols = [c for c in combined.columns if c not in excluded_cols]\n\n    # 5) Fill null => 0 and scale\n    df_test = filling_na(feature_cols, combined)\n    X_array = df_test.select(feature_cols).to_numpy()        # shape: [num_rows, input_dim]\n    X_scaled = loaded_scaler.transform(X_array).astype(np.float32)\n\n    # 6) Single-timestep inference for each row (or do a single batch pass):\n    row_ids = df_test['row_id'].to_list()\n    y_preds = np.zeros(len(row_ids), dtype=np.float32)\n\n    # We can do a single batch pass in one go:\n    # shape => [batch_size, seq_length=1, input_dim]\n    batch_size = X_scaled.shape[0]\n    input_dim = X_scaled.shape[1]\n\n    # Reshape to [batch_size, 1, input_dim]\n    X_tensor = torch.tensor(X_scaled.reshape(batch_size, 1, input_dim),\n                            dtype=torch.float32,\n                            device=device)\n\n    with torch.no_grad():\n        pred_tensor = loaded_model(X_tensor).view(-1).cpu().numpy()\n    y_preds[:] = pred_tensor\n\n    # 7) Build the final Polars DataFrame\n    out_df = pl.DataFrame({\n        'row_id': row_ids,\n        'responder_6': y_preds\n    })\n\n    # Basic shape checks\n    assert len(out_df) == len(combined), \"Number of predictions must match test rows.\"\n\n    return out_df\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T05:35:35.211094Z","iopub.execute_input":"2025-01-13T05:35:35.211440Z","iopub.status.idle":"2025-01-13T05:35:39.221802Z","shell.execute_reply.started":"2025-01-13T05:35:35.211408Z","shell.execute_reply":"2025-01-13T05:35:39.220932Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"inference_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":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-13T05:35:43.513806Z","iopub.execute_input":"2025-01-13T05:35:43.514285Z","iopub.status.idle":"2025-01-13T05:35:44.749113Z","shell.execute_reply.started":"2025-01-13T05:35:43.514253Z","shell.execute_reply":"2025-01-13T05:35:44.748387Z"}},"outputs":[],"execution_count":null}]}