{"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":"none","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","isSourceIdPinned":true,"modelInstanceId":193541,"modelId":215469},{"sourceId":226966,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":193544,"modelId":215472}],"dockerImageVersionId":30786,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"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 collections import defaultdict, deque\nfrom typing import Optional\n\n# GLOBALS\nloaded_model = None\nloaded_scaler = None\nlags_ = None\nrolling_windows = defaultdict(lambda: deque(maxlen=4)) \n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\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: [batch_size, seq_length, input_dim]\n        out, (hn, cn) = self.lstm(x)\n        out = out[:, -1, :]  # last time step\n        out = self.fc(out)\n        return out\n\n# Paths to your saved model and scaler\nMODEL_PATH = '/kaggle/input/lstm_lags_0063/pytorch/default/1/best_model_so_far.pth'\nSCALER_PATH = '/kaggle/input/scaler_lstm_lag_0063/scikitlearn/default/1/scaler.joblib'\n\ndef filling_na(feature_cols, df: pl.DataFrame) -> pl.DataFrame:\n    \"\"\"\n    Fill NA with 0\n    \"\"\"\n    df = df.fill_null(0)\n    return df\n\ndef predict(test: pl.DataFrame, lags: Optional[pl.DataFrame]) -> pl.DataFrame:\n    \"\"\"\n    Called by the Kaggle real-time API.\n    Must return Polars DF with columns ['row_id', 'responder_6'].\n    \"\"\"\n\n    global loaded_model, loaded_scaler, lags_, rolling_windows, device\n\n    # 1) Lazy-load model/scaler if not done\n    if loaded_model is None:\n        input_dim = 91\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 lags is not None => store them (renamed) for day boundary\n    #    The API only provides them at time_id=0 for new date_id.\n    if lags is not None:\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) Combine the test row(s) with cached lags_ if it exists\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) Identify feature columns\n    # Exclude 'responder_*', 'weight', 'row_id', 'is_scored'\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) Apply your rolling-mean fill + fill_null(0) approach\n    df_test = filling_na(feature_cols, combined)\n\n    # 6) Convert to NumPy + scale\n    X_array = df_test.select(feature_cols).to_numpy()\n    X_scaled = loaded_scaler.transform(X_array).astype(np.float32)\n\n    # We'll do row-by-row inference so we can handle seq_length=5\n    row_ids = df_test['row_id'].to_list()\n    symbol_ids = df_test['symbol_id'].to_list()\n    y_preds = np.zeros(len(row_ids), dtype=np.float32)\n\n    for i, row_id in enumerate(row_ids):\n        symbol = symbol_ids[i]\n        this_row_features = X_scaled[i]  # shape [input_dim]\n\n        # 7) Build the 5-step sequence from rolling buffer\n        prev_rows = list(rolling_windows[symbol])  # up to 4 old rows\n        full_sequence = prev_rows + [this_row_features]  # up to 5 steps\n        seq_len = len(full_sequence)\n\n        # if seq_len < 5, we either accept partial or pad. Let's accept partial\n        seq_np = np.stack(full_sequence, axis=0)  # shape [seq_len, input_dim]\n        seq_np = seq_np[np.newaxis, :, :]         # shape [batch=1, seq_len, input_dim]\n\n        # 8) LSTM forward pass\n        seq_tensor = torch.tensor(seq_np, device=device)\n        with torch.no_grad():\n            pred = loaded_model(seq_tensor).view(-1).item()\n        y_preds[i] = pred\n\n        # 9) Update rolling buffer with this row\n        rolling_windows[symbol].append(this_row_features)\n    \n    # 10) Build final output Polars DataFrame\n    out_df = pl.DataFrame({\n        'row_id': row_ids,\n        'responder_6': y_preds\n    })\n\n    # Ensure correct shape\n    assert len(out_df) == len(combined), \"Number of predictions must match test rows.\"\n\n    return out_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-12T04:59:32.697980Z","iopub.execute_input":"2025-01-12T04:59:32.698251Z","iopub.status.idle":"2025-01-12T04:59:37.559868Z","shell.execute_reply.started":"2025-01-12T04:59:32.698221Z","shell.execute_reply":"2025-01-12T04:59:37.558743Z"}},"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-12T04:59:46.130923Z","iopub.execute_input":"2025-01-12T04:59:46.131317Z","iopub.status.idle":"2025-01-12T04:59:47.176455Z","shell.execute_reply.started":"2025-01-12T04:59:46.131279Z","shell.execute_reply":"2025-01-12T04:59:47.175374Z"}},"outputs":[],"execution_count":null}]}