{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":59093,"databundleVersionId":7469972,"sourceType":"competition"},{"sourceId":235484897,"sourceType":"kernelVersion"}],"dockerImageVersionId":31041,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.preprocessing import LabelEncoder\nfrom sklearn.metrics import confusion_matrix\nimport seaborn as sns\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\nimport warnings\nfrom collections import defaultdict\nimport random\nimport uuid\n\nwarnings.filterwarnings('ignore')\n\n# Configuration\nTEST_PATH = \"/kaggle/working/\"\nBASE_PATH = \"/kaggle/input/hms-harmful-brain-activity-classification\"\nPREPROCESSED_PATH = \"/kaggle/input/preprocessing/preprocessed/eeg\"\nTRAIN_LABELS_PATH = os.path.join(BASE_PATH, \"train.csv\")\nMODEL_OUTPUT_PATH = os.path.join(TEST_PATH, \"models\")\nFEATURE_CACHE_PATH = os.path.join(TEST_PATH, \"feature_cache\")\n\nos.makedirs(MODEL_OUTPUT_PATH, exist_ok=True)\nos.makedirs(FEATURE_CACHE_PATH, exist_ok=True)\n\nCLASSES = ['Seizure', 'LPD', 'GPD', 'LRDA', 'GRDA', 'Other']\nN_CLASSES = len(CLASSES)\nBATCH_SIZE = 32\nEPOCHS = 20  # Increased for better convergence\nDEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {DEVICE}\")\n\n# KL Divergence function\ndef kl_divergence_torch(y_true, y_pred_proba, epsilon=1e-7):  # Slightly larger epsilon for stability\n    y_pred_proba = torch.clamp(y_pred_proba, epsilon, 1 - epsilon)\n    y_true = torch.clamp(y_true, epsilon, 1.0)\n    kl_div = torch.sum(y_true * torch.log(y_true / y_pred_proba), dim=1)\n    return torch.mean(kl_div).item()\n\n# Load and preprocess training labels\ntry:\n    train_df = pd.read_csv(TRAIN_LABELS_PATH)\n    print(f\"Loaded {len(train_df)} annotations with {len(train_df['eeg_id'].unique())} unique EEG IDs\")\nexcept Exception as e:\n    print(f\"Error loading training labels: {e}\")\n    exit()\n\n# Encode labels\nlabel_encoder = LabelEncoder()\nlabel_encoder.fit(CLASSES)\ntrain_df['label'] = label_encoder.transform(train_df['expert_consensus'])\n\n# Display original class distribution\nprint(\"Original class distribution:\")\nprint(train_df['expert_consensus'].value_counts(normalize=True))\n\n# Custom Dataset for EEG data\nclass EEGDataset(Dataset):\n    def __init__(self, eeg_ids, labels, data_path):\n        self.eeg_ids = eeg_ids\n        self.labels = labels\n        self.data_path = data_path\n    \n    def __len__(self):\n        return len(self.eeg_ids)\n    \n    def __getitem__(self, idx):\n        eeg_id = self.eeg_ids[idx]\n        eeg_path = os.path.join(self.data_path, f\"{eeg_id}.npy\")\n        try:\n            eeg_data = np.load(eeg_path).astype(np.float32)  # Shape: (19, 2500)\n            # Impute NaN/Inf with channel mean\n            if np.any(np.isnan(eeg_data)) or np.any(np.isinf(eeg_data)):\n                print(f\"Warning: NaN or Inf in EEG data for {eeg_id}, imputing with channel mean\")\n                for ch in range(eeg_data.shape[0]):\n                    channel = eeg_data[ch, :]\n                    mask = np.isnan(channel) | np.isinf(channel)\n                    if np.any(mask):\n                        channel[mask] = np.nanmean(channel)\n            # Standardize each channel\n            eeg_data = (eeg_data - np.mean(eeg_data, axis=1, keepdims=True)) / (np.std(eeg_data, axis=1, keepdims=True) + 1e-7)\n            # Transpose to (2500, 19) for LSTM: (timesteps, features)\n            eeg_data = eeg_data.T\n            label = self.labels[idx]\n            return torch.tensor(eeg_data, dtype=torch.float32), torch.tensor(label, dtype=torch.long)\n        except Exception as e:\n            print(f\"Error loading EEG data for {eeg_id}: {e}\")\n            return torch.zeros((2500, 19), dtype=torch.float32), torch.tensor(self.labels[idx], dtype=torch.long)\n\n# Load data\nprint(\"Loading data...\")\nsuccess_file_path = os.path.join(os.path.dirname(PREPROCESSED_PATH), \"success.csv\")\nif os.path.exists(success_file_path):\n    try:\n        success_df = pd.read_csv(success_file_path)\n        success_ids = set(success_df['eeg_id'].astype(str).tolist())\n        print(f\"Found success file with {len(success_ids)} successful preprocessing entries\")\n    except Exception as e:\n        print(f\"Error loading success file: {e}\")\n        success_ids = set(train_df['eeg_id'].astype(str).tolist())\nelse:\n    print(\"Success file not found, using all available EEG IDs\")\n    success_ids = set(train_df['eeg_id'].astype(str).tolist())\n\n# Filter samples based on success_ids\nvalid_samples = train_df[train_df['eeg_id'].astype(str).isin(success_ids)]\neeg_ids = valid_samples['eeg_id'].astype(str).tolist()\nlabels = valid_samples['label'].values\nprint(f\"Processing {len(eeg_ids)} valid samples\")\n\n# Train-validation split\ntrain_ids, val_ids, train_labels, val_labels = train_test_split(\n    eeg_ids, labels, test_size=0.2, stratify=labels, random_state=42\n)\nprint(f\"Training samples: {len(train_ids)}, Validation samples: {len(val_ids)}\")\n\n# Verify class distribution in training set\ntrain_class_counts = pd.Series(train_labels).value_counts(normalize=True)\ntrain_class_counts.index = [CLASSES[i] for i in train_class_counts.index]\nprint(\"\\nClass distribution in training set:\")\nprint(train_class_counts)\n\n# Create datasets and dataloaders\ntrain_dataset = EEGDataset(train_ids, train_labels, PREPROCESSED_PATH)\nval_dataset = EEGDataset(val_ids, val_labels, PREPROCESSED_PATH)\ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\nval_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)\n\n# LSTM Model with enhancements\nclass EEG_LSTM(nn.Module):\n    def __init__(self, input_size=19, hidden_size=128, num_layers=3, output_size=N_CLASSES):\n        super(EEG_LSTM, self).__init__()\n        self.conv1d = nn.Conv1d(in_channels=19, out_channels=32, kernel_size=3, padding=1)\n        self.lstm = nn.LSTM(32, hidden_size, num_layers, batch_first=True)\n        self.fc = nn.Linear(hidden_size, output_size)\n        self.dropout = nn.Dropout(0.4)  # Increased dropout\n    \n    def forward(self, x, temperature=1.5):  # Temperature for probability calibration\n        x = self.conv1d(x.transpose(1, 2)).transpose(1, 2)  # Conv1D: (batch, 2500, 32)\n        lstm_out, _ = self.lstm(x)\n        last_out = lstm_out[:, -1, :]\n        out = self.dropout(last_out)\n        out = self.fc(out) / temperature\n        return out\n\n# Initialize model, loss, and optimizer\nclass_counts = train_df['label'].value_counts().sort_index().values\nclass_weights = torch.tensor(1.0 / class_counts, dtype=torch.float32).to(DEVICE)\nmodel = EEG_LSTM(input_size=19, hidden_size=128, num_layers=3, output_size=N_CLASSES).to(DEVICE)\ncriterion = nn.CrossEntropyLoss(weight=class_weights, label_smoothing=0.1)  # Label smoothing added\noptimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-4)  # L2 regularization\nscheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=2)\n\n# Training loop\nprint(\"Training LSTM model...\")\ntrain_losses = []\nval_losses = []\ntrain_accuracies = []\nval_accuracies = []\ntrain_kls = []\nval_kls = []\nlrs = []\nbest_val_kl = float('inf')\nbest_model_path = os.path.join(MODEL_OUTPUT_PATH, \"lstm_best_model.pth\")\npatience = 3\ncounter = 0\n\nfor epoch in range(EPOCHS):\n    model.train()\n    train_loss = 0.0\n    train_correct = 0\n    train_total = 0\n    train_kl_sum = 0.0\n    \n    for batch_data, batch_labels in tqdm(train_loader, desc=f\"Epoch {epoch+1}/{EPOCHS}\"):\n        batch_data, batch_labels = batch_data.to(DEVICE), batch_labels.to(DEVICE)\n        \n        optimizer.zero_grad()\n        outputs = model(batch_data, temperature=1.5)  # Temperature scaling\n        loss = criterion(outputs, batch_labels)\n        loss.backward()\n        optimizer.step()\n        \n        train_loss += loss.item() * batch_data.size(0)\n        _, predicted = torch.max(outputs, 1)\n        train_total += batch_labels.size(0)\n        train_correct += (predicted == batch_labels).sum().item()\n        \n        # Compute KL divergence\n        pred_proba = torch.softmax(outputs, dim=1)\n        true_one_hot = torch.nn.functional.one_hot(batch_labels, N_CLASSES).float()\n        train_kl_sum += kl_divergence_torch(true_one_hot, pred_proba) * batch_data.size(0)\n    \n    train_loss /= len(train_dataset)\n    train_acc = train_correct / train_total\n    train_kl = train_kl_sum / len(train_dataset)\n    \n    # Validation\n    model.eval()\n    val_loss = 0.0\n    val_correct = 0\n    val_total = 0\n    val_kl_sum = 0.0\n    val_preds = []\n    val_true = []\n    with torch.no_grad():\n        for batch_data, batch_labels in val_loader:\n            batch_data, batch_labels = batch_data.to(DEVICE), batch_labels.to(DEVICE)\n            outputs = model(batch_data, temperature=1.5)\n            loss = criterion(outputs, batch_labels)\n            \n            val_loss += loss.item() * batch_data.size(0)\n            _, predicted = torch.max(outputs, 1)\n            val_total += batch_labels.size(0)\n            val_correct += (predicted == batch_labels).sum().item()\n            \n            # Collect predictions for confusion matrix\n            val_preds.extend(predicted.cpu().numpy())\n            val_true.extend(batch_labels.cpu().numpy())\n            \n            # Compute KL divergence\n            pred_proba = torch.softmax(outputs, dim=1)\n            true_one_hot = torch.nn.functional.one_hot(batch_labels, N_CLASSES).float()\n            val_kl_sum += kl_divergence_torch(true_one_hot, pred_proba) * batch_data.size(0)\n    \n    val_loss /= len(val_dataset)\n    val_acc = val_correct / val_total\n    val_kl = val_kl_sum / len(val_dataset)\n    \n    # Record learning rate\n    current_lr = optimizer.param_groups[0]['lr']\n    lrs.append(current_lr)\n    scheduler.step(val_kl)  # Schedule based on KL divergence\n    \n    train_losses.append(train_loss)\n    val_losses.append(val_loss)\n    train_accuracies.append(train_acc)\n    val_accuracies.append(val_acc)\n    train_kls.append(train_kl)\n    val_kls.append(val_kl)\n    \n    print(f\"Epoch {epoch+1}/{EPOCHS}\")\n    print(f\"Train Loss: {train_loss:.6f}, Train Accuracy: {train_acc:.6f}, Train KL Divergence: {train_kl:.6f}\")\n    print(f\"Val Loss: {val_loss:.6f}, Val Accuracy: {val_acc:.6f}, Val KL Divergence: {val_kl:.6f}\")\n    \n    # Early stopping based on validation KL divergence\n    if val_kl < best_val_kl:\n        best_val_kl = val_kl\n        counter = 0\n        torch.save(model.state_dict(), best_model_path)\n        print(f\"Saved best model at epoch {epoch+1} with val KL: {val_kl:.6f}\")\n    else:\n        counter += 1\n        if counter >= patience:\n            print(\"Early stopping triggered\")\n            break\n\n# Plot training and validation metrics\nplt.figure(figsize=(15, 10))\n\n# Loss plot\nplt.subplot(2, 2, 1)\nplt.plot(range(1, len(train_losses)+1), train_losses, label='Train Loss')\nplt.plot(range(1, len(val_losses)+1), val_losses, label='Val Loss')\nplt.title('Training and Validation Loss')\nplt.xlabel('Epoch')\nplt.ylabel('Loss')\nplt.legend()\n\n# Accuracy plot\nplt.subplot(2, 2, 2)\nplt.plot(range(1, len(train_accuracies)+1), train_accuracies, label='Train Accuracy')\nplt.plot(range(1, len(val_accuracies)+1), val_accuracies, label='Val Accuracy')\nplt.title('Training and Validation Accuracy')\nplt.xlabel('Epoch')\nplt.ylabel('Accuracy')\nplt.legend()\n\n# KL Divergence plot\nplt.subplot(2, 2, 3)\nplt.plot(range(1, len(train_kls)+1), train_kls, label='Train KL Divergence')\nplt.plot(range(1, len(val_kls)+1), val_kls, label='Val KL Divergence')\nplt.title('Training and Validation KL Divergence')\nplt.xlabel('Epoch')\nplt.ylabel('KL Divergence')\nplt.legend()\n\n# Learning Rate vs Epochs\nplt.subplot(2, 2, 4)\nplt.plot(range(1, len(lrs)+1), lrs, label='Learning Rate')\nplt.title('Learning Rate vs Epoch')\nplt.xlabel('Epoch')\nplt.ylabel('Learning Rate')\nplt.legend()\n\nplt.tight_layout()\nplt.savefig(os.path.join(MODEL_OUTPUT_PATH, 'training_metrics.png'), dpi=100, bbox_inches='tight')\nplt.close()\n\n# Confusion Matrix\ncm = confusion_matrix(val_true, val_preds)\nplt.figure(figsize=(10, 8))\nsns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=CLASSES, yticklabels=CLASSES)\nplt.title('Confusion Matrix (Validation Set)')\nplt.xlabel('Predicted')\nplt.ylabel('True')\nplt.savefig(os.path.join(MODEL_OUTPUT_PATH, 'confusion_matrix.png'), dpi=100, bbox_inches='tight')\nplt.close()\n\n# Load best model for submission\nmodel.load_state_dict(torch.load(best_model_path))\nmodel.eval()\n\n# Submission file creation\ndef create_prediction_file():\n    test_eeg_path = os.path.join(BASE_PATH, \"test_eegs\")\n    if not os.path.exists(test_eeg_path):\n        print(\"Warning: Test EEG path not found at\", test_eeg_path)\n        print(\"Creating dummy submission\")\n        submission_df = pd.DataFrame({'eeg_id': ['dummy_1', 'dummy_2']})\n        for cls in CLASSES:\n            submission_df[cls] = 1.0 / N_CLASSES\n        submission_path = os.path.join(TEST_PATH, \"submission.csv\")\n        submission_df.to_csv(submission_path, index=False)\n        print(\"Dummy submission file saved to:\", submission_path)\n        return submission_df\n    \n    test_files = [f.replace(\".parquet\", \"\") for f in os.listdir(test_eeg_path) if f.endswith(\".parquet\")]\n    if len(test_files) == 0:\n        print(\"No test files found in\", test_eeg_path)\n        return None\n\n    submission_df = pd.DataFrame({'eeg_id': test_files})\n    for cls in CLASSES:\n        submission_df[cls] = 0.0\n\n    print(\"Generating predictions for test data...\")\n    predictions_made = 0\n    \n    for eeg_id in tqdm(test_files[:100]):\n        try:\n            eeg_path = os.path.join(PREPROCESSED_PATH, f\"{eeg_id}.npy\")\n            if not os.path.exists(eeg_path):\n                print(f\"Test EEG file not found for {eeg_id} at {eeg_path}\")\n                raise FileNotFoundError\n            \n            eeg_data = np.load(eeg_path).astype(np.float32)  # Shape: (19, 2500)\n            if np.any(np.isnan(eeg_data)) or np.any(np.isinf(eeg_data)):\n                print(f\"Warning: NaN or Inf in test EEG data for {eeg_id}, imputing with channel mean\")\n                for ch in range(eeg_data.shape[0]):\n                    channel = eeg_data[ch, :]\n                    mask = np.isnan(channel) | np.isinf(channel)\n                    if np.any(mask):\n                        channel[mask] = np.nanmean(channel)\n            eeg_data = (eeg_data - np.mean(eeg_data, axis=1, keepdims=True)) / (np.std(eeg_data, axis=1, keepdims=True) + 1e-7)\n            eeg_data = eeg_data.T  # Shape: (2500, 19)\n            eeg_data = torch.tensor(eeg_data, dtype=torch.float32).unsqueeze(0).to(DEVICE)  # Shape: (1, 2500, 19)\n            \n            with torch.no_grad():\n                outputs = model(eeg_data, temperature=1.5)\n                probs = torch.softmax(outputs, dim=1).cpu().numpy()[0]\n            \n            for i, cls in enumerate(CLASSES):\n                submission_df.loc[submission_df['eeg_id'] == eeg_id, cls] = probs[i]\n            predictions_made += 1\n            \n        except Exception as e:\n            print(f\"Error processing test file {eeg_id}: {e}\")\n            uniform_prob = 1.0 / N_CLASSES\n            for cls in CLASSES:\n                submission_df.loc[submission_df['eeg_id'] == eeg_id, cls] = uniform_prob\n\n    submission_path = os.path.join(TEST_PATH, \"submission.csv\")\n    submission_df.to_csv(submission_path, index=False)\n    print(f\"Submission file saved to: {submission_path}\")\n    print(f\"Made predictions for {predictions_made}/{len(test_files)} test files\")\n    return submission_df\n\n# Create submission file\ntry:\n    submission_df = create_prediction_file()\nexcept Exception as e:\n    print(f\"Error creating submission file: {e}\")\n\nprint(\"\\nTraining completed!\")\nprint(f\"Best validation KL divergence: {best_val_kl:.6f}\")\nprint(f\"Last epoch train accuracy: {train_accuracies[-1]:.6f}, validation accuracy: {val_accuracies[-1]:.6f}\")\nprint(f\"Last epoch train KL divergence: {train_kls[-1]:.6f}, validation KL divergence: {val_kls[-1]:.6f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-30T16:22:40.787046Z","iopub.execute_input":"2025-05-30T16:22:40.788008Z","iopub.status.idle":"2025-05-30T18:47:54.430955Z","shell.execute_reply.started":"2025-05-30T16:22:40.787970Z","shell.execute_reply":"2025-05-30T18:47:54.430198Z"}},"outputs":[],"execution_count":null}]}