{"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":"none","dataSources":[{"sourceId":106809,"databundleVersionId":13056355,"sourceType":"competition"}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"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-11-09T13:14:10.863240Z","iopub.execute_input":"2025-11-09T13:14:10.863568Z","iopub.status.idle":"2025-11-09T13:14:11.528102Z","shell.execute_reply.started":"2025-11-09T13:14:10.863544Z","shell.execute_reply":"2025-11-09T13:14:11.527372Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Brain-to-Text '25: Baseline Submission Notebook\n\nThis notebook will set up the environment, load the data, and run the official baseline model to generate a valid `submission.csv`.\n\n**CRITICAL LIMITATION:**\nThe top-scoring baseline uses a 5-gram language model that requires **~300GB of RAM**, and the 3-gram model requires **~60GB**. A Kaggle notebook only has ~16GB (or ~30GB with High RAM).\n\nTherefore, this notebook will generate a submission using:\n1.  The **pre-trained RNN model** (provided by the hosts).\n2.  The **1-gram language model** (which fits in Kaggle's memory).\n\nThis will get you a valid score on the leaderboard. To improve, you will need to train a better neural model and/or use a high-RAM machine to run the larger language models.","metadata":{}},{"cell_type":"code","source":"import h5py\nimport os\n\n# Define paths from your CFG class\nDATA_DIR = \"/kaggle/input/brain-to-text-25/t15_copyTask_neuralData/hdf5_data_final\"\nSUBFOLDER = \"t15.2025.03.14\"\n\n# Path to the *test* file\ntest_file_path = os.path.join(DATA_DIR, SUBFOLDER, \"data_test.hdf5\")\n\nprint(f\"Inspecting test file: {test_file_path}\")\ntry:\n    with h5py.File(test_file_path, \"r\") as f:\n        # Get the first trial key (e.g., 'trial_0000')\n        first_trial_key = sorted(list(f.keys()))[0]\n        print(f\"Inspecting keys inside: {first_trial_key}\")\n        \n        trial_group = f[first_trial_key]\n        \n        # --- THIS IS THE IMPORTANT PART ---\n        print(f\"Keys found: {list(trial_group.keys())}\")\n        \n        for key in trial_group.keys():\n            print(f\"  • {key}: shape {trial_group[key].shape}, dtype {trial_group[key].dtype}\")\n\nexcept Exception as e:\n    print(f\"An error occurred: {e}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-09T13:14:11.529353Z","iopub.execute_input":"2025-11-09T13:14:11.529997Z","iopub.status.idle":"2025-11-09T13:14:11.909746Z","shell.execute_reply.started":"2025-11-09T13:14:11.529972Z","shell.execute_reply":"2025-11-09T13:14:11.908945Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install jiwer","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-09T18:59:21.898614Z","iopub.execute_input":"2025-11-09T18:59:21.899163Z","iopub.status.idle":"2025-11-09T18:59:29.209704Z","shell.execute_reply.started":"2025-11-09T18:59:21.899139Z","shell.execute_reply":"2025-11-09T18:59:29.208912Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# STEP 1: Imports and Configuration\n# ============================================================\nprint(\"Importing libraries...\")\nimport os\nimport yaml\nimport h5py\nimport torch\nimport numpy as np\nimport pandas as pd\nfrom torch import nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom tqdm.auto import tqdm\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\n# Configuration\nclass CFG:\n    # --- Model Hyperparameters ---\n    # Used for Transformer\n    N_HEAD = 8 \n    \n    # --- Training ---\n    EPOCHS = 5\n    LR = 1e-3\n    BATCH_SIZE = 32\n    DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    # --- Paths ---\n    DATA_DIR = \"/kaggle/input/brain-to-text-25/t15_copyTask_neuralData/hdf5_data_final\"\n    CHECKPOINT_PATH = \"/kaggle/input/brain-to-text-25/t15_pretrained_rnn_baseline/t15_pretrained_rnn_baseline/checkpoint/best_checkpoint\"\n    ARGS_PATH = \"/kaggle/input/brain-to-text-25/t15_pretrained_rnn_baseline/t15_pretrained_rnn_baseline/checkpoint/args.yaml\"\n    COMPETITION_TEST_PATH = \"/kaggle/input/brain-to-text-25/data_test.hdf5\"\n\nprint(f\"Running on device: {CFG.DEVICE}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-09T18:59:29.211457Z","iopub.execute_input":"2025-11-09T18:59:29.211664Z","iopub.status.idle":"2025-11-09T18:59:36.448941Z","shell.execute_reply.started":"2025-11-09T18:59:29.211643Z","shell.execute_reply":"2025-11-09T18:59:36.448247Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# STEP 2: Load Model Configuration (args.yaml)\n# ============================================================\nwith open(CFG.ARGS_PATH, \"r\") as f:\n    args = yaml.safe_load(f)\n\nprint(\"Loaded model config from args.yaml:\")\n# Use .get() for safety, providing defaults from your model init\nINPUT_SIZE = args.get(\"input_size\", 256)\nHIDDEN_SIZE = args.get(\"hidden_size\", 512)\nOUTPUT_SIZE = args.get(\"output_size\", 29)\nNUM_LAYERS = args.get(\"num_layers\", 1)\n\nprint(args)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-09T18:59:36.449700Z","iopub.execute_input":"2025-11-09T18:59:36.450141Z","iopub.status.idle":"2025-11-09T18:59:36.479758Z","shell.execute_reply.started":"2025-11-09T18:59:36.450118Z","shell.execute_reply":"2025-11-09T18:59:36.479166Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# BƯỚC 2.5: KIỂM TRA KEY CỦA TỆP HUẤN LUYỆN\n# ============================================================\nprint(\"Inspecting a TRAINING file to find the target key...\")\n\n# Thay đổi đường dẫn này nếu cần,\n# đây là một trong các thư mục huấn luyện\nTRAIN_SUBFOLDER = \"t15.2025.01.10\" \ntrain_file_path = os.path.join(CFG.DATA_DIR, TRAIN_SUBFOLDER, \"data_train.hdf5\")\n\ntry:\n    with h5py.File(train_file_path, \"r\") as f:\n        first_trial_key = sorted(list(f.keys()))[0]\n        print(f\"Inspecting keys inside trial: {first_trial_key}\")\n        \n        trial_group = f[first_trial_key]\n        print(f\"!!! TẤT CẢ CÁC KEY TÌM THẤY: {list(trial_group.keys())} !!!\")\n        \n        # In ra thông tin về các key có thể là target\n        for key in trial_group.keys():\n            if key != 'input_features':\n                print(f\"  • Tìm thấy key khả nghi: '{key}'\")\n                print(f\"    Shape: {trial_group[key].shape}, Dtype: {trial_group[key].dtype}\")\n                print(f\"    Sample data: {trial_group[key][:10]}\")\n                \nexcept Exception as e:\n    print(f\"Lỗi khi kiểm tra: {e}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-09T19:00:48.064354Z","iopub.execute_input":"2025-11-09T19:00:48.065091Z","iopub.status.idle":"2025-11-09T19:00:48.180687Z","shell.execute_reply.started":"2025-11-09T19:00:48.065059Z","shell.execute_reply":"2025-11-09T19:00:48.180025Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# STEP 3: Define Dataset Loader (Seq2Seq)\n# ============================================================\nfrom torch.utils.data import ConcatDataset, Dataset\n\ndef temporal_mask(data, mask_percentage=0.05, mask_value=0.0):\n    \"\"\"\n    Applies temporal masking to a 2D tensor [Sequence, Features].\n    \"\"\"\n    if not torch.is_tensor(data):\n        data = torch.tensor(data, dtype=torch.float32)\n        \n    seq_len, _ = data.shape\n    num_to_mask = int(seq_len * mask_percentage)\n    \n    if num_to_mask > 0:\n        mask_indices = torch.randperm(seq_len)[:num_to_mask]\n        data[mask_indices, :] = mask_value\n        \n    return data\n\nclass BrainDataset(Dataset):\n    \"\"\"\n    Reads data from a single HDF5 file (e.g., data_train.hdf5).\n    - input_key: The name of the HDF5 dataset for input features.\n    - target_key: The name of the HDF5 dataset for target sequences (indices).\n    - is_test: If True, __getitem__ also returns the trial_key.\n    \"\"\"\n    def __init__(self, hdf5_file, input_key=\"input_features\", target_key=\"phoneme_indices\", is_test=False, use_augmentation=False):\n        self.file_path = hdf5_file\n        self.input_key = input_key\n        self.target_key = target_key # Key for sequence targets\n        self.is_test = is_test\n        \n        # --- FIX 2: STORE THE PARAMETER ---\n        self.use_augmentation = use_augmentation \n        self.file = None # File handle\n        \n        try:\n            with h5py.File(self.file_path, \"r\") as f:\n                self.trial_keys = sorted(list(f.keys()))\n        except FileNotFoundError:\n            # Handle cases where a subfolder might be missing a split\n            print(f\"Warning: File not found {self.file_path}, creating empty dataset.\")\n            self.trial_keys = []\n\n    def __len__(self):\n        return len(self.trial_keys)\n\n    def __getitem__(self, idx):\n        if self.file is None:\n            self.file = h5py.File(self.file_path, \"r\")\n            \n        trial_key = self.trial_keys[idx]\n        trial_group = self.file[trial_key]\n        \n        x_data = trial_group[self.input_key][:]\n        x = torch.tensor(x_data, dtype=torch.float32)\n        \n        if self.use_augmentation and not self.is_test:\n            x = temporal_mask(x, mask_percentage=0.1)\n        \n        if self.target_key in trial_group:\n            # Assume targets are a 1D array of integer indices (for CTC)\n            y_data = trial_group[self.target_key][:]\n            y = torch.tensor(y_data, dtype=torch.long)\n        else:\n            # Create an empty long tensor as a placeholder for test/dummy targets\n            y = torch.tensor([], dtype=torch.long)\n        \n        if self.is_test:\n            return x, y, trial_key\n        else:\n            return x, y\n\ndef load_datasets():\n    \"\"\"\n    Scans all subfolders in CFG.DATA_DIR and creates combined\n    train, val, and test datasets from all found files.\n    \"\"\"\n    train_datasets = []\n    val_datasets = []\n    test_datasets = []\n\n    subfolders = [f.path for f in os.scandir(CFG.DATA_DIR) if f.is_dir()]\n    print(f\"Found {len(subfolders)} session folders.\")\n    \n    for subfolder_path in subfolders:\n        session_name = os.path.basename(subfolder_path)\n        \n        train_file = os.path.join(subfolder_path, \"data_train.hdf5\")\n        val_file = os.path.join(subfolder_path, \"data_val.hdf5\")\n        test_file = os.path.join(subfolder_path, \"data_test.hdf5\")\n\n        # --- PASS THE use_augmentation FLAG ---\n        # Only apply augmentations to the training set\n        train_set = BrainDataset(train_file, input_key=\"input_features\", target_key=\"seq_class_ids\", is_test=False, use_augmentation=True)\n        val_set = BrainDataset(val_file, input_key=\"input_features\", target_key=\"seq_class_ids\", is_test=False, use_augmentation=False)\n        test_set = BrainDataset(test_file, input_key=\"input_features\", target_key=\"seq_class_ids\", is_test=True, use_augmentation=False) \n        \n        if len(train_set) > 0:\n            train_datasets.append(train_set)\n        if len(val_set) > 0:\n            val_datasets.append(val_set)\n        if len(test_set) > 0:\n            test_datasets.append(test_set)\n            \n    # Combine all individual datasets into one large dataset\n    full_train_dataset = ConcatDataset(train_datasets)\n    full_val_dataset = ConcatDataset(val_datasets)\n    full_test_dataset = ConcatDataset(test_datasets)\n    \n    return full_train_dataset, full_val_dataset, full_test_dataset\n\nprint(\"Loading Train/Val/Test data from session folders...\")\ntrain_dataset, val_dataset, test_dataset = load_datasets()\nprint(\"=\"*40)\nprint(f\"Total Train samples: {len(train_dataset)}\")\nprint(f\"Total Val samples: {len(val_dataset)}\")\nprint(f\"Total Test samples: {len(test_dataset)}\")\nprint(\"=\"*40)\n\n# Check samples\nsample_x_train, sample_y_train = train_dataset[0]\nprint(f\"Train sample X shape: {sample_x_train.shape}\")\nprint(f\"Train sample Y (indices): {sample_y_train}\")\nprint(f\"Train sample Y shape: {sample_y_train.shape}, dtype: {sample_y_train.dtype}\")\n\nif len(test_dataset) > 0:\n    sample_x_test, sample_y_test, _ = test_dataset[0]\n    print(f\"Test sample X shape: {sample_x_test.shape}\")\n    print(f\"Test sample Y (dummy): {sample_y_test}\")\n    print(f\"Test sample Y shape: {sample_y_test.shape}, dtype: {sample_y_test.dtype}\")\nelse:\n    print(f\"Warning: Competition test file not found at {CFG.COMPETITION_TEST_PATH}\")\n    print(\"This is normal. The file will be present during submission.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-09T19:01:37.124851Z","iopub.execute_input":"2025-11-09T19:01:37.125484Z","iopub.status.idle":"2025-11-09T19:01:45.257642Z","shell.execute_reply.started":"2025-11-09T19:01:37.125459Z","shell.execute_reply":"2025-11-09T19:01:45.256941Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# STEP 3.5: INSPECT A TRAINING FILE\n# ============================================================\nprint(\"🕵️ Inspecting a training file to find the correct label key...\")\nimport h5py\nimport os\n\n# --- Find the FIRST available data_train.hdf5 file ---\ntrain_file_to_inspect = None\nsubfolders = [f.path for f in os.scandir(CFG.DATA_DIR) if f.is_dir()]\n\nfor subfolder_path in sorted(subfolders): # sort to get a consistent one\n    train_file = os.path.join(subfolder_path, \"data_train.hdf5\")\n    if os.path.exists(train_file):\n        train_file_to_inspect = train_file\n        break\n        \nif train_file_to_inspect:\n    print(f\"Inspecting file: {train_file_to_inspect}\")\n    try:\n        with h5py.File(train_file_to_inspect, \"r\") as f:\n            first_trial_key = sorted(list(f.keys()))[0]\n            print(f\"Inspecting keys inside trial: {first_trial_key}\")\n            \n            trial_group = f[first_trial_key]\n            print(f\"\\n--- 💡 ALL KEYS FOUND IN THIS TRIAL 💡 ---\")\n            for key in trial_group.keys():\n                print(f\"  • {key}\")\n            print(f\"-------------------------------------------\\n\")\n            print(\"Find the key that looks like labels (e.g., 'targets', 'labels', 'phonemes') and use it in the next step.\")\n            \n    except Exception as e:\n        print(f\"An error occurred: {e}\")\nelse:\n    print(\"Error: Could not find any data_train.hdf5 files to inspect.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-09T18:59:45.266032Z","iopub.execute_input":"2025-11-09T18:59:45.266382Z","iopub.status.idle":"2025-11-09T18:59:45.512264Z","shell.execute_reply.started":"2025-11-09T18:59:45.266346Z","shell.execute_reply":"2025-11-09T18:59:45.511547Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# STEP 4: Define Model (Adapter for Fine-tuning)\n# ============================================================\n\n# --- 1. DEFINE VOCABULARY ---\n# This is the list of the 40 phonemes.\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# We add +1 for the blank token\nOUTPUT_SIZE = len(VOCAB) + 1 # 40 + 1 = 41\nBLANK_ID = 0\n\n# --- 2. DEFINE MODEL PARAMETERS ---\n# These MUST match the shapes in your checkpoint\nDATA_INPUT_SIZE = 512       # From data\nADAPTER_OUTPUT_SIZE = 256   # <-- THIS IS THE MISSING VARIABLE\nHIDDEN_SIZE = 512           # The hidden size of your pretrained RNN\nNUM_LAYERS = 1              # The num_layers of your pretrained RNN\nIS_BIDIRECTIONAL = False\n\nprint(f\"Vocabulary Config: {len(VOCAB)} phonemes + 1 blank = {OUTPUT_SIZE} classes.\")\nprint(f\"Model Config: Data(512) -> Adapter(256) -> RNN(256, {HIDDEN_SIZE}) -> FC({HIDDEN_SIZE}, {OUTPUT_SIZE})\")\n\n\n### --- Model 1: Flexible Recurrent Model (Seq2Seq) --- ###\nclass RecurrentModel(nn.Module):\n    # --- MODIFIED: Added adapter_output_size ---\n    def __init__(self, model_type, data_input_size, adapter_output_size, \n                 hidden_size, output_size, num_layers, bidirectional):\n        super().__init__()\n        self.hidden_size = hidden_size\n        self.bidirectional = bidirectional\n        \n        # --- ADDED: Adapter Layer ---\n        self.adapter_layer = nn.Linear(data_input_size, adapter_output_size)\n        \n        rnn_args = {\n            'input_size': adapter_output_size, # <-- Use adapter output\n            'hidden_size': hidden_size,\n            'num_layers': num_layers,\n            'batch_first': True,\n            'bidirectional': bidirectional\n        }\n        \n        if model_type == \"LSTM\": self.rnn = nn.LSTM(**rnn_args)\n        elif model_type == \"GRU\": self.rnn = nn.GRU(**rnn_args)\n        elif model_type == \"RNN\": self.rnn = nn.RNN(**rnn_args)\n        else: raise ValueError(\"Invalid model_type\")\n\n        fc_in_features = hidden_size * 2 if bidirectional else hidden_size\n        self.fc = nn.Linear(fc_in_features, output_size) # output_size is 41\n\n    def forward(self, x):\n        # x is [B, S, 512]\n        x = self.adapter_layer(x) # [B, S, 256]\n        out, _ = self.rnn(x)      # [B, S, 512]\n        out = self.fc(out)        # [B, S, 41]\n        return nn.functional.log_softmax(out, dim=2)\n\n        \n### --- Model 2: Transformer Encoder Model (Seq2Seq) --- ###\nclass TransformerEncModel(nn.Module):\n    # --- MODIFIED: Added adapter_output_size ---\n    def __init__(self, data_input_size, adapter_output_size, n_head, num_layers, \n                 dim_feedforward, output_size):\n        super().__init__()\n        \n        # --- ADDED: Adapter Layer ---\n        self.adapter_layer = nn.Linear(data_input_size, adapter_output_size)\n        self.d_model = adapter_output_size # d_model is now 256\n        \n        encoder_layer = nn.TransformerEncoderLayer(\n            d_model=self.d_model, nhead=n_head, \n            dim_feedforward=dim_feedforward,\n            batch_first=True, dropout=0.1\n        )\n        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers)\n        self.fc = nn.Linear(self.d_model, output_size) # output_size is 41\n\n    def forward(self, x):\n        # x is [B, S, 512]\n        x = self.adapter_layer(x) # [B, S, 256]\n        out = self.transformer_encoder(x) # [B, S, 256]\n        out = self.fc(out)                # [B, S, 41]\n        return nn.functional.log_softmax(out, dim=2)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-09T13:14:31.386496Z","iopub.execute_input":"2025-11-09T13:14:31.386691Z","iopub.status.idle":"2025-11-09T13:14:31.397465Z","shell.execute_reply.started":"2025-11-09T13:14:31.386676Z","shell.execute_reply":"2025-11-09T13:14:31.396718Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# STEP 6: Prepare Dataloaders\n# ============================================================\nimport torch.nn.utils.rnn as rnn_utils\n\ndef custom_collate(batch):\n    \"\"\"\n    Custom collate function for CTC (Sequence-to-Sequence).\n    Pads both x (inputs) and y (targets) and returns their original lengths.\n    `batch` is a list of tuples: (x, y) or (x, y, key)\n    \"\"\"\n    # Check if it's a test batch (has 3 items: x, y, key)\n    is_test = len(batch[0]) == 3\n\n    if is_test:\n        xs, ys, keys = zip(*batch)\n    else:\n        xs, ys = zip(*batch) # Train/val batch\n        \n    # These are the unpadded lengths, required by CTCLoss\n    x_lengths = torch.tensor([len(x) for x in xs], dtype=torch.long)\n    y_lengths = torch.tensor([len(y) for y in ys], dtype=torch.long)\n    \n    # 1. Pad the 'x' sequences (inputs)\n    padded_xs = rnn_utils.pad_sequence(xs, batch_first=True, padding_value=0.0)\n    \n    # 2. Pad the 'y' sequences (targets)\n    # We use padding_value=0. This assumes '0' is your 'blank' token index.\n    # This will also correctly handle the empty 'y' tensors from the test set.\n    padded_ys = rnn_utils.pad_sequence(ys, batch_first=True, padding_value=0)\n    \n    if is_test:\n        return padded_xs, padded_ys, x_lengths, y_lengths, keys\n    else:\n        return padded_xs, padded_ys, x_lengths, y_lengths\n\n# Create the DataLoaders using this new collate function\ntrain_loader = DataLoader(\n    train_dataset, \n    batch_size=CFG.BATCH_SIZE, \n    shuffle=True, \n    collate_fn=custom_collate\n)\nval_loader = DataLoader(\n    val_dataset, \n    batch_size=CFG.BATCH_SIZE, \n    shuffle=False, \n    collate_fn=custom_collate \n)\ntest_loader = DataLoader(\n    test_dataset, \n    batch_size=CFG.BATCH_SIZE, \n    shuffle=False, \n    collate_fn=custom_collate\n)\n\nprint(\"DataLoaders created with CTC-ready padding.\")\n\n# --- Let's check a sample batch ---\ntry:\n    x_batch, y_batch, x_len, y_len = next(iter(train_loader))\n    print(\"\\nChecking one batch from train_loader:\")\n    print(f\"  x_batch shape: {x_batch.shape}\")\n    print(f\"  y_batch shape: {y_batch.shape}\")\n    print(f\"  x_lengths shape: {x_len.shape}, sample: {x_len[:5]}\")\n    print(f\"  y_lengths shape: {y_len.shape}, sample: {y_len[:5]}\")\nexcept Exception as e:\n    print(f\"\\nCould not get batch from train_loader (is it empty?): {e}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-09T13:14:31.398360Z","iopub.execute_input":"2025-11-09T13:14:31.398979Z","iopub.status.idle":"2025-11-09T13:14:32.910768Z","shell.execute_reply.started":"2025-11-09T13:14:31.398951Z","shell.execute_reply":"2025-11-09T13:14:32.909963Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# STEP 7: Experiment Setup (for CTC / Seq2Seq)\n# ============================================================\nimport jiwer # For Word/Character/Phoneme Error Rate\nimport shutil\n# ---\n# 1. Automatically build the TOKEN_MAP from the VOCAB\n# ---\nTOKEN_MAP = {i + 1: phoneme for i, phoneme in enumerate(VOCAB)}\nTOKEN_MAP[BLANK_ID] = \"\" # BLANK_ID was set to 0 in STEP 4\n\nprint(\"Token map created:\")\nprint(f\"  Index 0: '{TOKEN_MAP[0]}' (BLANK)\")\nprint(f\"  Index 1: '{TOKEN_MAP[1]}' (e.g., AA)\")\nprint(f\"  Index 40: '{TOKEN_MAP[40]}' (e.g., |)\")\n\n# 2. Define the models AND configurations you want to test\n# We'll run each model twice: once with the pretrained checkpoint, once from scratch.\n# Format: (model_name, use_checkpoint_boolean)\nexperiments_to_run = [\n    (\"RNN\", True),\n    (\"RNN\", False),\n    (\"LSTM\", True),\n    (\"LSTM\", False),\n    (\"GRU\", True),\n    (\"GRU\", False),\n    (\"TRANSFORMER\", True),\n    (\"TRANSFORMER\", False),\n]\n\n# 3. A dictionary to store results from all experiments\nall_experiment_results = {}\n\n# 4. Define the decoders\ndef greedy_decoder(logits, token_map):\n    pred_indices = torch.argmax(logits, dim=-1)\n    collapsed_indices = torch.unique_consecutive(pred_indices)\n    final_indices = [idx.item() for idx in collapsed_indices if idx.item() != BLANK_ID]\n    \n    # --- MODIFICATION ---\n    # Join with a space to treat each phoneme as a \"word\"\n    phonemes = [token_map.get(i, \"?\") for i in final_indices]\n    text = \" \".join(phonemes)\n    # --- END MODIFICATION ---\n    \n    return text\n\ndef decode_true_target(target_indices, token_map):\n    phonemes = [token_map.get(i.item(), \"?\") for i in target_indices]\n    text = \" \".join(phonemes)\n    return text\n\n# 5. Define the new training and validation functions\ndef train_one_epoch(epoch, model, train_loader, criterion, optimizer):\n    model.train()\n    running_loss = 0.0\n    \n    for x, y, x_lengths, y_lengths in tqdm(train_loader, desc=f\"Epoch {epoch} [Train]\", leave=False):\n        x, y, x_lengths, y_lengths = x.to(CFG.DEVICE), y.to(CFG.DEVICE), x_lengths.to(CFG.DEVICE), y_lengths.to(CFG.DEVICE)\n        \n        optimizer.zero_grad()\n        y_pred = model(x) # [B, S, 41]\n        \n        y_pred_for_loss = y_pred.permute(1, 0, 2) # [S, B, 41]\n        \n        loss = criterion(y_pred_for_loss, y, x_lengths, y_lengths)\n        \n        if torch.isinf(loss) or torch.isnan(loss):\n            print(\"Warning: Skipping batch with inf/nan loss\")\n            continue\n            \n        loss.backward()\n        optimizer.step()\n        running_loss += loss.item() * x.size(0)\n        \n    return running_loss / len(train_loader.dataset)\n\ndef validate_one_epoch(epoch, model, val_loader, criterion, token_map):\n    model.eval()\n    val_loss = 0.0\n    all_pred_texts = []\n    all_true_texts = []\n    \n    with torch.no_grad():\n        for x, y, x_lengths, y_lengths in tqdm(val_loader, desc=f\"Epoch {epoch} [Val]\", leave=False):\n            x, y, x_lengths, y_lengths = x.to(CFG.DEVICE), y.to(CFG.DEVICE), x_lengths.to(CFG.DEVICE), y_lengths.to(CFG.DEVICE)\n\n            y_pred = model(x) # [B, S, 41]\n            y_pred_for_loss = y_pred.permute(1, 0, 2) # [S, B, 41]\n            loss = criterion(y_pred_for_loss, y, x_lengths, y_lengths)\n            val_loss += loss.item() * x.size(0)\n\n            for i in range(x.size(0)):\n                pred_logits = y_pred[i, :x_lengths[i], :]\n                true_indices = y[i, :y_lengths[i]]\n                \n                pred_text = greedy_decoder(pred_logits, token_map)\n                true_text = decode_true_target(true_indices, token_map)\n                \n                all_pred_texts.append(pred_text)\n                all_true_texts.append(true_text)\n\n    # --- MODIFICATION ---\n    # Use .wer() (Word Error Rate) to calculate Phoneme Error Rate\n    error_rate = jiwer.wer(all_true_texts, all_pred_texts)\n    # --- END MODIFICATION ---\n    \n    return val_loss / len(val_loader.dataset), error_rate\n\nprint(\"Experiment setup for CTC complete.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-09T13:14:32.911688Z","iopub.execute_input":"2025-11-09T13:14:32.912092Z","iopub.status.idle":"2025-11-09T13:14:32.945993Z","shell.execute_reply.started":"2025-11-09T13:14:32.912066Z","shell.execute_reply":"2025-11-09T13:14:32.945231Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# STEP 8: Run Experiment Loop (with Fine-tuning)\n# ============================================================\n\n# --- 8.1: Load the checkpoint data ---\nprint(\"Loading checkpoint weights from file...\")\ncheckpoint = torch.load(CFG.CHECKPOINT_PATH, map_location=CFG.DEVICE, weights_only=False)\nprint(\"Checkpoint data loaded.\")\n\n\n# --- 8.2: Start the main experiment loop ---\n# This loop iterates through the (model_name, use_checkpoint) tuples defined in STEP 7\nfor model_name, use_checkpoint in experiments_to_run:\n    experiment_name = f\"{model_name}_{'pretrained' if use_checkpoint else 'scratch'}\"\n    \n    print(f\"\\n{'='*20} 🚀 STARTING EXPERIMENT: {experiment_name} {'='*20}\\n\")\n    \n    # --- 8.3: Initialize Model ---\n    model = None\n    if model_name == \"TRANSFORMER\":\n        print(\"Initializing TransformerEncModel...\")\n        model = TransformerEncModel(\n            data_input_size=DATA_INPUT_SIZE,\n            adapter_output_size=ADAPTER_OUTPUT_SIZE,\n            n_head=CFG.N_HEAD,\n            num_layers=NUM_LAYERS,\n            dim_feedforward=HIDDEN_SIZE,\n            output_size=OUTPUT_SIZE\n        )\n    elif model_name in [\"RNN\", \"LSTM\", \"GRU\"]:\n        print(f\"Initializing RecurrentModel (Type: {model_name})...\")\n        model = RecurrentModel(\n            model_type=model_name,\n            data_input_size=DATA_INPUT_SIZE,\n            adapter_output_size=ADAPTER_OUTPUT_SIZE, # From STEP 4\n            hidden_size=HIDDEN_SIZE,\n            output_size=OUTPUT_SIZE,\n            num_layers=NUM_LAYERS,\n            bidirectional=IS_BIDIRECTIONAL\n        )\n    \n    model = model.to(CFG.DEVICE)\n    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n    print(f\"Total Trainable Parameters: {total_params:,}\")\n\n    # --- 8.4: Load Pretrained Weights ---\n    if use_checkpoint:\n        print(\"Applying pretrained weights (strict=False)...\")\n        missing_keys, unexpected_keys = model.load_state_dict(\n            checkpoint['model_state_dict'], \n            strict=False\n        )\n        print(f\"  > Weights loaded. Missing keys (good): {missing_keys}\")\n        print(f\"  > Unexpected keys (should be empty): {unexpected_keys}\")\n    else:\n        print(\"Training from scratch. No checkpoint loaded.\")\n\n\n    # --- 8.5: Define Loss and Optimizer ---\n    criterion = nn.CTCLoss(blank=BLANK_ID, zero_infinity=True) \n    optimizer = torch.optim.Adam(model.parameters(), lr=CFG.LR)\n    \n    # --- 8.6: Training Loop ---\n    history = {'train_loss': [], 'val_loss': [], 'error_rate': []}\n    best_error_rate = float(\"inf\")\n    \n    # Save file using the new experiment_name\n    best_model_path = f\"/kaggle/working/best_model_{experiment_name}.pth\"\n    \n    print(\"Starting training...\")\n\n    for epoch in range(1, CFG.EPOCHS + 1):\n        train_loss = train_one_epoch(epoch, model, train_loader, criterion, optimizer)\n        val_loss, error_rate = validate_one_epoch(epoch, model, val_loader, criterion, TOKEN_MAP)\n        \n        history['train_loss'].append(train_loss)\n        history['val_loss'].append(val_loss)\n        history['error_rate'].append(error_rate)\n        \n        print(f\"Epoch {epoch}/{CFG.EPOCHS} | Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f} | Error Rate (PER): {error_rate:.4f}\")\n        \n        if error_rate < best_error_rate:\n            best_error_rate = error_rate\n            torch.save(model.state_dict(), best_model_path)\n            print(f\"✅ Saved new best model for {experiment_name} (PER: {error_rate:.4f})!\")\n\n    # --- 8.7: Test / Prediction Loop ---\n    print(f\"Running predictions for {experiment_name}...\")\n    \n    # Load the correct best model file\n    model.load_state_dict(torch.load(best_model_path))\n    model.eval()\n    \n    all_pred_texts = []\n    all_trial_keys = []\n    \n    if len(test_loader.dataset) > 0:\n        with torch.no_grad():\n            for x, y, x_lengths, y_lengths, keys in tqdm(test_loader, desc=f\"Testing {experiment_name}\", leave=False):\n                x, x_lengths = x.to(CFG.DEVICE), x_lengths.to(CFG.DEVICE)\n                y_pred = model(x)\n                \n                for i in range(x.size(0)):\n                    pred_logits = y_pred[i, :x_lengths[i], :] \n                    pred_text = greedy_decoder(pred_logits, TOKEN_MAP)\n                    all_pred_texts.append(pred_text)\n                    all_trial_keys.append(keys[i])\n\n    # --- 8.8: Generate Submission ---\n    print(f\"Generating submission file for {experiment_name}...\")\n    submission_df = pd.DataFrame({'id': all_trial_keys, 'text': all_pred_texts})\n    \n    # Format text for submission (e.g., \"AA B | K\")\n    submission_df['text'] = submission_df['text'].str.strip()\n    \n    # Save submission using the new experiment_name\n    submission_path = f\"/kaggle/working/submission_{experiment_name}.csv\"\n    submission_df.to_csv(submission_path, index=False)\n    print(f\"✅ Submission file saved to {submission_path}\")\n\n    # --- 8.9: Store results ---\n    \n    # Save results using the new experiment_name\n    all_experiment_results[experiment_name] = {\n        'history': history,\n        'best_val_loss': min(history['val_loss']),\n        'best_error_rate': best_error_rate,\n        'total_params': total_params\n    }\n\nprint(\"\\n🎉 All experiments complete! 🎉\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-09T13:14:32.947952Z","iopub.execute_input":"2025-11-09T13:14:32.948164Z","iopub.status.idle":"2025-11-09T14:09:12.004280Z","shell.execute_reply.started":"2025-11-09T13:14:32.948149Z","shell.execute_reply":"2025-11-09T14:09:12.003697Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# STEP 9: Compare Experiment Results (for CTC)\n# ============================================================\n\nprint(\"📊 Experiment Results Summary\\n\")\nsns.set(style=\"whitegrid\", font_scale=1.1)\n\n# --- 1. Print Summary Table ---\nresults_df = pd.DataFrame.from_dict(all_experiment_results, orient='index')\nresults_df = results_df.drop(columns='history') # Drop history for clean table\nprint(results_df.to_markdown(floatfmt=\".6f\"))\n\n\n# --- Define a color map for model types ---\nmodel_colors = {\n    \"RNN\": \"blue\",\n    \"LSTM\": \"green\",\n    \"GRU\": \"red\",\n    \"TRANSFORMER\": \"purple\"\n}\n\n\n# --- 2. Plot Error Rate Curves ---\nplt.figure(figsize=(16, 9)) \nfor model_name, results in all_experiment_results.items():\n    \n    base_model = model_name.split('_')[0] # e.g., \"RNN\"\n    is_pretrained = \"pretrained\" in model_name\n    \n    color = model_colors.get(base_model, 'black') # Get color from map\n    linestyle = '-' if is_pretrained else '--'    # Solid for pretrained, dashed for scratch\n\n    error_history = results['history']['error_rate']\n    \n    # --- Added color and linestyle to plot call ---\n    plt.plot(\n        error_history, \n        label=f\"{model_name} (Best: {results['best_error_rate']:.4f})\", \n        lw=2, \n        color=color, \n        linestyle=linestyle\n    )\n\nplt.title('Model Comparison: Validation Error Rate (PER)')\nplt.xlabel('Epoch')\nplt.ylabel('Phoneme Error Rate (PER) (lower is better)')\nplt.legend(bbox_to_anchor=(1.05, 1), loc='upper left') # Move legend outside\nplt.grid(True, which='both', linestyle='--', linewidth=0.5)\nplt.tight_layout()\nplt.show()\n\n\n# --- 3. Plot Validation Loss Curves ---\nplt.figure(figsize=(16, 9))\nfor model_name, results in all_experiment_results.items():\n    base_model = model_name.split('_')[0] \n    is_pretrained = \"pretrained\" in model_name\n    color = model_colors.get(base_model, 'black')\n    linestyle = '-' if is_pretrained else '--'\n    \n    val_loss_history = results['history']['val_loss']\n\n    plt.plot(\n        val_loss_history, \n        label=f\"{model_name} (Best: {results['best_val_loss']:.4f})\", \n        lw=2,\n        color=color,\n        linestyle=linestyle\n    )\n\nplt.title('Model Comparison: Validation Loss (CTC)')\nplt.xlabel('Epoch')\nplt.ylabel('Validation CTCLoss')\nplt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')\nplt.grid(True, which='both', linestyle='--', linewidth=0.5)\nplt.tight_layout()\nplt.show()\n\n# --- 4. Plot Training Loss Curves ---\nplt.figure(figsize=(16, 9))\nfor model_name, results in all_experiment_results.items():\n    base_model = model_name.split('_')[0] \n    is_pretrained = \"pretrained\" in model_name\n    color = model_colors.get(base_model, 'black')\n    linestyle = '-' if is_pretrained else '--'\n\n    train_loss_history = results['history']['train_loss']\n    \n    plt.plot(\n        train_loss_history, \n        label=f\"{model_name}\", \n        lw=2,\n        color=color,\n        linestyle=linestyle\n    )\n\nplt.title('Model Comparison: Training Loss (CTC)')\nplt.xlabel('Epoch')\nplt.ylabel('Training CTCLoss')\nplt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')\nplt.grid(True, which='both', linestyle='--', linewidth=0.5)\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-09T14:09:12.005156Z","iopub.execute_input":"2025-11-09T14:09:12.005455Z","iopub.status.idle":"2025-11-09T14:09:13.064507Z","shell.execute_reply.started":"2025-11-09T14:09:12.005437Z","shell.execute_reply":"2025-11-09T14:09:13.063717Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# STEP 10: Generate Final Submission\n# ============================================================\nimport shutil # Make sure shutil is imported\n\nprint(\"Generating final submission file...\\n\")\n\nif not all_experiment_results:\n    print(\"Warning: `all_experiment_results` is empty. Please run STEP 8 first.\")\nelse:\n    # --- MODIFIED: Section 1 ---\n    # 1. Find the best model based on LOWEST PER, using LOWEST VAL_LOSS as a tie-breaker\n    best_model_name = None\n    best_model_error = float(\"inf\")\n    best_model_val_loss = float(\"inf\") # <-- New variable to track val_loss\n\n    for model_name, results in all_experiment_results.items():\n        current_error = results['best_error_rate']\n        current_val_loss = results['best_val_loss']\n        \n        # Check if this model is better:\n        # 1. Is the error rate strictly lower?\n        # 2. OR is the error rate the same, but the val_loss is lower?\n        if current_error < best_model_error or \\\n           (current_error == best_model_error and current_val_loss < best_model_val_loss):\n            \n            # This is our new best model\n            best_model_error = current_error\n            best_model_val_loss = current_val_loss\n            best_model_name = model_name\n\n    # --- MODIFIED: Print statement now shows both scores ---\n    print(f\"🏆 The best overall model is: {best_model_name}\")\n    print(f\"   > Best Phoneme Error Rate (PER): {best_model_error:.6f}\")\n    print(f\"   > Best Validation Loss: {best_model_val_loss:.6f}\")\n    # --- End of modifications ---\n\n    # 2. Define the paths\n    # This is the file saved by STEP 8 for the best model\n    best_submission_path = f\"/kaggle/working/submission_{best_model_name}.csv\"\n    \n    # This is the standard file Kaggle looks for\n    final_submission_path = \"/kaggle/working/submission.csv\"\n\n    # 3. Copy the best model's submission to \"submission.csv\"\n    try:\n        shutil.copyfile(best_submission_path, final_submission_path)\n        print(f\"\\n✅ Successfully copied '{best_submission_path}' to '{final_submission_path}'\")\n        \n        # 4. Display the head of the final submission\n        final_df = pd.read_csv(final_submission_path)\n        print(\"\\nFinal submission file head:\")\n        print(final_df.head())\n        \n    except FileNotFoundError:\n        print(f\"Error: Could not find '{best_submission_path}'. Make sure STEP 8 ran correctly.\")\n    except Exception as e:\n        print(f\"An error occurred while copying the file: {e}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}