{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":106809,"databundleVersionId":13056355,"sourceType":"competition"}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-22T07:44:26.067627Z","iopub.execute_input":"2025-12-22T07:44:26.06794Z","iopub.status.idle":"2025-12-22T07:44:28.004176Z","shell.execute_reply.started":"2025-12-22T07:44:26.067922Z","shell.execute_reply":"2025-12-22T07:44:28.003592Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Install required packages (for Kaggle environment)\n!pip install jiwer pyctcdecode scikit-learn\n\nimport os\nimport h5py\nimport numpy as np\nfrom pathlib import Path\nfrom tqdm import tqdm\nfrom sklearn.decomposition import IncrementalPCA\nimport torch\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.nn.utils.rnn import pad_sequence\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.nn.functional import gelu, silu  # activation functions\nfrom jiwer import wer  # for WER metric\n\n# Use GPU if available\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")\n\n# Configuration\nCONFIG = {\n    'data_dir': '/kaggle/input/brain-to-text-25/t15_copyTask_neuralData/hdf5_data_final',\n    'batch_size': 16,            # GPU memory-friendly batch size\n    'num_epochs': 30,            # number of training epochs\n    'input_dim': 128,            # post-PCA feature dimension\n    'hidden_dim': 512,           # hidden dimension in Conformer\n    'ff_expansion': 4,           # Feed-forward expansion factor\n    'num_heads': 8,              # Multi-head attention heads\n    'num_layers': 3,             # number of Conformer blocks\n    'conv_kernel_size': 15,      # convolution kernel size in conv module\n    'dropout': 0.1,              # dropout probability\n    'learning_rate': 3e-4,       # peak learning rate for OneCycleLR\n    'blank_index': 0,            # index of CTC blank token in vocab\n    'time_mask_width': 20,       # max width for time masking augmentation\n}\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T07:44:28.005808Z","iopub.execute_input":"2025-12-22T07:44:28.008271Z","iopub.status.idle":"2025-12-22T07:44:40.473216Z","shell.execute_reply.started":"2025-12-22T07:44:28.008246Z","shell.execute_reply":"2025-12-22T07:44:40.472462Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def load_hdf5_files(file_list):\n    data = {'neural': [], 'n_steps': [], 'sent': []}\n    for filepath in file_list:\n        with h5py.File(filepath, 'r') as hf:\n            for trial_name in hf:\n                trial = hf[trial_name]\n                neural = np.array(trial['input_features'])  # ← fix here\n                n_steps = trial.attrs['n_time_steps']\n                sentence = trial.attrs.get('sentence_label')\n                if sentence is None:\n                    sentence = \"\"\n                elif isinstance(sentence, bytes):\n                    sentence = sentence.decode('utf-8')\n                data['neural'].append(neural)\n                data['n_steps'].append(n_steps)\n                data['sent'].append(sentence.lower())\n    print(f\"Loaded {len(data['neural'])} samples from {len(file_list)} files.\")\n    return data\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T07:44:40.474984Z","iopub.execute_input":"2025-12-22T07:44:40.475398Z","iopub.status.idle":"2025-12-22T07:44:40.480655Z","shell.execute_reply.started":"2025-12-22T07:44:40.475374Z","shell.execute_reply":"2025-12-22T07:44:40.480059Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# Find all data_train/val/test files\ndata_path = Path(CONFIG['data_dir'])\ntrain_files = sorted(data_path.glob(\"**/data_train.hdf5\"))\nval_files   = sorted(data_path.glob(\"**/data_val.hdf5\"))\ntest_files  = sorted(data_path.glob(\"**/data_test.hdf5\"))\n\ntrain_data = load_hdf5_files(train_files)\nval_data   = load_hdf5_files(val_files)\ntest_data  = load_hdf5_files(test_files)\n\n# Fit Incremental PCA on the training neural data (stack frames to one batch)\npca = IncrementalPCA(n_components=CONFIG['input_dim'], batch_size=1000)\nfor arr in train_data['neural']:\n    pca.partial_fit(arr)\nprint(\"PCA training complete.\")\n\n# Apply PCA transform to all datasets\ndef apply_pca(data_dict):\n    for i in range(len(data_dict['neural'])):\n        arr = data_dict['neural'][i]\n        data_dict['neural'][i] = pca.transform(arr)\n\napply_pca(train_data)\napply_pca(val_data)\napply_pca(test_data)\nprint(\"Applied PCA to train/val/test features.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T07:44:40.48152Z","iopub.execute_input":"2025-12-22T07:44:40.481788Z","iopub.status.idle":"2025-12-22T08:09:23.8705Z","shell.execute_reply.started":"2025-12-22T07:44:40.481768Z","shell.execute_reply":"2025-12-22T08:09:23.869701Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import h5py\nfrom pathlib import Path\n\n# Pick one example train file to inspect\nsample_file = list(Path(\"/kaggle/input/brain-to-text-25/t15_copyTask_neuralData/hdf5_data_final/\").rglob(\"*.hdf5\"))[0]\n\nwith h5py.File(sample_file, 'r') as f:\n    print(\"File:\", sample_file)\n    for trial_name in f:\n        print(f\" Trial: {trial_name}\")\n        trial = f[trial_name]\n        for key in trial:\n            print(f\"  Dataset: {key}\")\n        for attr in trial.attrs:\n            print(f\"  Attribute: {attr} -> {trial.attrs[attr]}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T08:09:23.87156Z","iopub.execute_input":"2025-12-22T08:09:23.872251Z","iopub.status.idle":"2025-12-22T08:09:24.208893Z","shell.execute_reply.started":"2025-12-22T08:09:23.87223Z","shell.execute_reply":"2025-12-22T08:09:24.208284Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class BrainDataset(Dataset):\n    def __init__(self, data, char2idx=None, normalize=True):\n        self.neural = data['neural']\n        self.lengths = data['n_steps']\n        self.sentences = data['sent']\n        self.normalize = normalize\n        \n        # Build or use existing vocabulary (include '<BLANK>' as index 0)\n        if char2idx is None:\n            chars = set(\"\".join(self.sentences))\n            chars.discard(\"\")  # remove empty string if any\n            chars = sorted(chars)\n            self.char2idx = {'<BLANK>': CONFIG['blank_index']}\n            for i,ch in enumerate(chars, start=1):\n                self.char2idx[ch] = i\n        else:\n            self.char2idx = char2idx\n        self.idx2char = {v:k for k,v in self.char2idx.items()}\n        self.vocab_size = len(self.char2idx)\n        \n    def __len__(self):\n        return len(self.neural)\n    \n    def __getitem__(self, idx):\n        x = self.neural[idx]  # numpy array [time, input_dim]\n        n_steps = self.lengths[idx]\n        x = x[:n_steps]  # remove any padding frames\n        \n        # Normalize per-sample (zero mean, unit std)\n        if self.normalize:\n            mean = x.mean()\n            std  = x.std() + 1e-8\n            x = (x - mean) / std\n        \n        # Convert to torch Tensor\n        x = torch.FloatTensor(x)  # shape (T, input_dim)\n        \n        # Apply Gaussian noise augmentation (during training only)\n        # (We will apply in training loop if needed)\n        # Apply time masking augmentation:\n        # Mask a random time segment to 0\n        if self.normalize and np.random.rand() < 0.5:\n            t = x.size(0)\n            mask_len = np.random.randint(0, CONFIG['time_mask_width']+1)\n            if t > 0 and mask_len > 0:\n                start = np.random.randint(0, t - mask_len + 1)\n                x[start:start+mask_len, :] = 0.0\n        \n        # Prepare target indices (CTC does not require <BLANK> in targets)\n        sentence = self.sentences[idx] or \"\"\n        target = [ self.char2idx.get(ch, CONFIG['blank_index']) for ch in sentence ]\n        target = torch.LongTensor(target) if target else torch.LongTensor([])\n        \n        return {\n            'neural': x, \n            'length': x.size(0),\n            'target': target,\n            'target_length': len(target),\n            'sentence': sentence\n        }\nMAX_LENGTH = 800  # set to 600 if needed\n\ndef collate_fn(batch):\n    \"\"\"\n    batch: list of dicts {'neural': ..., 'sent': ...}\n    Returns padded tensors: inputs, targets, lengths, target_lengths\n    \"\"\"\n    batch_inputs = []\n    input_lengths = []\n    batch_targets = []\n    target_lengths = []\n\n    for sample in batch:\n        neural = torch.tensor(sample['neural'], dtype=torch.float32)\n\n        # 🔧 Truncate if too long\n        if neural.shape[0] > MAX_LENGTH:\n            neural = neural[:MAX_LENGTH, :]\n\n        batch_inputs.append(neural)\n        input_lengths.append(neural.shape[0])\n\n        target = torch.tensor(sample['target'], dtype=torch.int64)\n        batch_targets.append(target)\n        target_lengths.append(len(target))\n\n    # Pad sequences\n    padded_inputs = nn.utils.rnn.pad_sequence(batch_inputs, batch_first=True)\n    padded_targets = torch.cat(batch_targets)\n\n    return padded_inputs, torch.tensor(input_lengths), padded_targets, torch.tensor(target_lengths)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T08:09:24.209661Z","iopub.execute_input":"2025-12-22T08:09:24.209918Z","iopub.status.idle":"2025-12-22T08:09:24.22158Z","shell.execute_reply.started":"2025-12-22T08:09:24.209901Z","shell.execute_reply":"2025-12-22T08:09:24.220921Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Build datasets; include train vocab in val/test for consistency\ntrain_ds = BrainDataset(train_data, normalize=True)\nchar2idx = train_ds.char2idx\nidx2char = train_ds.idx2char\n\nval_ds   = BrainDataset(val_data, char2idx=char2idx, normalize=True)\ntest_ds  = BrainDataset(test_data, char2idx=char2idx, normalize=True)  # no labels anyway\n\ntrain_loader = DataLoader(train_ds, batch_size=CONFIG['batch_size'], shuffle=True, collate_fn=collate_fn)\nval_loader   = DataLoader(val_ds, batch_size=CONFIG['batch_size'], shuffle=False, collate_fn=collate_fn)\nprint(f\"Vocab size (including blank): {len(char2idx)}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T08:09:24.222413Z","iopub.execute_input":"2025-12-22T08:09:24.22295Z","iopub.status.idle":"2025-12-22T08:09:24.245684Z","shell.execute_reply.started":"2025-12-22T08:09:24.222932Z","shell.execute_reply":"2025-12-22T08:09:24.245053Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class FeedForward(nn.Module):\n    \"\"\"Position-wise Feed-Forward layer (2 linear layers with SiLU and dropout).\"\"\"\n    def __init__(self, dim, expansion_factor=4, dropout=0.1):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(dim, dim * expansion_factor),\n            nn.SiLU(),\n            nn.Dropout(dropout),\n            nn.Linear(dim * expansion_factor, dim),\n            nn.Dropout(dropout)\n        )\n    def forward(self, x):\n        return self.net(x)\n\nclass ConvolutionModule(nn.Module):\n    \"\"\"Convolution Module in Conformer block (GLU, depthwise conv, etc.).\"\"\"\n    def __init__(self, dim, kernel_size=15, dropout=0.1):\n        super().__init__()\n        self.layer_norm = nn.LayerNorm(dim)\n        self.pointwise_conv1 = nn.Conv1d(dim, dim*2, kernel_size=1)\n        self.depthwise_conv = nn.Conv1d(dim, dim, kernel_size=kernel_size, groups=dim,\n                                        padding=(kernel_size-1)//2)\n        self.batch_norm = nn.BatchNorm1d(dim)\n        self.activation = nn.SiLU()\n        self.pointwise_conv2 = nn.Conv1d(dim, dim, kernel_size=1)\n        self.dropout = nn.Dropout(dropout)\n    def forward(self, x):\n        \"\"\"\n        x: (B, T, D)\n        returns: (B, T, D) after conv and residual\n        \"\"\"\n        res = x\n        # LayerNorm over feature dim\n        x = self.layer_norm(x)\n        # Conv1d expects (B, D, T)\n        x = x.transpose(1, 2)\n        # Pointwise conv + GLU\n        x = self.pointwise_conv1(x)  # (B, 2*D, T)\n        x = F.glu(x, dim=1)      # (B, D, T)\n        # Depthwise conv + BN + SiLU\n        x = self.depthwise_conv(x)\n        x = self.batch_norm(x)\n        x = self.activation(x)\n        # Pointwise conv\n        x = self.pointwise_conv2(x)\n        x = self.dropout(x)\n        x = x.transpose(1, 2)  # back to (B, T, D)\n        return res + x\n\nclass ConformerBlock(nn.Module):\n    \"\"\"One Conformer block (FF - MHA - Conv - FF, with residuals).\"\"\"\n    def __init__(self, dim, num_heads, ff_expansion, conv_kernel, dropout=0.1):\n        super().__init__()\n        self.ff1 = FeedForward(dim, expansion_factor=ff_expansion, dropout=dropout)\n        self.attn = nn.MultiheadAttention(dim, num_heads, dropout=dropout, batch_first=True)\n        self.conv = ConvolutionModule(dim, kernel_size=conv_kernel, dropout=dropout)\n        self.ff2 = FeedForward(dim, expansion_factor=ff_expansion, dropout=dropout)\n        # Layer norms for pre-norm, and final post-norm\n        self.norm_ff1 = nn.LayerNorm(dim)\n        self.norm_attn = nn.LayerNorm(dim)\n        self.norm_conv = nn.LayerNorm(dim)\n        self.norm_ff2 = nn.LayerNorm(dim)\n        self.norm_final = nn.LayerNorm(dim)\n    \n    def forward(self, x, attn_mask=None):\n        # x: (B, T, D)\n        # Feed-forward (1st half, Macaron)\n        res = x\n        x = self.norm_ff1(x)\n        x = res + 0.5 * self.ff1(x)\n        # Self-attention\n        res = x\n        x = self.norm_attn(x)\n        # MultiheadAttention with shape (B,T,D)\n        attn_out, _ = self.attn(x, x, x, attn_mask=attn_mask)\n        x = res + attn_out\n        # Convolution module\n        res = x\n        x = self.norm_conv(x)\n        x = res + self.conv(x)\n        # Feed-forward (2nd half)\n        res = x\n        x = self.norm_ff2(x)\n        x = res + 0.5 * self.ff2(x)\n        # Final LayerNorm\n        x = self.norm_final(x)\n        return x\n\nclass ConformerCTC(nn.Module):\n    \"\"\"Conformer Encoder for CTC (no decoder, outputs char probabilities).\"\"\"\n    def __init__(self, input_dim, hidden_dim, vocab_size, \n                 num_layers=3, num_heads=8, ff_expansion=4, \n                 conv_kernel=31, dropout=0.1):\n        super().__init__()\n        self.input_proj = nn.Linear(input_dim, hidden_dim)\n        self.pos_enc = nn.Sequential(\n            nn.Dropout(dropout),\n            # Positional encoding could be added here if needed (sinusoidal/learnable)\n        )\n        self.layers = nn.ModuleList([\n            ConformerBlock(hidden_dim, num_heads, ff_expansion, conv_kernel, dropout)\n            for _ in range(num_layers)\n        ])\n        self.fc_out = nn.Linear(hidden_dim, vocab_size)  # to character logits\n\n    def forward(self, x, lengths=None):\n        \"\"\"\n        x: (B, T, input_dim)\n        lengths: (B,) lengths of each sequence\n        Returns: log_probs (T, B, C), output_lengths (no change here)\n        \"\"\"\n        # Input projection\n        x = self.input_proj(x)  # (B, T, hidden_dim)\n        x = self.pos_enc(x)     # (B, T, hidden_dim)\n        \n        # NOTE: For padded input, one could create an attention mask here to ignore padding in attention.\n        attn_mask = None\n        \n        # Pass through Conformer layers\n        for layer in self.layers:\n            x = layer(x, attn_mask=attn_mask)\n        \n        # Final linear + log softmax\n        logits = self.fc_out(x)             # (B, T, C)\n        log_probs = nn.functional.log_softmax(logits, dim=-1)\n        # Permute to (T, B, C) for CTC loss\n        log_probs = log_probs.transpose(0, 1)\n        return log_probs, lengths\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T08:09:24.246362Z","iopub.execute_input":"2025-12-22T08:09:24.246638Z","iopub.status.idle":"2025-12-22T08:09:24.276795Z","shell.execute_reply.started":"2025-12-22T08:09:24.246621Z","shell.execute_reply":"2025-12-22T08:09:24.275698Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Initialize model, loss, optimizer, scheduler\nvocab_size = len(char2idx)\nmodel = ConformerCTC(input_dim=CONFIG['input_dim'], hidden_dim=CONFIG['hidden_dim'],\n                    vocab_size=vocab_size, num_layers=CONFIG['num_layers'],\n                    num_heads=CONFIG['num_heads'], ff_expansion=CONFIG['ff_expansion'],\n                    conv_kernel=CONFIG['conv_kernel_size'], dropout=CONFIG['dropout']).to(device)\n\nctc_loss = nn.CTCLoss(blank=CONFIG['blank_index'], zero_infinity=True)\noptimizer = optim.AdamW(model.parameters(), lr=CONFIG['learning_rate'], weight_decay=1e-4)\nscheduler = optim.lr_scheduler.OneCycleLR(optimizer, max_lr=CONFIG['learning_rate'],\n                                          steps_per_epoch=len(train_loader),\n                                          epochs=CONFIG['num_epochs'])\n\nscaler = torch.cuda.amp.GradScaler() \nwith torch.cuda.amp.autocast(): \n\n\n    best_wer = float('inf')\n    best_model = None\n\nfor epoch in range(1, CONFIG['num_epochs']+1):\n    model.train()\n    total_loss = 0.0\n    for batch in tqdm(train_loader, desc=f\"Epoch {epoch}/{CONFIG['num_epochs']} - Training\"):\n        inputs = batch['neural'].to(device)         # (B, T, D)\n        input_lengths = batch['lengths'].to(device)\n        targets = batch['target'].to(device)        # (B, T_tgt)\n        target_lengths = batch['target_lengths'].to(device)\n        \n        optimizer.zero_grad()\n        with torch.cuda.amp.autocast():\n            log_probs, output_lengths = model(inputs, input_lengths)\n            # CTC loss expects (T, B, C) and lengths\n            loss = ctc_loss(log_probs, targets, output_lengths, target_lengths)\n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n        scheduler.step()\n        total_loss += loss.item()\n    avg_loss = total_loss / len(train_loader)\n    print(f\"Epoch {epoch} training loss: {avg_loss:.4f}\")\n    \n    # Validation with greedy decoding\n    model.eval()\n    preds = []\n    refs  = []\n    with torch.no_grad():\n        for batch in val_loader:\n            inputs = batch['neural'].to(device)\n            input_lengths = batch['lengths'].to(device)\n            log_probs, output_lengths = model(inputs, input_lengths)  # (T, B, C)\n            # Greedy decode\n            # Permute log_probs to (B, T, C) and take argmax\n            probs = log_probs.cpu().transpose(0,1)  # (B, T, C)\n            batch_preds = []\n            for seq in probs:\n                seq = seq.argmax(dim=-1).numpy().tolist()\n                # CTC collapse (remove repeats and blanks)\n                text = []\n                prev = CONFIG['blank_index']\n                for idx in seq:\n                    if idx != CONFIG['blank_index'] and idx != prev:\n                        text.append(idx2char[idx])\n                    prev = idx\n                batch_preds.append(\"\".join(text))\n            preds.extend(batch_preds)\n            refs.extend([s.lower() for s in batch['sentences']])\n    # Compute WER\n    val_wer = wer(refs, preds)\n    print(f\"Epoch {epoch} validation WER: {val_wer*100:.2f}%\")\n    torch.cuda.empty_cache()\n    # Save best model\n    if val_wer < best_wer:\n        best_wer = val_wer\n        best_model = model.state_dict()\n        print(f\"New best model (WER {val_wer*100:.2f}%).\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T08:09:24.277363Z","iopub.status.idle":"2025-12-22T08:09:24.277674Z","shell.execute_reply.started":"2025-12-22T08:09:24.277514Z","shell.execute_reply":"2025-12-22T08:09:24.277531Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}