{"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":9925758,"sourceType":"datasetVersion","datasetId":6100740},{"sourceId":10386024,"sourceType":"datasetVersion","datasetId":6434157}],"dockerImageVersionId":30787,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!cp /kaggle/input/tabm-swish/tabm_reference_swish.py ./","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"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\nfrom tabm_reference_swish import Model\nimport rtdl_num_embeddings ","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')\narch_type = 'tabm'\nbins = None\ntask_type = 'regression'\nn_classes = None\nprocess_labels = False\ncompile_model = False\nkey_cols = ['date_id','time_id','symbol_id']\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)]\nresponders = [f'responder_{i}' for i in range(9)]\n\norigin_feats = origin_feats + time_feats\n\nn_cont_features = len(origin_feats)\n\ncat_cardinalities = []\n\nmodel = 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=None,\n    num_embeddings=None,\n    arch_type=arch_type,\n    k=32,\n).cuda()\nmodel.load_state_dict(torch.load('/kaggle/input/js-swa-pure-wide/swa_pure_tabm_wide.pth'))\noptimizer = torch.optim.AdamW(model.parameters(), lr=3e-5, weight_decay=0.0003)\n\nbase_loss_fn = F.mse_loss\n\ndef loss_fn(y_pred: Tensor, y_true: Tensor) -> Tensor:\n    # TabM produces k predictions per object. Each of them must be trained separately.\n    # (regression)     y_pred.shape == (batch_size, k)\n    # (classification) y_pred.shape == (batch_size, k, n_classes)\n    k = y_pred.shape[-1 if task_type == 'regression' else -2]\n    return base_loss_fn(y_pred.flatten(0, 1), y_true.repeat_interleave(k)) / (y_true.squeeze() ** 2.0).sum()","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\n\ndef predict(test, lags):\n    global model\n    global data_, lags_\n    global optimizer\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['responder_6'] = train['responder_6_lag_1']\n            train = train.loc[~train.responder_6.isna()].reset_index(drop=True)\n\n        if len(train) > 0:\n            data = {\n                'train': {'x_cont': torch.as_tensor(train[origin_feats].values, device=device).float(), \n                          'y': torch.as_tensor(train['responder_6'].values, device=device).float()},\n            }\n    \n            batch_size = 512\n            for batch_idx in torch.randperm(len(data['train']['y']), device=device).split(batch_size):\n                model.train()\n                optimizer.zero_grad()\n                loss = loss_fn(model(data['train']['x_cont'][batch_idx]).squeeze(-1).float(), \n                               data['train']['y'][batch_idx]\n                              )\n                loss.backward()\n                optimizer.step()\n            del data\n\n        del data_, train\n        data_ = []\n\n    model.eval();\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['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\n    test=test.fillna(0)    \n    data_.append(test[['date_id','time_id','symbol_id'] + origin_feats])\n    \n    test=test[origin_feats].values\n    #print(test.shape)\n    eps=1e-10\n    with torch.no_grad():\n        preds = model(torch.from_numpy(test).float().to(device)).mean(1).squeeze(-1).float().detach().cpu().numpy()\n    test_preds=np.clip(preds, -5+eps, 5-eps )\n    #print(test_preds.shape)\n    predictions = predictions.with_columns(pl.Series('responder_6', test_preds.ravel()))\n    #print(predictions)\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}]}