{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":106809,"databundleVersionId":13056355,"sourceType":"competition"},{"sourceId":14180943,"sourceType":"datasetVersion","datasetId":9040339}],"dockerImageVersionId":31236,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## 🧠 Brain-To-Text — Extensions\n\nThis notebook builds on the **Brain-To-Text Baseline** by **Piotr Jurczyk** 🔗 https://www.kaggle.com/code/piotrjurczyk/brain-to-text-baseline\n\n- The dataset comes from a **saved execution of the original notebook** (best model + train/val/test predictions), enabling **fast EDA** without rerunning long training.\n- **Data loading, preprocessing, and splits** are kept identical to ensure full reproducibility.\n\n\n### ✨ Added Contributions\n- **Beam search decoding**\n- **Data augmentation**\n\n### 🚀 Next Improvements that I suggest\n- Fine-tune **beam search parameters**\n- Integrate a **Transformer / Language Model** during decoding to:\n  - Select the best sentence among beam candidates\n  - Perform **word-by-word orthographic correction**  \n    *(tested locally with **LLaMA 3 2B**, carefully constrained to avoid reformulation)*\n\n\n\n---\n\n\n## Note to the Kaggle community 🎄\n\n\nI am taking a short **Kaggle break for Christmas**, so I will not actively iterate on this notebook for a while.\n\n\nAnyone interested is **very welcome to reuse, adapt, or extend this code** for their own experiments or submissions.\n\n\nHappy holidays and good luck! 🎅✨","metadata":{}},{"cell_type":"code","source":"\"\"\"\n================================================================================\nBRAIN-TO-TEXT '25 : TRAINING WITH DATA AUGMENTATION + TEST INFERENCE\n================================================================================\nPipeline:\n1. Model training with strong data augmentation (60 epochs, ~2h GPU)\n2. Beam search inference on TEST with fixed params (bw=50, prune=-12, min=-8)\n3. Generation of submission_final.csv for Kaggle\n\nEstimated time: 2h15 (2h training + 15min test)\nExpected gain: WER 29.5% → 25–27% (data augmentation)\n================================================================================\n\"\"\"\n\n# ============================================================================\n# IMPORTS\n# ============================================================================\n\nimport os\nimport torch\nimport torch.nn as nn\nfrom torch.nn.utils.rnn import pad_sequence, pack_padded_sequence, pad_packed_sequence\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import OneCycleLR\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nimport h5py\nfrom glob import glob\nimport json\nfrom datetime import datetime\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# Dependency installation (Kaggle-compatible)\nprint(\"Installing dependencies...\")\nos.system('pip install -q jiwer')\nos.system('pip install -q pyctcdecode')\n\nprint(f\"\\n{'#'*80}\")\nprint(f\"##  BRAIN-TO-TEXT '25 - TRAIN + DATA AUGMENTATION\")\nprint(f\"{'#'*80}\\n\")\n\n\n# ============================================================================\n# DATA AUGMENTATION MODULES\n# ============================================================================\n\nclass SpecAugmentStrong(nn.Module):\n    \"\"\"Strong SpecAugment implementation (time + frequency masking).\"\"\"\n    def __init__(self, prob=0.8, time_mask_param=60, feature_mask_param=50,\n                 num_time_masks=2, num_feat_masks=2):\n        super().__init__()\n        self.prob = prob\n        self.time_mask_param = time_mask_param\n        self.feature_mask_param = feature_mask_param\n        self.num_time_masks = num_time_masks\n        self.num_feat_masks = num_feat_masks\n\n    def forward(self, x):\n        # Apply augmentation only during training and with given probability\n        if not self.training or torch.rand(1) > self.prob:\n            return x\n        \n        b, c, t = x.size()\n        x_aug = x.clone()\n        \n        # Time masking\n        for _ in range(self.num_time_masks):\n            mask_len = torch.randint(0, min(self.time_mask_param, t // 2), (1,)).item()\n            if mask_len > 0 and t > mask_len:\n                t0 = torch.randint(0, t - mask_len, (1,)).item()\n                x_aug[:, :, t0:t0 + mask_len] = 0\n        \n        # Feature masking\n        for _ in range(self.num_feat_masks):\n            mask_feat = torch.randint(0, min(self.feature_mask_param, c // 2), (1,)).item()\n            if mask_feat > 0 and c > mask_feat:\n                f0 = torch.randint(0, c - mask_feat, (1,)).item()\n                x_aug[:, f0:f0 + mask_feat, :] = 0\n        \n        return x_aug\n\n\nclass NoiseInjection(nn.Module):\n    \"\"\"Gaussian noise injection.\"\"\"\n    def __init__(self, prob=0.5, noise_std=0.05):\n        super().__init__()\n        self.prob = prob\n        self.noise_std = noise_std\n    \n    def forward(self, x):\n        if not self.training or torch.rand(1) > self.prob:\n            return x\n        noise = torch.randn_like(x) * self.noise_std\n        return x + noise\n\n\nclass TimeWarping(nn.Module):\n    \"\"\"Temporal warping augmentation.\"\"\"\n    def __init__(self, prob=0.4, warp_factor_range=(0.9, 1.1)):\n        super().__init__()\n        self.prob = prob\n        self.warp_min, self.warp_max = warp_factor_range\n    \n    def forward(self, x):\n        if not self.training or torch.rand(1) > self.prob:\n            return x\n        \n        b, c, t = x.size()\n        warp_factor = self.warp_min + torch.rand(1).item() * (self.warp_max - self.warp_min)\n        new_t = int(t * warp_factor)\n        new_t = max(10, min(new_t, t * 2))\n        \n        # Interpolate along time axis\n        x_warped = torch.nn.functional.interpolate(\n            x, size=new_t, mode='linear', align_corners=False\n        )\n        \n        # Crop or pad back to original length\n        if new_t < t:\n            x_warped = torch.nn.functional.pad(x_warped, (0, t - new_t))\n        elif new_t > t:\n            x_warped = x_warped[:, :, :t]\n        \n        return x_warped\n\n\n# ============================================================================\n# MODEL WITH DATA AUGMENTATION\n# ============================================================================\n\nclass BrainCTCModel(nn.Module):\n    \"\"\"CTC-based Brain-to-Text model with strong data augmentation.\"\"\"\n    def __init__(self, input_dim=512, hidden_dim=512, num_layers=3, \n                 vocab_size=50, dropout=0.5):\n        super().__init__()\n        \n        # Augmentation pipeline (applied only during training)\n        self.augmentations = nn.Sequential(\n            NoiseInjection(prob=0.5, noise_std=0.05),\n            TimeWarping(prob=0.4, warp_factor_range=(0.9, 1.1)),\n            SpecAugmentStrong(\n                prob=0.8,\n                time_mask_param=60,\n                feature_mask_param=50,\n                num_time_masks=2,\n                num_feat_masks=2\n            )\n        )\n        \n        # Convolutional feature extractor\n        self.cnn = nn.Sequential(\n            nn.Conv1d(input_dim, 256, kernel_size=11, stride=2, padding=5),\n            nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(dropout),\n            nn.Conv1d(256, 256, kernel_size=5, stride=2, padding=2),\n            nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(dropout)\n        )\n        \n        # Bidirectional LSTM encoder\n        self.lstm = nn.LSTM(\n            256, hidden_dim, num_layers,\n            batch_first=True,\n            bidirectional=True,\n            dropout=dropout if num_layers > 1 else 0\n        )\n        \n        # Output projection to vocabulary size\n        self.fc = nn.Linear(hidden_dim * 2, vocab_size)\n    \n    def forward(self, x, lengths):\n        # Input: (B, T, C) → (B, C, T)\n        x = x.transpose(1, 2)\n        \n        # Apply augmentations during training only\n        if self.training:\n            x = self.augmentations(x)\n        \n        # CNN feature extraction\n        x = self.cnn(x)\n        x = x.transpose(1, 2)\n        \n        # Update sequence lengths after strided convolutions\n        cnn_lengths = (lengths.cpu() // 4).clamp(min=1)\n        \n        # Pack sequences for LSTM\n        x_packed = pack_padded_sequence(\n            x, cnn_lengths, batch_first=True, enforce_sorted=False\n        )\n        lstm_out, _ = self.lstm(x_packed)\n        lstm_out, _ = pad_packed_sequence(lstm_out, batch_first=True)\n        \n        # Linear projection + log-softmax for CTC\n        logits = self.fc(lstm_out)\n        log_probs = torch.log_softmax(logits, dim=-1)\n        \n        # Output shape: (T, B, V)\n        return log_probs.transpose(0, 1), cnn_lengths\n\n\n# ============================================================================\n# DATASET\n# ============================================================================\n\nclass BrainDataset(Dataset):\n    \"\"\"Brain-to-Text dataset loader.\"\"\"\n    def __init__(self, split='train'):\n        self.samples = []\n        \n        data_dir = '/kaggle/input/brain-to-text-25/t15_copyTask_neuralData/hdf5_data_final/'\n        pattern = f'{data_dir}/**/data_{split}.hdf5'\n        files = sorted(glob(pattern, recursive=True))\n        \n        print(f\"Loading {split} split: {len(files)} files\")\n        \n        for filepath in tqdm(files, desc=split):\n            try:\n                with h5py.File(filepath, 'r') as f:\n                    for trial_key in f.keys():\n                        trial = f[trial_key]\n                        \n                        # Neural features\n                        neural = trial['input_features'][:]\n                        n_steps = trial.attrs['n_time_steps']\n                        neural = neural[:n_steps]\n                        neural = (neural - neural.mean()) / (neural.std() + 1e-8)\n                        \n                        # Sentence label\n                        sentence = trial.attrs.get('sentence_label', '')\n                        if isinstance(sentence, bytes):\n                            sentence = sentence.decode('utf-8')\n                        \n                        self.samples.append({\n                            'neural': torch.FloatTensor(neural),\n                            'sentence': sentence.lower(),\n                            'length': len(neural)\n                        })\n            except Exception:\n                continue\n        \n        print(f\"✓ {len(self.samples)} samples loaded\")\n    \n    def __len__(self):\n        return len(self.samples)\n    \n    def __getitem__(self, idx):\n        return self.samples[idx]\n\n\n# ============================================================================\n# COLLATE FUNCTION\n# ============================================================================\n\ndef collate_fn(batch, char2idx):\n    \"\"\"Pads features and encodes text targets for CTC training.\"\"\"\n    features = [s['neural'] for s in batch]\n    sentences = [s['sentence'] for s in batch]\n    lengths = torch.LongTensor([s['length'] for s in batch])\n    \n    # Pad neural features\n    features_padded = pad_sequence(features, batch_first=True)\n    \n    # Encode sentences into character indices\n    targets = []\n    target_lengths = []\n    for sentence in sentences:\n        encoded = [char2idx.get(c, char2idx.get('<BLANK>')) for c in sentence]\n        targets.extend(encoded)\n        target_lengths.append(len(encoded))\n    \n    targets = torch.LongTensor(targets)\n    target_lengths = torch.LongTensor(target_lengths)\n    \n    return features_padded, targets, lengths, target_lengths\n\n\n# ============================================================================\n# VOCABULARY\n# ============================================================================\n\ndef build_vocabulary():\n    \"\"\"Builds character-level vocabulary.\"\"\"\n    # Special characters\n    chars = [' ', '!', \"'\", ',', '-', '.', ';', '?', '[', ']']\n    \n    # Letters a–z\n    chars += [chr(i) for i in range(ord('a'), ord('z') + 1)]\n    \n    # Curly apostrophe (Unicode U+2019)\n    chars += ['\\u2019']\n    \n    char2idx = {'<BLANK>': 0}\n    for i, c in enumerate(chars, 1):\n        char2idx[c] = i\n    \n    idx2char = {v: k for k, v in char2idx.items()}\n    \n    return char2idx, idx2char\n\n\n# ============================================================================\n# TRAINING\n# ============================================================================\n\ndef train_model():\n    \"\"\"Full training loop with validation and checkpointing.\"\"\"\n    print(f\"\\n{'='*80}\")\n    print(\"TRAINING WITH DATA AUGMENTATION\")\n    print(f\"{'='*80}\\n\")\n    \n    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n    print(f\"Device: {device}\\n\")\n    \n    # Vocabulary\n    char2idx, idx2char = build_vocabulary()\n    print(f\"Vocabulary size: {len(char2idx)} characters\\n\")\n    \n    # Datasets and loaders\n    train_dataset = BrainDataset('train')\n    val_dataset = BrainDataset('val')\n    \n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=64,\n        shuffle=True,\n        collate_fn=lambda b: collate_fn(b, char2idx),\n        num_workers=2,\n        pin_memory=True\n    )\n    \n    val_loader = DataLoader(\n        val_dataset,\n        batch_size=64,\n        shuffle=False,\n        collate_fn=lambda b: collate_fn(b, char2idx),\n        num_workers=2,\n        pin_memory=True\n    )\n    \n    # Model\n    model = BrainCTCModel(\n        input_dim=512,\n        hidden_dim=512,\n        num_layers=3,\n        vocab_size=len(char2idx),\n        dropout=0.5\n    ).to(device)\n    \n    # Optimizer and scheduler\n    optimizer = AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)\n    scheduler = OneCycleLR(\n        optimizer,\n        max_lr=1e-3,\n        steps_per_epoch=len(train_loader),\n        epochs=60\n    )\n    \n    # CTC loss\n    ctc_loss = nn.CTCLoss(blank=0, zero_infinity=True)\n    \n    print(\"Configuration:\")\n    print(\"  - Epochs: 60\")\n    print(\"  - Batch size: 64\")\n    print(\"  - Learning rate: 1e-3\")\n    print(\"  - Augmentations: SpecAugment + Noise + TimeWarp\")\n    print(f\"{'='*80}\\n\")\n    \n    best_wer = float('inf')\n    \n    # Training loop\n    for epoch in range(60):\n        model.train()\n        train_loss = 0.0\n        \n        pbar = tqdm(train_loader, desc=f'Epoch {epoch+1}/60')\n        for features, targets, lengths, target_lengths in pbar:\n            features = features.to(device)\n            targets = targets.to(device)\n            \n            optimizer.zero_grad()\n            \n            log_probs, cnn_lengths = model(features, lengths)\n            loss = ctc_loss(log_probs, targets, cnn_lengths, target_lengths)\n            \n            loss.backward()\n            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n            optimizer.step()\n            scheduler.step()\n            \n            train_loss += loss.item()\n            pbar.set_postfix({'loss': f'{loss.item():.4f}'})\n        \n        # Validation every 10 epochs\n        if (epoch + 1) % 10 == 0:\n            model.eval()\n            from jiwer import wer\n            \n            all_preds = []\n            all_targets = []\n            \n            with torch.no_grad():\n                for features, _, lengths, _ in val_loader:\n                    features = features.to(device)\n                    log_probs, cnn_lengths = model(features, lengths)\n                    \n                    for b in range(log_probs.size(1)):\n                        logits = log_probs[:cnn_lengths[b], b].cpu().numpy()\n                        pred_indices = np.argmax(logits, axis=1)\n                        \n                        # Greedy CTC decoding\n                        decoded = []\n                        prev = None\n                        for idx in pred_indices:\n                            if idx != 0 and idx != prev:\n                                decoded.append(idx2char.get(idx, ''))\n                            prev = idx\n                        \n                        all_preds.append(''.join(decoded))\n                        all_targets.append(\n                            val_dataset.samples[len(all_preds)-1]['sentence']\n                        )\n            \n            val_wer = wer(\n                [t.lower() for t in all_targets],\n                [p.lower() for p in all_preds]\n            ) * 100\n            \n            print(f\"\\nEpoch {epoch+1} - WER: {val_wer:.2f}%\")\n            \n            # Save best model\n            if val_wer < best_wer:\n                best_wer = val_wer\n                torch.save({\n                    'epoch': epoch + 1,\n                    'model_state_dict': model.state_dict(),\n                    'optimizer_state_dict': optimizer.state_dict(),\n                    'wer': val_wer,\n                    'char2idx': char2idx,\n                    'idx2char': idx2char,\n                    'config': {\n                        'hidden_dim': 512,\n                        'num_layers': 3,\n                        'dropout': 0.5\n                    }\n                }, 'best_model_augmented.pt')\n                print(f\"✓ Model saved (WER: {val_wer:.2f}%)\")\n    \n    print(f\"\\n{'='*80}\")\n    print(\"TRAINING COMPLETED\")\n    print(f\"Best WER: {best_wer:.2f}%\")\n    print(f\"{'='*80}\\n\")\n    \n    return char2idx, idx2char\n\n\n# ============================================================================\n# BEAM SEARCH TEST INFERENCE\n# ============================================================================\n\ndef inference_test(char2idx):\n    \"\"\"Runs beam search inference on the TEST set and creates submission file.\"\"\"\n    print(f\"\\n{'='*80}\")\n    print(\"TEST INFERENCE WITH BEAM SEARCH\")\n    print(f\"{'='*80}\\n\")\n    \n    from pyctcdecode import build_ctcdecoder\n    \n    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n    \n    # Load trained model\n    checkpoint = torch.load('best_model_augmented.pt', map_location=device)\n    \n    model = BrainCTCModel(\n        input_dim=512,\n        hidden_dim=512,\n        num_layers=3,\n        vocab_size=len(char2idx),\n        dropout=0.5\n    ).to(device)\n    \n    model.load_state_dict(checkpoint['model_state_dict'])\n    model.eval()\n    \n    print(f\"✓ Model loaded (val WER: {checkpoint['wer']:.2f}%)\")\n    \n    # Build CTC decoder vocabulary\n    vocab_list = [''] * len(char2idx)\n    for char, idx in char2idx.items():\n        vocab_list[idx] = '' if char == '<BLANK>' else char\n    \n    decoder = build_ctcdecoder(labels=vocab_list)\n    print(\"✓ Decoder created (beam_width=50)\\n\")\n    \n    # Load TEST set\n    test_dataset = BrainDataset('test')\n    \n    all_preds = []\n    print(\"Running inference on TEST set...\")\n    \n    with torch.no_grad():\n        for sample in tqdm(test_dataset.samples):\n            features = sample['neural'].unsqueeze(0).to(device)\n            length = torch.LongTensor([sample['length']])\n            \n            log_probs, cnn_lengths = model(features, length)\n            logits = log_probs[:cnn_lengths[0], 0].cpu().numpy()\n            probs = np.exp(logits)\n            \n            try:\n                beam_results = decoder.decode_beams(\n                    probs,\n                    beam_width=50,\n                    beam_prune_logp=-12.0,\n                    token_min_logp=-8.0\n                )\n                pred_text = beam_results[0][0] if beam_results else \"\"\n            except Exception:\n                pred_text = \"\"\n            \n            all_preds.append(pred_text)\n    \n    # Save Kaggle submission\n    with open('submission.csv', 'w', encoding='utf-8') as f:\n        f.write('id,text\\n')\n        for i, pred in enumerate(all_preds):\n            pred_clean = pred.replace('\"', '\"\"')\n            f.write(f'{i},\"{pred_clean}\"\\n')\n    \n    print(f\"\\n✓ Submission saved: submission.csv ({len(all_preds)} predictions)\")\n    \n    print(\"\\nExamples:\")\n    for i in range(min(5, len(all_preds))):\n        print(f\"  {i}: {all_preds[i][:60]}\")\n    \n    print(f\"\\n{'='*80}\")\n    print(\"🚀 Ready for Kaggle submission!\")\n    print(f\"{'='*80}\\n\")\n\n\n# ============================================================================\n# MAIN ENTRY POINT\n# ============================================================================\n\ndef main():\n    start = datetime.now()\n    \n    # Training phase\n    char2idx, idx2char = train_model()\n    \n    # Test inference phase\n    inference_test(char2idx)\n    \n    duration = (datetime.now() - start).total_seconds() / 60\n    print(f\"\\n⏱️  Total runtime: {duration:.1f} min\")\n\n\nif __name__ == \"__main__\":\n    try:\n        main()\n    except Exception as e:\n        print(f\"\\n❌ ERROR: {e}\")\n        import traceback\n        traceback.print_exc()\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-18T17:27:53.612276Z","iopub.execute_input":"2025-12-18T17:27:53.613014Z","iopub.status.idle":"2025-12-18T18:47:58.948939Z","shell.execute_reply.started":"2025-12-18T17:27:53.612987Z","shell.execute_reply":"2025-12-18T18:47:58.947942Z"}},"outputs":[],"execution_count":null}]}