{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":104854,"databundleVersionId":12637092,"sourceType":"competition"},{"sourceId":248784268,"sourceType":"kernelVersion"}],"dockerImageVersionId":31041,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"### >>> [Yandex Research TabM Implentation](https://github.com/yandex-research/tabm) <<<","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19"}},{"cell_type":"code","source":"%%capture\n!pip install rtdl_num_embeddings\n!pip install tabm","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-04T10:31:23.251064Z","iopub.execute_input":"2025-07-04T10:31:23.251349Z","iopub.status.idle":"2025-07-04T10:32:41.312516Z","shell.execute_reply.started":"2025-07-04T10:31:23.251326Z","shell.execute_reply":"2025-07-04T10:32:41.311483Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import math\nimport random\nfrom copy import deepcopy\nfrom typing import Any, Literal, NamedTuple\n\nimport numpy as np\nimport pandas as pd\nimport rtdl_num_embeddings \nimport scipy.special\nimport sklearn.datasets\nimport sklearn.metrics\nimport sklearn.model_selection\nimport sklearn.preprocessing\nimport tabm\nimport torch\nimport torch.nn.functional as F\nimport torch.optim\nfrom torch import Tensor\nfrom sklearn.metrics import classification_report\nimport warnings\nwarnings.filterwarnings('ignore')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-04T10:33:14.312598Z","iopub.execute_input":"2025-07-04T10:33:14.313288Z","iopub.status.idle":"2025-07-04T10:33:20.700625Z","shell.execute_reply.started":"2025-07-04T10:33:14.313257Z","shell.execute_reply":"2025-07-04T10:33:20.699843Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Features were cooked here -> 🅰️ [Alfa Features Kitchen](https://www.kaggle.com/code/samvelkoch/alfa-features-kitchen/notebook?scriptVersionId=248784268) | 🔥👨‍🍳","metadata":{}},{"cell_type":"code","source":"seed = 42\nrandom.seed(seed)\nnp.random.seed(seed + 1)\ntorch.manual_seed(seed + 2)\n\n# === Load Data ===\nprint(\"Loading data...\")\ntrain_df = pd.read_parquet('/kaggle/input/alfa-features-kitchen/SET_2_train_data.parquet')\ntest_df = pd.read_parquet('/kaggle/input/alfa-features-kitchen/SET_2_test_data.parquet')\n\nprint(f\"Train data shape: {train_df.shape}\")\nprint(f\"Test data shape: {test_df.shape}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-04T10:33:22.422059Z","iopub.execute_input":"2025-07-04T10:33:22.422886Z","iopub.status.idle":"2025-07-04T10:33:24.152553Z","shell.execute_reply.started":"2025-07-04T10:33:22.422859Z","shell.execute_reply":"2025-07-04T10:33:24.151800Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# === Custom WMAE metric ===\nweights = {0: 1, 1: 0.72, 2: 0.52, 3: 0.37, 4: 0.27, 5: 0.19, 6: 0.14}\n\ndef wmae_metric_func(y_true, y_pred):\n    y_true = np.array(y_true).astype(int)\n    y_pred = np.array(y_pred).astype(int)\n    abs_errors = np.abs(y_true - y_pred)\n    weighted_errors = [weights[int(y)] * err for y, err in zip(y_true, abs_errors)]\n    return np.mean(weighted_errors)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-04T10:33:27.384342Z","iopub.execute_input":"2025-07-04T10:33:27.385045Z","iopub.status.idle":"2025-07-04T10:33:27.389365Z","shell.execute_reply.started":"2025-07-04T10:33:27.385019Z","shell.execute_reply":"2025-07-04T10:33:27.388704Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# === Prepare dataset for TabM ===\ntask_type: Literal['multiclass'] = 'multiclass'\nn_classes = 7\n\n# Prepare features (exclude target and client_num)\nfeature_cols = [col for col in train_df.columns if col not in ['target', 'client_num']]\nX_num = train_df[feature_cols].values.astype(np.float32)\nY = train_df['target'].values.astype(np.int64)\n\n# Test data features\nX_test = test_df[feature_cols].values.astype(np.float32)\n\nprint(f\"Feature columns: {len(feature_cols)}\")\nprint(f\"Target distribution:\\n{pd.Series(Y).value_counts().sort_index()}\")\n\n# Validate that Y contains values 0 to n_classes-1\nassert set(Y.tolist()) == set(range(n_classes)), (\n    'Classification labels must form the range [0, 1, ..., n_classes - 1]'\n)\n\n# === Split the dataset ===\n# Create train/val split for cross-validation\ntrain_idx, val_idx = sklearn.model_selection.train_test_split(\n    np.arange(len(Y)), train_size=0.85, random_state=seed, stratify=Y\n)\n\n# Prepare data structure\ndata_numpy = {\n    'train': {'x_num': X_num[train_idx], 'y': Y[train_idx]},\n    'val': {'x_num': X_num[val_idx], 'y': Y[val_idx]},\n    'test': {'x_num': X_test, 'y': np.zeros(len(X_test), dtype=np.int64)}  # Dummy labels for test\n}\n\n# No categorical features in our case\ncat_cardinalities = []\nX_cat = None\n\nn_num_features = X_num.shape[1]\n\nprint(f\"\\nData splits:\")\nfor part, part_data in data_numpy.items():\n    for key, value in part_data.items():\n        print(f'{part:<5}    {key:<5}    {value.shape!r:<15}    {value.dtype}')\n\n# === Feature preprocessing ===\nprint(\"\\nApplying feature preprocessing...\")\n\n# Check for constant features and remove them\nx_num_train_numpy = data_numpy['train']['x_num']\n\n# Find non-constant features\nfeature_std = np.std(x_num_train_numpy, axis=0)\nnon_constant_mask = feature_std > 1e-8  # Features with meaningful variance\nconstant_features = np.sum(~non_constant_mask)\n\nif constant_features > 0:\n    print(f\"Found {constant_features} constant features, removing them...\")\n    # Filter out constant features from all datasets\n    for part in data_numpy:\n        data_numpy[part]['x_num'] = data_numpy[part]['x_num'][:, non_constant_mask]\n    \n    # Update feature info\n    n_num_features = data_numpy['train']['x_num'].shape[1]\n    print(f\"Features after filtering: {n_num_features}\")\n\n# Advanced preprocessing strategy with QuantileTransformer\nx_num_train_numpy = data_numpy['train']['x_num']\nnoise = (\n    np.random.default_rng(0)\n    .normal(0.0, 1e-5, x_num_train_numpy.shape)\n    .astype(x_num_train_numpy.dtype)\n)\n\npreprocessing = sklearn.preprocessing.QuantileTransformer(\n    n_quantiles=max(min(len(train_idx) // 30, 1000), 10),\n    output_distribution='normal',\n    subsample=10**9,\n).fit(x_num_train_numpy + noise)\n\n# Apply preprocessing to all parts\nfor part in data_numpy:\n    data_numpy[part]['x_num'] = preprocessing.transform(data_numpy[part]['x_num'])\n\nprint(\"Preprocessing completed.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-04T10:33:30.102749Z","iopub.execute_input":"2025-07-04T10:33:30.102988Z","iopub.status.idle":"2025-07-04T10:33:38.110161Z","shell.execute_reply.started":"2025-07-04T10:33:30.102972Z","shell.execute_reply":"2025-07-04T10:33:38.109552Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# === PyTorch setup ===\ndevice = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")\n\n# Convert data to tensors\ndata = {\n    part: {k: torch.as_tensor(v, device=device) for k, v in data_numpy[part].items()}\n    for part in data_numpy\n}\n\n# AMP settings\namp_dtype = (\n    torch.bfloat16\n    if torch.cuda.is_available() and torch.cuda.is_bf16_supported()\n    else torch.float16\n    if torch.cuda.is_available()\n    else None\n)\namp_enabled = False and amp_dtype is not None\ngrad_scaler = torch.cuda.amp.GradScaler() if amp_dtype is torch.float16 else None\n\ncompile_model = False\n\nprint(f'Device:        {device.type.upper()}')\nprint(f'AMP:           {amp_enabled}{f\" ({amp_dtype})\" if amp_enabled else \"\"}')\nprint(f'torch.compile: {compile_model}')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-04T10:33:38.111204Z","iopub.execute_input":"2025-07-04T10:33:38.111459Z","iopub.status.idle":"2025-07-04T10:33:38.376045Z","shell.execute_reply.started":"2025-07-04T10:33:38.111436Z","shell.execute_reply":"2025-07-04T10:33:38.375333Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# === Model creation ===\nprint(\"\\nCreating TabM model...\")\n\n# Use PiecewiseLinearEmbeddings for best performance\ntry:\n    num_embeddings = rtdl_num_embeddings.PiecewiseLinearEmbeddings(\n        rtdl_num_embeddings.compute_bins(data['train']['x_num'], n_bins=48),\n        d_embedding=16,\n        activation=False,\n        version='B',\n    )\n    print(\"Using PiecewiseLinearEmbeddings\")\nexcept ValueError as e:\n    print(f\"PiecewiseLinearEmbeddings failed: {e}\")\n    print(\"Falling back to PeriodicEmbeddings\")\n    # Fallback to PeriodicEmbeddings if PiecewiseLinear fails\n    num_embeddings = rtdl_num_embeddings.PeriodicEmbeddings(n_num_features, lite=False)\n\nmodel = tabm.TabM.make(\n    n_num_features=n_num_features,\n    cat_cardinalities=cat_cardinalities,\n    d_out=n_classes,  # 7 classes\n    num_embeddings=num_embeddings,\n).to(device)\n\noptimizer = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=3e-4)\n\nif compile_model:\n    model = torch.compile(model)\n    evaluation_mode = torch.no_grad\nelse:\n    evaluation_mode = torch.inference_mode\n\nprint(f\"Model created with {sum(p.numel() for p in model.parameters())} parameters\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-04T10:33:40.864033Z","iopub.execute_input":"2025-07-04T10:33:40.864304Z","iopub.status.idle":"2025-07-04T10:33:45.984411Z","shell.execute_reply.started":"2025-07-04T10:33:40.864282Z","shell.execute_reply":"2025-07-04T10:33:45.983665Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# === Training setup ===\nshare_training_batches = True\n\n@torch.autocast(device.type, enabled=amp_enabled, dtype=amp_dtype)\ndef apply_model(part: str, idx: Tensor) -> Tensor:\n    return model(\n        data[part]['x_num'][idx],\n        data[part]['x_cat'][idx] if 'x_cat' in data[part] else None,\n    ).float()\n\nbase_loss_fn = F.cross_entropy\n\ndef loss_fn(y_pred: Tensor, y_true: Tensor) -> Tensor:\n    # TabM produces k predictions. Each of them must be trained separately.\n    # Classification: (batch_size, k, n_classes) -> (batch_size * k, n_classes)\n    y_pred = y_pred.flatten(0, 1)\n    \n    if share_training_batches:\n        # (batch_size,) -> (batch_size * k,)\n        y_true = y_true.repeat_interleave(model.backbone.k)\n    else:\n        # (batch_size, k) -> (batch_size * k,)\n        y_true = y_true.flatten(0, 1)\n    \n    return base_loss_fn(y_pred, y_true)\n\n@evaluation_mode()\ndef evaluate(part: str, return_predictions: bool = False) -> float | tuple[float, np.ndarray, np.ndarray]:\n    model.eval()\n    \n    eval_batch_size = 8096\n    y_pred_logits = torch.cat([\n        apply_model(part, idx)\n        for idx in torch.arange(len(data[part]['y']), device=device).split(eval_batch_size)\n    ]).cpu().numpy()\n    \n    # For classification, compute mean in probability space\n    y_pred_probs = scipy.special.softmax(y_pred_logits, axis=-1)\n    y_pred_mean_probs = y_pred_probs.mean(1)  # Average across k predictions\n    y_pred_classes = y_pred_mean_probs.argmax(1)\n    \n    y_true = data[part]['y'].cpu().numpy()\n    \n    if part == 'test':\n        # For test set, we don't have true labels, so return dummy accuracy\n        accuracy = 0.0\n        wmae_score = 0.0\n    else:\n        accuracy = sklearn.metrics.accuracy_score(y_true, y_pred_classes)\n        wmae_score = wmae_metric_func(y_true, y_pred_classes)\n    \n    # Return negative WMAE as score (higher is better for early stopping)\n    score = -wmae_score if part != 'test' else accuracy\n    \n    if return_predictions:\n        return score, y_pred_mean_probs, y_pred_classes\n    return score\n\nprint(f'Validation score before training: {evaluate(\"val\"):.4f}')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-04T10:33:45.985578Z","iopub.execute_input":"2025-07-04T10:33:45.986014Z","iopub.status.idle":"2025-07-04T10:33:46.480174Z","shell.execute_reply.started":"2025-07-04T10:33:45.985992Z","shell.execute_reply":"2025-07-04T10:33:46.479590Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# === Training loop ===\nprint(\"\\nStarting training...\")\n\nn_epochs = 1000\ntrain_size = len(train_idx)\nbatch_size = 512  # Increased batch size for better stability\nepoch_size = math.ceil(train_size / batch_size)\n\nepoch = -1\nmetrics = {'val': -math.inf, 'test': -math.inf}\n\ndef make_checkpoint() -> dict[str, Any]:\n    return deepcopy({\n        'model': model.state_dict(),\n        'optimizer': optimizer.state_dict(),\n        'epoch': epoch,\n        'metrics': metrics,\n    })\n\nbest_checkpoint = make_checkpoint()\n\n# Early stopping\npatience = 25  # Increased patience for better convergence\nremaining_patience = patience\n\nfor epoch in range(n_epochs):\n    # Training\n    batches = (\n        torch.randperm(train_size, device=device).split(batch_size)\n        if share_training_batches\n        else (\n            torch.rand((train_size, model.backbone.k), device=device)\n            .argsort(dim=0)\n            .split(batch_size, dim=0)\n        )\n    )\n    \n    model.train()\n    epoch_loss = 0.0\n    n_batches = 0\n    \n    for batch_idx in batches:\n        optimizer.zero_grad()\n        y_batch = data['train']['y'][batch_idx]\n        loss = loss_fn(apply_model('train', batch_idx), y_batch)\n        \n        if grad_scaler is None:\n            loss.backward()\n            optimizer.step()\n        else:\n            grad_scaler.scale(loss).backward()\n            grad_scaler.step(optimizer)\n            grad_scaler.update()\n        \n        epoch_loss += loss.item()\n        n_batches += 1\n    \n    # Evaluation\n    metrics = {'val': evaluate('val')}\n    val_score_improved = metrics['val'] > best_checkpoint['metrics']['val']\n    \n    if epoch % 5 == 0 or val_score_improved:\n        print(\n            f'{\"*\" if val_score_improved else \" \"}'\n            f' [epoch] {epoch:<3}'\n            f' [loss] {epoch_loss/n_batches:.4f}'\n            f' [val] {metrics[\"val\"]:.4f}'\n        )\n    \n    if val_score_improved:\n        best_checkpoint = make_checkpoint()\n        remaining_patience = patience\n    else:\n        remaining_patience -= 1\n    \n    if remaining_patience < 0:\n        print(f\"Early stopping at epoch {epoch}\")\n        break\n\n# === Load best model and make final predictions ===\nmodel.load_state_dict(best_checkpoint['model'])\n\nprint('\\n=== Final Evaluation ===')\nprint(f'Best epoch: {best_checkpoint[\"epoch\"]}')\nprint(f'Best validation score: {best_checkpoint[\"metrics\"][\"val\"]:.4f}')\n\n# Get validation predictions for analysis\nval_score, val_probs, val_classes = evaluate('val', return_predictions=True)\nval_true = data['val']['y'].cpu().numpy()\n\nprint(f\"\\nValidation Results:\")\nprint(f\"WMAE Score: {wmae_metric_func(val_true, val_classes):.6f}\")\nprint(f\"Accuracy: {sklearn.metrics.accuracy_score(val_true, val_classes):.4f}\")\n\nprint(\"\\nValidation Classification Report:\")\nprint(classification_report(val_true, val_classes))\n\n# Get test predictions\n_, test_probs, test_classes = evaluate('test', return_predictions=True)\n\nprint(f\"\\nTest predictions distribution:\")\nprint(pd.Series(test_classes).value_counts().sort_index())\n\n# === Save results ===\nprint(\"\\nSaving results...\")\n\n# Main submission file\nsubmission_df = pd.DataFrame({\n    'client_num': test_df['client_num'].values,\n    'target': test_classes\n})\n\nsubmission_df.to_csv('submission.csv', index=False)\nprint(\"Saved submission_tabm.csv\")\n\n# Save probabilities for ensemble\nprobs_df = pd.DataFrame(test_probs, columns=[f'prob_class_{i}' for i in range(n_classes)])\nprobs_df['client_num'] = test_df['client_num'].values\n\n# Reorder columns to have client_num first\ncols = ['client_num'] + [f'prob_class_{i}' for i in range(n_classes)]\nprobs_df = probs_df[cols]\n\nprobs_df.to_csv('test_probabilities_tabm.csv', index=False)\nprint(\"Saved test_probabilities_tabm.csv\")\n\n# Also save validation probabilities for stacking\nval_probs_df = pd.DataFrame(val_probs, columns=[f'prob_class_{i}' for i in range(n_classes)])\nval_probs_df['client_num'] = train_df.iloc[val_idx]['client_num'].values\nval_probs_df['true_target'] = val_true\n\ncols_val = ['client_num', 'true_target'] + [f'prob_class_{i}' for i in range(n_classes)]\nval_probs_df = val_probs_df[cols_val]\n\nval_probs_df.to_csv('val_probabilities_tabm.csv', index=False)\nprint(\"Saved val_probabilities_tabm.csv\")\n\nprint(f\"\\n=== Summary ===\")\nprint(f\"Training completed in {best_checkpoint['epoch']} epochs\")\nprint(f\"Final validation WMAE: {wmae_metric_func(val_true, val_classes):.6f}\")\nprint(f\"Final validation accuracy: {sklearn.metrics.accuracy_score(val_true, val_classes):.4f}\")\nprint(f\"Test predictions saved to submission.csv\")\nprint(f\"Test probabilities saved to test_probabilities_tabm.csv\")\nprint(f\"Validation probabilities saved to val_probabilities_tabm.csv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-04T10:34:14.120868Z","iopub.execute_input":"2025-07-04T10:34:14.121402Z","iopub.status.idle":"2025-07-04T10:36:29.818767Z","shell.execute_reply.started":"2025-07-04T10:34:14.121372Z","shell.execute_reply":"2025-07-04T10:36:29.817609Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission_df.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-04T10:36:34.121011Z","iopub.execute_input":"2025-07-04T10:36:34.121667Z","iopub.status.idle":"2025-07-04T10:36:34.128017Z","shell.execute_reply.started":"2025-07-04T10:36:34.121647Z","shell.execute_reply":"2025-07-04T10:36:34.127458Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"probs_df.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-04T10:36:35.540848Z","iopub.execute_input":"2025-07-04T10:36:35.541118Z","iopub.status.idle":"2025-07-04T10:36:35.551256Z","shell.execute_reply.started":"2025-07-04T10:36:35.541098Z","shell.execute_reply":"2025-07-04T10:36:35.550613Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"That's all folks! Have fun... 🙏\n\n[Samvel Kocharyan](https://www.linkedin.com/in/samvelkoch/), 2025","metadata":{}}]}