{"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":"gpu","dataSources":[{"sourceId":84493,"databundleVersionId":9871156,"sourceType":"competition"},{"sourceId":10443910,"sourceType":"datasetVersion","datasetId":6464480}],"dockerImageVersionId":30787,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import polars as pl\nimport pandas as pd\nimport numpy as np\n\nimport gc\nimport os\nimport kaggle_evaluation.jane_street_inference_server\nimport warnings #avoid some negligible errors\nwarnings.filterwarnings('ignore')\nimport random\n\nimport torch\nfrom torch import Tensor\nimport torch.nn as nn\nimport torch.nn.functional as F\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2024-11-04T10:28:22.733154Z","iopub.execute_input":"2024-11-04T10:28:22.733603Z","iopub.status.idle":"2024-11-04T10:28:25.576544Z","shell.execute_reply.started":"2024-11-04T10:28:22.733547Z","shell.execute_reply":"2024-11-04T10:28:25.575376Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n\ntime_feats = ['sin_time_id', 'cos_time_id', 'sin_time_id_halfday', 'cos_time_id_halfday']\norigin_feats = [f'feature_0{i}' if i<10 else f'feature_{i}' for i in range(79)] + time_feats\nresponders = [f'responder_{i}_lag_1' for i in range(9)]\n\nbase_loss_fn = F.mse_loss\n\ndef loss_fn(y_pred: Tensor, y_true: Tensor) -> Tensor:\n    return F.mse_loss(y_pred, y_true) #/ (y_true ** 2.0).sum()\n\nclass DeepNetworkWithSkipConnections(nn.Module):\n    def __init__(self, hidden_size, output_size):\n        super(DeepNetworkWithSkipConnections, self).__init__()\n       \n        # Создаем 6 слоев\n        self.layers = nn.ModuleList([\n            nn.Sequential(\n                nn.Linear(hidden_size, hidden_size),\n                nn.Hardswish(),\n                #nn.Dropout(0.1)\n            ) for _ in range(3)\n        ])\n        \n        self.output_layer = nn.Linear(hidden_size, output_size)\n        \n    def forward(self, x):        \n        # Проходим через все слои со skip connections\n        for layer in self.layers:\n            # Сохраняем входное значение для skip connection\n            residual = x\n            # Пропускаем через текущий слой\n            x = layer(x)\n            # Добавляем skip connection\n            x = x + residual\n            \n        x = self.output_layer(x)\n        return x\n\nclass SymbolTimePredictor(nn.Module):\n    def __init__(\n        self,\n        feature_dim,\n        hidden_dim = 128,\n        hidden_size = 64,\n        num_heads = 4,\n        num_encoder_layers=1,\n        dropout = 0.0\n    ):\n        super().__init__()\n\n        self.hidden_dim = hidden_dim\n        self.hidden_size = hidden_size\n\n        self.feature_projection = nn.Linear(feature_dim, hidden_dim)\n        # Encoder layer с multi-head attention\n        encoder_layer = nn.TransformerEncoderLayer(\n            d_model=hidden_dim,\n            nhead=num_heads,\n            dim_feedforward=hidden_dim*4,\n            dropout=dropout,\n            batch_first=True\n        )\n        self.transformer_encoder = nn.TransformerEncoder(\n            encoder_layer,\n            num_layers=num_encoder_layers\n        )\n        \n        # Выходной слой для предсказания цены\n        self.rnn = nn.GRU(hidden_dim, hidden_size, num_layers=3, batch_first=True) #num_layers=3, dropout=0.1, \n\n        self.fc = DeepNetworkWithSkipConnections(hidden_size, 1)\n\n    def forward(self, features, hidden):\n        \"\"\"\n        features: tensor shape (batch_size, seq_length, feature_dim)\n        attention_mask: tensor shape (batch_size, seq_length)\n        \"\"\"\n\n        #print(features.shape)\n        x = torch.swapaxes(features, 0, 1)\n        #print(x.shape)\n        #shape: time, symbol, feature\n\n        \n        # Проекция признаков\n        x = self.feature_projection(x)  # (time, seq_length, hidden_dim)\n        #print(x.shape)\n        \n        # Применение transformer encoder\n        \"\"\"\n        if attention_mask is not None:\n            # Преобразование маски для transformer\n            attention_mask = attention_mask.float().masked_fill(\n                attention_mask == 0, float('-inf')\n            ).masked_fill(attention_mask == 1, float(0.0))\n        #, src_key_padding_mask=attention_mask\n        \"\"\"\n        x = self.transformer_encoder(x)\n        #print(x.shape)\n\n        x = torch.swapaxes(x, 0, 1)\n        #print(x.shape)\n        #shape: symbol, time, hidden_dim\n\n        out, hidden = self.rnn(x, hidden)\n        #bs, seq_len, hidden\n        out = self.fc(out)\n        #bs, seq_len, 1\n        return out.squeeze(), hidden\n\nmodel = SymbolTimePredictor(len(origin_feats)).to(device)\nmodel.load_state_dict(torch.load('/kaggle/input/js-mse-rnn/s_swa_model_mse.pth'))\noptimizer = torch.optim.AdamW(model.parameters(), lr=3e-5, weight_decay=0.0003)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-04T10:29:50.805407Z","iopub.execute_input":"2024-11-04T10:29:50.805845Z","iopub.status.idle":"2024-11-04T10:29:50.833712Z","shell.execute_reply.started":"2024-11-04T10:29:50.805807Z","shell.execute_reply":"2024-11-04T10:29:50.832447Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"data_ = []\nlags_ = None\nhs  = None\n\ndef predict(test, lags):\n    global model\n    global data_, lags_\n    global optimizer\n    global hs\n\n    if lags is not None:\n        train = []\n        if len(data_) > 0:\n            train = pd.concat(data_).reset_index(drop=True)\n            train = train.merge(lags.to_pandas(), on=['symbol_id','time_id'])\n            train = train.loc[~train[responders].isna().any(axis=1)].reset_index(drop=True)\n\n        if len(train) > 0:\n            time_ids = int(train.time_id.max() + 1)\n            model.train()\n            df = train.sort_values(['symbol_id','time_id']).reset_index(drop=True)\n            X_rnn = df[origin_feats].values.reshape((len(df)//time_ids, time_ids, len(origin_feats)))\n            X_rnn = torch.from_numpy(X_rnn).float().cuda()\n            y = df['responder_6_lag_1'].values.reshape((len(df)//time_ids, time_ids))\n            y = torch.from_numpy(y).float().cuda()\n            out, _ = model(X_rnn, None)\n            optimizer.zero_grad()\n            loss = loss_fn(out, y)\n            loss.backward()\n            optimizer.step()\n            del df, X_rnn, y, out\n\n        del data_, train\n        data_ = []\n        hs = None\n\n    model.eval();\n    predictions = test.select(\n        'row_id',\n        pl.lit(0.0).alias('responder_6'),\n    )\n    \n    orig_test_row_id = test.to_pandas()[['row_id']]\n    test=test.to_pandas().sort_values('symbol_id').reset_index(drop=True)\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['feature_09'] = test['feature_09'] / 90.0 - 0.5\n    test['feature_10'] = test['feature_10'] / 12.0\n    test['feature_11'] = test['feature_11'] / 600.0 - 0.5\n    test=test.fillna(0)\n   \n    data_.append(test)\n    \n    eps=1e-10\n    with torch.no_grad():\n        preds, hs = model(torch.from_numpy(test[origin_feats].values).float().to(device).unsqueeze(1),hs)\n        hs = hs.detach()\n        preds = preds.squeeze().detach().cpu().numpy()\n    test_preds=np.clip(preds, -5+eps, 5-eps )\n\n    test['responder_6'] = test_preds.ravel()\n    orig_test_row_id = orig_test_row_id.merge(test[['row_id','responder_6']], \n                                              on='row_id', how='left')\n    test_preds = orig_test_row_id['responder_6'].values\n    \n    predictions = predictions.with_columns(pl.Series('responder_6', test_preds))\n    return predictions","metadata":{"trusted":true},"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":"2024-11-04T10:29:53.192175Z","iopub.execute_input":"2024-11-04T10:29:53.192624Z","iopub.status.idle":"2024-11-04T10:29:53.266595Z","shell.execute_reply.started":"2024-11-04T10:29:53.192579Z","shell.execute_reply":"2024-11-04T10:29:53.265533Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}