{"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"}],"dockerImageVersionId":30786,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import numpy as np\nimport polars as pl\nfrom pathlib import Path\nimport gc\nimport os\nfrom typing import List, Union, Dict, Any\n\nimport torch\nfrom torch import Tensor\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2024-12-21T00:46:26.91282Z","iopub.execute_input":"2024-12-21T00:46:26.913284Z","iopub.status.idle":"2024-12-21T00:46:36.109605Z","shell.execute_reply.started":"2024-12-21T00:46:26.913239Z","shell.execute_reply":"2024-12-21T00:46:36.108353Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"BASE_PATH = Path('/kaggle/input/jane-street-real-time-market-data-forecasting/train.parquet')\n\ntrain_ds = pl.concat([\n    pl.scan_parquet(BASE_PATH / f'partition_id={i}' / 'part-0.parquet')\n    for i in range(5, 10)\n]).sort(['date_id', 'time_id', 'symbol_id']).fill_nan(0).fill_null(0)\n\nmax_date = train_ds.select(pl.col('date_id').max()).collect().item()\nval_ds = train_ds.filter(pl.col('date_id') > max_date - 30)\ntrain_ds = train_ds.filter(pl.col('date_id') <= max_date - 30)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T00:46:36.111018Z","iopub.execute_input":"2024-12-21T00:46:36.111732Z","iopub.status.idle":"2024-12-21T00:46:36.149257Z","shell.execute_reply.started":"2024-12-21T00:46:36.111679Z","shell.execute_reply":"2024-12-21T00:46:36.147647Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_device():\n    if torch.cuda.is_available():\n        return 'cuda'\n    return 'cpu'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T00:46:36.151224Z","iopub.execute_input":"2024-12-21T00:46:36.151765Z","iopub.status.idle":"2024-12-21T00:46:36.172747Z","shell.execute_reply.started":"2024-12-21T00:46:36.151721Z","shell.execute_reply":"2024-12-21T00:46:36.171363Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class BaseDataset(Dataset):\n    \n    def __init__(self, dataset: pl.DataFrame):\n        super(BaseDataset, self).__init__()   \n        self.dataset = dataset\n        feature_cols = [f'feature_{i:02d}' for i in range(79)]\n        self.X = torch.FloatTensor(self.dataset.select(feature_cols).collect().to_numpy().astype(np.float32))\n        self.y = torch.FloatTensor(self.dataset.select(['responder_6']).collect().to_numpy().flatten().astype(np.float32))\n        self.weights = torch.FloatTensor(self.dataset.select(['weight']).collect().to_numpy().flatten().astype(np.float32))        \n    \n    def __len__(self):\n        return self.X.shape[0]\n    \n    def __getitem__(self, idx):       \n        return self.X[idx], self.y[idx], self.weights[idx]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T00:46:36.17541Z","iopub.execute_input":"2024-12-21T00:46:36.176166Z","iopub.status.idle":"2024-12-21T00:46:36.205133Z","shell.execute_reply.started":"2024-12-21T00:46:36.176125Z","shell.execute_reply":"2024-12-21T00:46:36.203867Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class WeightedMSELoss(nn.Module):\n    def __init__(self):\n        super(WeightedMSELoss, self).__init__()\n    \n    def forward(self, predictions: Tensor, targets: Tensor, weights: Tensor) -> Tensor:\n        squared_diff = (predictions - targets) ** 2\n        weighted_squared_diff = weights * squared_diff\n        return weighted_squared_diff.sum() / weights.sum()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class Mlp(nn.Module):\n    def __init__(self, input_features, hidden_dims=[], dropout_rate=0.1, output_dim=1, \n                 use_tanh=False, final_mult=1.0, bn_momentum=0.1, \n                 initial_bn=True, use_dropout=True, use_norm=True):\n        super(Mlp, self).__init__()\n        self.final_mult = final_mult\n        self.use_tanh = use_tanh\n        self.bn_momentum = bn_momentum\n        \n        layers = []\n        if initial_bn:\n            layers.append(nn.BatchNorm1d(input_features, momentum=bn_momentum))\n        if use_dropout:\n            layers.append(nn.Dropout(dropout_rate))\n            \n        in_features = input_features\n        for hidden_dim in hidden_dims:\n            layers.append(nn.Linear(in_features, hidden_dim))\n            if use_norm:\n                layers.append(nn.BatchNorm1d(hidden_dim, momentum=bn_momentum))\n            layers.append(nn.SiLU())\n            if use_dropout:\n                layers.append(nn.Dropout(dropout_rate))\n            in_features = hidden_dim\n            \n        layers.append(nn.Linear(in_features, output_dim))\n        if self.use_tanh:\n            layers.append(nn.Tanh())\n        \n        self.model = nn.Sequential(*layers)\n        \n    def forward(self, x):\n        return self.final_mult * self.model(x)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def evaluate_model(model, val_dl, device):\n    ss_res = 0.0\n    ss_tot = 0.0\n    for x, targets, w in val_dl:\n        with torch.no_grad():\n            y_out = model(x.to(device)).squeeze()\n        w = w.to(device)\n        targets = targets.to(device)\n        ss_res += (w * (y_out - targets) ** 2).sum().cpu()\n        ss_tot += (w * (targets ** 2)).sum().cpu()\n    return 1 - ss_res / ss_tot","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T00:46:56.089699Z","iopub.execute_input":"2024-12-21T00:46:56.090146Z","iopub.status.idle":"2024-12-21T00:46:56.097455Z","shell.execute_reply.started":"2024-12-21T00:46:56.090101Z","shell.execute_reply":"2024-12-21T00:46:56.095902Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def train_with_es(model, optimizer, train_dl, val_dl, epochs_max, gradient_clipping, loss_fn, is_weighted_loss, output_dir, es_patience, device):\n    \n    save_path = os.path.join(output_dir, 'best_model.pth')\n    torch.save(model.state_dict(), save_path)\n    best_score = evaluate_model(model, val_dl, device).item()\n    print(f'Initial weighted r2: {best_score}')\n    best_epoch = -1\n    for epoch in range(epochs_max):\n        model.train()\n        for x, targets, w in train_dl:\n            optimizer.zero_grad()\n            y_out = model.forward(x.to(device)).squeeze()\n            if is_weighted_loss:\n                loss = loss_fn(y_out, targets.to(device), w.to(device))\n            else:\n                loss = loss_fn(y_out, targets.to(device))\n            loss.backward()\n            torch.nn.utils.clip_grad_norm_(model.parameters(), gradient_clipping)\n            optimizer.step()\n            \n        model.eval()\n        score = evaluate_model(model, val_dl, device).item()\n        print(f'Epoch {epoch} weighted r2: {score}')\n        if score > best_score:\n            torch.save(model.state_dict(), save_path)\n            best_epoch = epoch\n            best_score = score\n        elif epoch - best_epoch >= es_patience:\n            print(f'Stopping after {epoch} epochs')\n            break\n        \n    model.load_state_dict(torch.load(save_path, weights_only=True))\n    model = model.to(device)\n    return model, score","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T00:46:56.099262Z","iopub.execute_input":"2024-12-21T00:46:56.099749Z","iopub.status.idle":"2024-12-21T00:46:56.133121Z","shell.execute_reply.started":"2024-12-21T00:46:56.099693Z","shell.execute_reply":"2024-12-21T00:46:56.131701Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_dataset = BaseDataset(train_ds)\nval_dataset = BaseDataset(val_ds)\n\ntrain_dataloader = DataLoader(train_dataset, batch_size=1024, shuffle=True, num_workers=3)\nval_dataloader = DataLoader(val_dataset, batch_size=8192, shuffle=False, num_workers=3)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T00:46:36.20663Z","iopub.execute_input":"2024-12-21T00:46:36.207023Z","iopub.status.idle":"2024-12-21T00:46:56.088259Z","shell.execute_reply.started":"2024-12-21T00:46:36.206987Z","shell.execute_reply":"2024-12-21T00:46:56.087071Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if torch.cuda.is_available():\n    device = 'cuda:0'\nelse:\n    device = 'cpu'","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = Mlp(\n    input_features=79,\n    output_dim=1,\n    hidden_dims=[512, 256],\n    dropout_rate=0.2,\n    bn_momentum=0.01)\nmodel = model.to(device)\n\noptimizer = torch.optim.AdamW(model.parameters(), lr=0.00093, weight_decay=0.0016748563011872717, betas=(0.88, 0.979))\nloss_fn = WeightedMSELoss()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"checkpoint_dir = '/kaggle/working/model'\nos.makedirs(checkpoint_dir)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model, _ = train_with_es(model, optimizer, train_dataloader, val_dataloader, 100, 50, loss_fn, True, checkpoint_dir, 5, device)\nmodel","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T00:46:56.220179Z","iopub.execute_input":"2024-12-21T00:46:56.220941Z","iopub.status.idle":"2024-12-21T00:49:50.316066Z","shell.execute_reply.started":"2024-12-21T00:46:56.220866Z","shell.execute_reply":"2024-12-21T00:49:50.314765Z"}},"outputs":[],"execution_count":null}]}