{"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,"isSourceIdPinned":false,"sourceType":"competition"}],"dockerImageVersionId":31154,"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-10-31T18:27:11.124827Z","iopub.execute_input":"2025-10-31T18:27:11.125405Z","iopub.status.idle":"2025-10-31T18:27:11.192608Z","shell.execute_reply.started":"2025-10-31T18:27:11.125382Z","shell.execute_reply":"2025-10-31T18:27:11.191983Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %% [code]\n# ====================================================\n# Step 1: Install Libraries\n# ====================================================\nprint(\"Installing required libraries...\")\n# We need the pyctcdecode library and its kenlm dependency\n# to run the beam search with the language model.\n!pip install pyctcdecode kenlm -q\nprint(\"Installation complete.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-31T18:27:11.193792Z","iopub.execute_input":"2025-10-31T18:27:11.193959Z","iopub.status.idle":"2025-10-31T18:27:14.341877Z","shell.execute_reply.started":"2025-10-31T18:27:11.193946Z","shell.execute_reply":"2025-10-31T18:27:14.340705Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ====================================================\n# Step 2: Imports & Configuration\n# ====================================================\nimport os\nimport h5py\nimport glob # Import glob to find files\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\n\nclass CFG:\n    DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n    \n    # --- Data Paths ---\n    DATA_DIR = \"/kaggle/input/brain-to-text-25/t15_copyTask_neuralData\"\n    HDF5_DIR = os.path.join(DATA_DIR, \"hdf5_data_final\")\n    \n    TEST_FILES = sorted(glob.glob(os.path.join(HDF5_DIR, \"t15.*\", \"data_test.hdf5\")))\n    \n    # --- Model Path ---\n    MODEL_WEIGHTS_PATH = \"/kaggle/input/brain-to-text-25/t15_pretrained_rnn_baseline/t15_pretrained_rnn_baseline/checkpoint/best_checkpoint\"\n\n    # --- Model Parameters (for 'best_checkpoint') ---\n    N_FEATURES = 7168\n    N_PHONEMES = 41\n    HIDDEN_SIZE = 768\n    N_LAYERS = 5\n    DROPOUT = 0.2\n    \n    # --- Decoding Parameters ---\n    BATCH_SIZE = 32\n    NUM_WORKERS = 2\n    \n# --- Sanity Check ---\nprint(f\"Using device: {CFG.DEVICE}\")\nprint(f\"Found {len(CFG.TEST_FILES)} test files:\")\n\nif not CFG.TEST_FILES:\n    raise FileNotFoundError(\"No test HDF5 files found.\")\n\nprint(f\"Checking for model file at: {CFG.MODEL_WEIGHTS_PATH}\")\nif not os.path.exists(CFG.MODEL_WEIGHTS_PATH):\n    raise FileNotFoundError(f\"Pre-trained 'best_checkpoint' not found.\")\nprint(\"All necessary files found.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-31T18:45:27.493909Z","iopub.execute_input":"2025-10-31T18:45:27.494482Z","iopub.status.idle":"2025-10-31T18:45:27.607394Z","shell.execute_reply.started":"2025-10-31T18:45:27.494450Z","shell.execute_reply":"2025-10-31T18:45:27.606758Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %% [code]\n# ====================================================\n# Step 3: Define Phoneme Vocabulary\n# ====================================================\n\n# This is the list of the 40 phonemes the model predicts.\n# Index 0 (the 'BLANK' token) is handled automatically by the decoder.\n# This list comes from the competition's \"Data\" page description.\nVOCAB = [\n    'AA', 'AE', 'AH', 'AO', 'AW', 'AY', 'B', 'CH', 'D', 'DH', 'EH', 'ER', \n    'EY', 'F', 'G', 'HH', 'IH', 'IY', 'JH', 'K', 'L', 'M', 'N', 'NG', 'OW',\n    'OY', 'P', 'R', 'S', 'SH', 'T', 'TH', 'UH', 'UW', 'V', 'W', 'Y', 'Z',\n    'ZH', '|'  # '|' is the word boundary token\n]\n\nprint(f\"Vocabulary size: {len(VOCAB)} phonemes\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-31T18:45:29.797512Z","iopub.execute_input":"2025-10-31T18:45:29.797769Z","iopub.status.idle":"2025-10-31T18:45:29.803121Z","shell.execute_reply.started":"2025-10-31T18:45:29.797749Z","shell.execute_reply":"2025-10-31T18:45:29.802210Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ====================================================\n# Step 4: Define Model, Dataset, and DataLoader\n# ====================================================\n\n# --- 1. Model Definition (for 'best_checkpoint') ---\nclass NeuralDecoderRNN(nn.Module):\n    def __init__(self, n_features, n_phonemes, hidden_size, n_layers, dropout):\n        super().__init__()\n        self.gru = nn.GRU(\n            input_size=n_features,\n            hidden_size=hidden_size,\n            num_layers=n_layers,\n            bidirectional=False,\n            batch_first=True,\n            dropout=dropout\n        )\n        self.out = nn.Linear(hidden_size, n_phonemes) \n\n    def forward(self, x):\n        gru_out, _ = self.gru(x)\n        logits = self.out(gru_out)\n        return logits.permute(1, 0, 2)\n\n# --- 2. Test Dataset ---\nclass TestDataset(Dataset):\n    def __init__(self, hdf5_paths):\n        self.hdf5_paths = hdf5_paths\n        # This will store (file_idx, trial_key) to find data later\n        self.trial_keys = []    \n        # We must keep the HDF5 files open\n        self.open_files = {}    \n\n        print(\"Loading trial keys from all test files...\")\n        for i, h5_path in enumerate(tqdm(hdf5_paths, desc=\"Loading files\")):\n            try:\n                # Open the file and store the handle\n                f = h5py.File(h5_path, 'r')\n                self.open_files[i] = f\n                # Add all trial keys from this file\n                for key in f.keys():\n                    self.trial_keys.append((i, key))\n            except Exception as e:\n                print(f\"Error opening {h5_path}: {e}\")\n        \n        print(f\"Found {len(self.trial_keys)} total test samples.\")\n        \n        if len(self.trial_keys) == 0:\n            raise ValueError(\"No test samples found in any HDF5 file.\")\n        \n        # We know from before the key is 'input_features'\n        self.data_key = 'input_features' \n\n    def __len__(self):\n        return len(self.trial_keys)\n\n    def __getitem__(self, idx):\n        # Get the file index and the trial key\n        file_idx, key = self.trial_keys[idx]\n        \n        # Get the data from the correct open file\n        trial_group = self.open_files[file_idx][key]\n        neural_data = torch.tensor(trial_group[self.data_key][:], dtype=torch.float32) \n        \n        # Return the original key (e.g., 'trial_0000') as the ID\n        return neural_data, key\n\n# --- 3. Collate Function ---\ndef collate_fn_test(batch):\n    neural_data, trial_ids = zip(*batch)\n    neural_padded = nn.utils.rnn.pad_sequence(neural_data, batch_first=True, padding_value=0.0)\n    return neural_padded, trial_ids\n\n# --- 4. Create Test DataLoader ---\nprint(\"Initializing test dataset...\")\n# Pass the list of files from CFG\ntest_dataset = TestDataset(CFG.TEST_FILES) \n\ntest_loader = DataLoader(\n    test_dataset,\n    batch_size=CFG.BATCH_SIZE,\n    shuffle=False, # Must be False for submission\n    collate_fn=collate_fn_test,\n    num_workers=CFG.NUM_WORKERS\n)\nprint(\"Test data loaded.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-31T18:45:31.058182Z","iopub.execute_input":"2025-10-31T18:45:31.058782Z","iopub.status.idle":"2025-10-31T18:45:41.946396Z","shell.execute_reply.started":"2025-10-31T18:45:31.058760Z","shell.execute_reply":"2025-10-31T18:45:41.945587Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ====================================================\n# Step 5: Load Model\n# ====================================================\n\n# --- 1. Load Pre-trained Stage 1 Model ---\nprint(\"Loading pre-trained model...\")\n# This initializes the 7168-feature model\nmodel = NeuralDecoderRNN(\n    n_features=CFG.N_FEATURES,\n    n_phonemes=CFG.N_PHONEMES,\n    hidden_size=CFG.HIDDEN_SIZE,\n    n_layers=CFG.N_LAYERS,\n    dropout=0.0\n).to(CFG.DEVICE)\n\nprint(\"Loading checkpoint file...\")\ncheckpoint_dict = torch.load(CFG.MODEL_WEIGHTS_PATH, map_location=CFG.DEVICE, weights_only=False)\nmodel_state_dict = checkpoint_dict['model_state_dict']\n\nstate_dict = {}\nfor key, value in model_state_dict.items():\n    if key.startswith('_orig_mod.'):\n        new_key = key[len('_orig_mod.'):]\n        state_dict[new_key] = value\n    else:\n        state_dict[key] = value\n\n# Load with strict=False to ignore mismatched 'day_weights' keys\nmodel.load_state_dict(state_dict, strict=False)\n\nmodel.eval()\nprint(\"Model loaded successfully.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-31T18:45:46.185383Z","iopub.execute_input":"2025-10-31T18:45:46.185897Z","iopub.status.idle":"2025-10-31T18:45:46.956450Z","shell.execute_reply.started":"2025-10-31T18:45:46.185872Z","shell.execute_reply":"2025-10-31T18:45:46.955767Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ====================================================\n# Step 6: Run Inference Loop and Generate Submission\n# ====================================================\n\n# This will store ONLY the sentence strings, in the correct order\nall_predictions = [] \nBLANK_TOKEN_ID = 40 \n\nprint(f\"Starting inference on {len(test_loader)} batches...\")\n\nwith torch.no_grad():\n    for batch in tqdm(test_loader, desc=\"Decoding Test Set\"):\n        neural_padded, trial_ids = batch # trial_ids is a list of keys\n        \n        neural_padded = neural_padded.to(CFG.DEVICE)\n        \n        B, T, N_data = neural_padded.shape\n        if N_data != CFG.N_FEATURES:\n            padded_input = torch.zeros(B, T, CFG.N_FEATURES, device=CFG.DEVICE)\n            padded_input[:, :, :N_data] = neural_padded\n        else:\n            padded_input = neural_padded\n        \n        # 2. Get model predictions\n        logits = model(padded_input) \n        \n        # 3. Perform Greedy Decoding\n        pred_indices = torch.argmax(logits, dim=2)\n        pred_indices_batch = pred_indices.permute(1, 0).cpu().numpy()\n        \n        # 5. Decode each sequence in the batch manually\n        for i, indices in enumerate(pred_indices_batch):\n            collapsed_indices = []\n            for j, idx in enumerate(indices):\n                if j == 0 or idx != indices[j-1]:\n                    collapsed_indices.append(idx)\n            \n            decoded_phonemes = []\n            for idx in collapsed_indices:\n                if idx != BLANK_TOKEN_ID:\n                    decoded_phonemes.append(VOCAB[idx])\n            \n            sentence = \"\".join(decoded_phonemes)\n            sentence = sentence.replace('|', ' ').strip()\n            sentence_formatted = sentence.lower()\n            \n            all_predictions.append(sentence_formatted)\n\nprint(\"Inference complete.\")\nprint(f\"Total predictions generated: {len(all_predictions)}\")\n\n# --- Create and Save Submission File ---\nprint(\"Creating submission.csv...\")\n\n# Create the IDs from 0 to N-1\nsubmission_ids = list(range(len(all_predictions)))\n\n# Create the DataFrame\nsubmission_df = pd.DataFrame({\n    \"id\": submission_ids,\n    \"text\": all_predictions \n})\n# -----------------------------------------------------\n\nprint(f\"Final submission.csv has {len(submission_df)} rows.\")\nsubmission_df.to_csv(\"submission.csv\", index=False)\n\nprint(\"submission.csv created successfully!\")\ndisplay(submission_df.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-31T19:01:18.488898Z","iopub.execute_input":"2025-10-31T19:01:18.489445Z","iopub.status.idle":"2025-10-31T19:01:56.843255Z","shell.execute_reply.started":"2025-10-31T19:01:18.489413Z","shell.execute_reply":"2025-10-31T19:01:56.840889Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}