{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.10.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":87331,"databundleVersionId":10191418,"sourceType":"competition"}],"dockerImageVersionId":30919,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Library Imports & Configuration","metadata":{}},{"cell_type":"code","source":"import os\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport random\nimport torch\nimport torchaudio\n\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import accuracy_score, f1_score\nfrom transformers import (\n    ASTFeatureExtractor,\n    ASTForAudioClassification,\n    TrainingArguments,\n    Trainer,\n    set_seed,\n)","metadata":{"execution":{"iopub.status.busy":"2025-04-24T04:50:19.570786Z","iopub.execute_input":"2025-04-24T04:50:19.571148Z","iopub.status.idle":"2025-04-24T04:50:43.385419Z","shell.execute_reply.started":"2025-04-24T04:50:19.571113Z","shell.execute_reply":"2025-04-24T04:50:43.384474Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Hyperparameter Sweep on W&B: https://wandb.ai/ncduy0303/comp562-cough-sound-classification/sweeps/mtu964xg\n# Current best run id: fresh-sweep-11 or electric-sweep-47 (we will round some of the found hyperparameters for easier reading)\n\nCONFIG = {\n    \"seed\": 42,\n\n    \"hf_model_name\": \"MIT/ast-finetuned-audioset-10-10-0.4593\",\n    \"saved_model_path\": \"ast-finetuned-cough-classifier\",\n    \"run_name\": \"ast-cosine-warmup-mix-up-crop-noise-mask-aug\", # for wandb logging\n\n    \"num_labels\": 3,\n    \"sample_rate\": 16000,   # Should match the pretrained model's sample rate\n    \"num_mel_bins\": 128,    # Frequency dimension of the spectrograms\n    \"max_length\": 1024,     # Temporal dimension of the spectrograms\n\n    \"do_random_crop\": True, # Whether to random crop before truncation\n    \"time_shift_pct\": 0.0,  # Fraction of the audio length to shift\n    \"add_noise\": True,      # Whether to add noise\n    \"noise_std\": 1e-3,      # Standard deviation of the noise\n    # \"freq_mask_frac\": 0.09, # Fraction of frequency bins to mask\n    # \"time_mask_frac\": 0.05, # Number of time steps to mask\n\n    \"mixup_alpha\": 10,      # Mixup alpha parameter\n    \"mixup_rate\" : 0.5,     # Probability of applying mixup\n\n    \"data_dir\": \"/kaggle/input/airs-ai-in-respiratory-sounds\",\n\n    # TrainingArguments\n    \"num_epochs\": 20,\n    \"batch_size\": 16,\n    \"lr_scheduler_type\": \"cosine\",\n    \"warmup_ratio\": 0.1,\n    \"learning_rate\": 5e-5,\n    \"weight_decay\": 0.01,\n    \"class_weights\": [1, 1, 1],\n    \"label_smoothing_factor\": 0.0,\n    \"fp16\": False,\n    \"report_to\": \"wandb\",\n}\n\n# CONFIG[\"freq_mask_param\"] = int(CONFIG[\"num_mel_bins\"] * CONFIG[\"freq_mask_frac\"])\n# CONFIG[\"time_mask_param\"] = int(CONFIG[\"max_length\"] * CONFIG[\"time_mask_frac\"])\nCONFIG[\"freq_mask_param\"] = 48\nCONFIG[\"time_mask_param\"] = 192\n\n# Set random seed for reproducibility\nset_seed(CONFIG[\"seed\"])","metadata":{"execution":{"iopub.status.busy":"2025-04-24T04:50:43.386492Z","iopub.execute_input":"2025-04-24T04:50:43.387176Z","iopub.status.idle":"2025-04-24T04:50:43.398936Z","shell.execute_reply.started":"2025-04-24T04:50:43.387143Z","shell.execute_reply":"2025-04-24T04:50:43.398117Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nsecret_value_0 = user_secrets.get_secret(\"wandb_key\")\n\nimport wandb\nwandb.login(key=secret_value_0)\n\nimport os\nos.environ[\"WANDB_PROJECT\"]=\"comp562-cough-sound-classification\"\nos.environ[\"WANDB_LOG_MODEL\"]=\"end\"\nos.environ[\"WANDB_WATCH\"]=\"false\"","metadata":{"execution":{"iopub.status.busy":"2025-04-24T04:51:11.184943Z","iopub.execute_input":"2025-04-24T04:51:11.185229Z","iopub.status.idle":"2025-04-24T04:51:19.717991Z","shell.execute_reply.started":"2025-04-24T04:51:11.185206Z","shell.execute_reply":"2025-04-24T04:51:19.717282Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Data","metadata":{}},{"cell_type":"markdown","source":"### Getting the Data","metadata":{}},{"cell_type":"code","source":"DATA_DIR = CONFIG[\"data_dir\"]\nAUDIO_DIR = os.path.join(DATA_DIR, \"sounds\", \"sounds\")\nTRAIN_CSV_PATH = os.path.join(DATA_DIR, \"train.csv\")\nTEST_CSV_PATH = os.path.join(DATA_DIR, \"test.csv\")","metadata":{"execution":{"iopub.status.busy":"2025-04-24T04:51:23.189156Z","iopub.execute_input":"2025-04-24T04:51:23.189825Z","iopub.status.idle":"2025-04-24T04:51:23.193950Z","shell.execute_reply.started":"2025-04-24T04:51:23.189791Z","shell.execute_reply":"2025-04-24T04:51:23.193005Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def is_corrupted(candidateID):\n    \"\"\"\n    Returns True if the .wav file for the given patient cannot be opened by torchaudio.\n    Returns False if it's a valid audio file.\n    \"\"\"\n    audio_path = os.path.join(AUDIO_DIR, candidateID, \"cough.wav\")\n\n    # If file doesn't exist, consider it \"corrupted\"\n    if not os.path.exists(audio_path):\n        return True\n\n    try:\n        # Try loading with torchaudio\n        torchaudio.load(audio_path)\n        # If it succeeds, we consider it not corrupted\n        return False\n    except Exception as e:\n        # If there's an error, it's likely corrupted\n        print(f\"Could not read {audio_path} with torchaudio. Error: {e}\")\n        return True\n\ndf_train = pd.read_csv(TRAIN_CSV_PATH)\ndf_test  = pd.read_csv(TEST_CSV_PATH)\n\n# Remove the corrupted patient from the training set\ndf_train = df_train[df_train[\"candidateID\"] != \"5ee582f2832c2\"]\n\n# # 1) Collect all patient IDs in train & test\n# all_patients = set(df_train[\"candidateID\"].unique()) | set(df_test[\"candidateID\"].unique())\n\n# # 2) Identify corrupted patients\n# corrupted_patients = [pid for pid in all_patients if is_corrupted(pid)]\n\n# # 3) Remove these from both train & test\n# df_train_clean = df_train[~df_train[\"candidateID\"].isin(corrupted_patients)].copy()\n# df_test_clean  = df_test[~df_test[\"candidateID\"].isin(corrupted_patients)].copy()\n\n# print(f\"Total corrupted patients detected: {len(corrupted_patients)}\")\n# print(f\"Clean train set size: {len(df_train_clean)}, Original train set size: {len(df_train)}\")\n# print(f\"Clean test set size:  {len(df_test_clean)}, Original test set size: {len(df_test)}\")\n\n# df_train = df_train_clean\n# df_test = df_test_clean","metadata":{"execution":{"iopub.status.busy":"2025-04-24T04:51:23.333287Z","iopub.execute_input":"2025-04-24T04:51:23.333520Z","iopub.status.idle":"2025-04-24T04:51:23.374085Z","shell.execute_reply.started":"2025-04-24T04:51:23.333500Z","shell.execute_reply":"2025-04-24T04:51:23.373226Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Dataset Class","metadata":{}},{"cell_type":"code","source":"class CoughASTDataset(Dataset):\n    def __init__(\n        self,\n        df: pd.DataFrame,\n        audio_dir: str,\n        feature_extractor: ASTFeatureExtractor,\n        sample_rate: int = 16000,\n        training: bool = True,\n        # ── waveform‑domain aug ────────────────────────────────\n        do_random_crop: bool = False,\n        time_shift_pct: float = 0.3,\n        add_noise: bool = False,\n        noise_std: float = 1e-3,\n        # ── spec‑domain aug (SpecAugment) ─────────────────────\n        freq_mask_param: int = 48,\n        time_mask_param: int = 192,\n    ):\n        \"\"\"\n        Parameters\n        ----------\n        df : DataFrame      expects a 'candidateID' column and a 'disease' int label.\n        audio_dir : str     root that contains <candidateID>/cough.wav\n        feature_extractor   HF ASTFeatureExtractor (fbank, normalization, padding/truncation)\n        sample_rate : int   MUST match the extractor & model (16 kHz for MIT/ast‑*)\n        training : bool     switch to `False` for val/test ⇒ no augmentation\n        \"\"\"\n        self.df                = df.reset_index(drop=True)\n        self.audio_dir         = audio_dir\n        self.feat_extractor    = feature_extractor\n        self.sr                = sample_rate\n        self.training          = training\n\n        # ─ waveform‑level augmenters ─\n        self._do_random_crop   = do_random_crop\n        self._num_samples      = self.feat_extractor.max_length * 160\n        self._time_shift_pct   = time_shift_pct\n        self._add_noise        = add_noise\n        self._noise_std        = noise_std\n\n        # ─ SpecAugment blocks (instantiated once, faster than per‑item) ─\n        self._freq_mask = torchaudio.transforms.FrequencyMasking(freq_mask_param)\n        self._time_mask = torchaudio.transforms.TimeMasking(time_mask_param)\n\n        # one‑off resampler if files aren’t already 16 kHz\n        self._resampler   = torchaudio.transforms.Resample(orig_freq=self.sr, new_freq=self.sr)\n\n    # --------------------------------------------------------\n    # helpers\n    # --------------------------------------------------------\n    def _random_crop(self, wav: torch.Tensor) -> torch.Tensor:\n        \"\"\"Randomly crop a fixed-length segment from the waveform.\"\"\"\n        num_samples = self._num_samples\n        total_len   = wav.shape[-1]\n\n        if total_len <= num_samples:\n            return wav\n        else:\n            start = random.randint(0, total_len - num_samples)\n            return wav[..., start:start + num_samples]\n    \n    def _random_time_shift(self, wav: torch.Tensor) -> torch.Tensor:\n        \"\"\"Roll the waveform left / right by up to `time_shift_pct`.\"\"\"\n        max_shift = int(wav.shape[-1] * self._time_shift_pct)\n        shift     = random.randint(-max_shift, max_shift)\n        return torch.roll(wav, shifts=shift, dims=-1)\n\n    # --------------------------------------------------------\n    # core\n    # --------------------------------------------------------\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx: int):\n        cid   = self.df.loc[idx, \"candidateID\"]\n        label = self.df.loc[idx, \"disease\"] if \"disease\" in self.df.columns else None\n\n        # ────────── load waveform ──────────\n        wav_path          = os.path.join(self.audio_dir, cid, \"cough.wav\")\n        waveform, orig_sr = torchaudio.load(wav_path)\n\n        # mono + resample\n        if waveform.shape[0] > 1:\n            waveform = waveform.mean(dim=0, keepdim=True)\n        if orig_sr != self.sr:\n            waveform = self._resampler(waveform)\n\n        # ────────── waveform‑domain aug ──────────\n        if self.training:\n            # Crop a random chunk of fixed length\n            if self._do_random_crop:\n                waveform = self._random_crop(waveform)\n            # Apply small time shift to that chunk\n            waveform = self._random_time_shift(waveform)\n            # additive noise\n            if self._add_noise:\n                waveform += torch.randn_like(waveform) * self._noise_std\n\n        # ────────── extract log‑Mel patches with HF helper ──────────\n        feats = self.feat_extractor(\n            waveform.squeeze().numpy(),\n            sampling_rate=self.sr,\n            return_tensors=\"pt\",\n        )[\"input_values\"].squeeze(0)       # tensor [frames, mel]\n\n        # ────────── SpecAugment (Freq & Time masking) ──────────\n        if self.training:\n            feats = feats.transpose(0, 1)  # torchaudio maskers expect [mel, frames]\n            feats = feats.unsqueeze(0)     # torchaudio maskers expect 3‑D\n            feats = self._freq_mask(feats)\n            feats = self._time_mask(feats)\n            feats = feats.squeeze(0)\n            feats = feats.transpose(0, 1)\n\n        return {\n            \"input_values\": feats,         # [frames, mel]\n            \"label\":        torch.tensor(label, dtype=torch.long) if label is not None else None,\n        }","metadata":{"execution":{"iopub.status.busy":"2025-04-24T04:51:24.707036Z","iopub.execute_input":"2025-04-24T04:51:24.707322Z","iopub.status.idle":"2025-04-24T04:51:24.718590Z","shell.execute_reply.started":"2025-04-24T04:51:24.707298Z","shell.execute_reply":"2025-04-24T04:51:24.717826Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Creating Dataset","metadata":{}},{"cell_type":"code","source":"# Load the feature extractor\nfeature_extractor = ASTFeatureExtractor.from_pretrained(\n    CONFIG['hf_model_name'], \n    max_length=CONFIG['max_length'],\n)\nfeature_extractor","metadata":{"execution":{"iopub.status.busy":"2025-04-24T04:51:26.250922Z","iopub.execute_input":"2025-04-24T04:51:26.251207Z","iopub.status.idle":"2025-04-24T04:51:26.605615Z","shell.execute_reply.started":"2025-04-24T04:51:26.251184Z","shell.execute_reply":"2025-04-24T04:51:26.604820Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Split the dataset into training and validation sets\ntrain_df, val_df = train_test_split(\n    df_train, \n    test_size=0.2, \n    stratify=df_train['disease'],\n    random_state=CONFIG['seed'],\n)\n\n# Create datasets\ntrain_ds = CoughASTDataset(\n    train_df, AUDIO_DIR, feature_extractor, training=True,\n    sample_rate=CONFIG[\"sample_rate\"],\n    do_random_crop=CONFIG[\"do_random_crop\"],\n    time_shift_pct=CONFIG[\"time_shift_pct\"],\n    add_noise=CONFIG[\"add_noise\"],\n    noise_std=CONFIG[\"noise_std\"],\n    freq_mask_param=CONFIG[\"freq_mask_param\"],\n    time_mask_param=CONFIG[\"time_mask_param\"],\n)\nval_ds = CoughASTDataset(\n    val_df, AUDIO_DIR, feature_extractor, training=False,\n    sample_rate=CONFIG[\"sample_rate\"],\n    do_random_crop=CONFIG[\"do_random_crop\"],\n    time_shift_pct=CONFIG[\"time_shift_pct\"],\n    add_noise=CONFIG[\"add_noise\"],\n    noise_std=CONFIG[\"noise_std\"],\n    freq_mask_param=CONFIG[\"freq_mask_param\"],\n    time_mask_param=CONFIG[\"time_mask_param\"],\n)","metadata":{"execution":{"iopub.status.busy":"2025-04-24T04:51:26.606929Z","iopub.execute_input":"2025-04-24T04:51:26.607249Z","iopub.status.idle":"2025-04-24T04:51:26.619713Z","shell.execute_reply.started":"2025-04-24T04:51:26.607214Z","shell.execute_reply":"2025-04-24T04:51:26.618790Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Check if Dataset is working\nidx = random.randint(0, len(train_ds) - 1)\nsample = train_ds[idx]\nprint(sample)\nprint(sample[\"input_values\"].shape) # (time_frames, mel_bins)\n\n# Plot the specialtrogram for sanity check\nplt.figure(figsize=(16, 8))\nplt.suptitle(f\"Spectrograms for train index: {idx}\", fontsize=18)\nfor i in range(4):\n    sample = train_ds[idx]  # Each call applies augmentation if training=True\n    plt.subplot(2, 2, i + 1)\n    plt.imshow(sample[\"input_values\"].numpy().T, origin=\"lower\", aspect=\"auto\")\n    plt.title(f\"Augmentation #{i+1}\")\n    plt.xlabel(\"Time Frames\")\n    plt.ylabel(\"Mel-Frequency Bins\")\n    plt.colorbar(label=\"Normalized Amplitude\")\nplt.tight_layout(rect=[0, 0, 1, 0.96])\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2025-04-24T04:51:28.247165Z","iopub.execute_input":"2025-04-24T04:51:28.247468Z","iopub.status.idle":"2025-04-24T04:51:30.012591Z","shell.execute_reply.started":"2025-04-24T04:51:28.247441Z","shell.execute_reply":"2025-04-24T04:51:30.011591Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"# Vision Transformer Model\nmodel = ASTForAudioClassification.from_pretrained(\n    CONFIG['hf_model_name'],\n    num_labels=CONFIG['num_labels'],\n    num_mel_bins=CONFIG['num_mel_bins'],\n    max_length=CONFIG['max_length'],\n    ignore_mismatched_sizes=True\n)\nmodel","metadata":{"execution":{"iopub.status.busy":"2025-04-24T04:51:33.767499Z","iopub.execute_input":"2025-04-24T04:51:33.767819Z","iopub.status.idle":"2025-04-24T04:51:36.471146Z","shell.execute_reply.started":"2025-04-24T04:51:33.767793Z","shell.execute_reply":"2025-04-24T04:51:36.470265Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"# Use accuracy and macro F1 score as metrics\ndef compute_metrics(eval_pred):\n    logits, labels = eval_pred\n    predictions = np.argmax(logits, axis=-1)\n    return {\n        \"accuracy\": accuracy_score(labels, predictions),\n        \"f1_macro\": f1_score(labels, predictions, average='macro')\n    }","metadata":{"execution":{"iopub.status.busy":"2025-04-24T04:51:36.472334Z","iopub.execute_input":"2025-04-24T04:51:36.472666Z","iopub.status.idle":"2025-04-24T04:51:36.476876Z","shell.execute_reply.started":"2025-04-24T04:51:36.472634Z","shell.execute_reply":"2025-04-24T04:51:36.476033Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class MixupTrainer(Trainer):\n    \"\"\"\n    Implements mixup inside `compute_loss`.\n\n    Parameters\n    ----------\n    weights : list of float\n              class weights for the loss function.\n    mixup_alpha : float (α in the Beta(α, α) distribution);\n                   0 ⇒ no mixup. Default = 10 same as the original AST paper.\n    mixup_rate  : float (probability of applying mixup);\n                   0 ⇒ no mixup. Default = 0.5 same as the original AST paper.\n    \"\"\"\n    def __init__(self, *args, weights: list = None,\n                 mixup_alpha: float = 10, mixup_rate: float = 0.5, **kwargs):\n        super().__init__(*args, **kwargs)\n        self.mixup_alpha = mixup_alpha\n        self.mixup_rate = mixup_rate\n        self.loss_fn = torch.nn.CrossEntropyLoss(label_smoothing=self.args.label_smoothing_factor)\n        if weights is not None:\n            self.loss_fn = torch.nn.CrossEntropyLoss(\n                weight=torch.tensor(weights).float().to(self.args.device),\n                label_smoothing=self.args.label_smoothing_factor\n            )\n\n    # --------------------------------------------------------\n    # override loss computation\n    # --------------------------------------------------------\n    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):\n        labels = inputs.pop(\"labels\")                     # [B]\n        input_values = inputs[\"input_values\"]             # [B, …]\n\n        if model.training and self.mixup_alpha > 0.0 and random.random() < self.mixup_rate:\n            lam    = np.random.beta(self.mixup_alpha, self.mixup_alpha)\n            batch  = labels.size(0)\n            perm   = torch.randperm(batch).to(labels.device)\n\n            # mix the inputs\n            mixed_inputs = lam * input_values + (1.0 - lam) * input_values[perm]\n\n            # convert labels → one‑hot → mix\n            num_classes  = model.config.num_labels\n            one_hot      = torch.nn.functional.one_hot(labels, num_classes).float()\n            one_hot_perm = torch.nn.functional.one_hot(labels[perm], num_classes).float()\n            mixed_labels = lam * one_hot + (1.0 - lam) * one_hot_perm\n\n            outputs = model(**{\"input_values\": mixed_inputs})\n            loss    = self.loss_fn(outputs.logits, mixed_labels)\n\n        else:\n            outputs = model(**inputs)\n            loss    = self.loss_fn(outputs.logits, labels)\n\n        return (loss, outputs) if return_outputs else loss","metadata":{"execution":{"iopub.status.busy":"2025-04-24T04:51:37.485869Z","iopub.execute_input":"2025-04-24T04:51:37.486183Z","iopub.status.idle":"2025-04-24T04:51:37.494239Z","shell.execute_reply.started":"2025-04-24T04:51:37.486155Z","shell.execute_reply":"2025-04-24T04:51:37.493185Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Define training arguments and trainer\ntraining_args = TrainingArguments(\n    output_dir=\"./results\",\n    eval_strategy=\"epoch\",\n    save_strategy=\"epoch\",\n    lr_scheduler_type=CONFIG[\"lr_scheduler_type\"],\n    warmup_ratio=CONFIG[\"warmup_ratio\"],\n    learning_rate=CONFIG['learning_rate'],\n    per_device_train_batch_size=CONFIG['batch_size'],\n    per_device_eval_batch_size=CONFIG['batch_size'],\n    num_train_epochs=CONFIG['num_epochs'],\n    weight_decay=CONFIG[\"weight_decay\"],\n    label_smoothing_factor=CONFIG[\"label_smoothing_factor\"],\n    logging_dir='./logs',\n    logging_steps=10,\n    load_best_model_at_end=True,\n    metric_for_best_model=\"f1_macro\",\n    remove_unused_columns=True,\n    save_total_limit=1,\n    fp16=CONFIG[\"fp16\"],\n    report_to=CONFIG[\"report_to\"],\n    run_name=CONFIG[\"run_name\"]\n)\n\ntrainer = MixupTrainer(\n    model=model,\n    args=training_args,\n    train_dataset=train_ds,\n    eval_dataset=val_ds,\n    compute_metrics=compute_metrics,\n    weights=CONFIG[\"class_weights\"],\n    mixup_alpha=CONFIG[\"mixup_alpha\"],\n    mixup_rate=CONFIG[\"mixup_rate\"],\n)","metadata":{"execution":{"iopub.status.busy":"2025-04-24T04:51:39.713048Z","iopub.execute_input":"2025-04-24T04:51:39.713369Z","iopub.status.idle":"2025-04-24T04:51:40.119412Z","shell.execute_reply.started":"2025-04-24T04:51:39.713338Z","shell.execute_reply":"2025-04-24T04:51:40.118774Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Train the model\ntrainer.train()","metadata":{"execution":{"iopub.status.busy":"2025-04-24T04:51:42.139383Z","iopub.execute_input":"2025-04-24T04:51:42.139677Z","iopub.status.idle":"2025-04-24T04:53:44.840100Z","shell.execute_reply.started":"2025-04-24T04:51:42.139653Z","shell.execute_reply":"2025-04-24T04:53:44.839248Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Evaluate the final model on the validation set (without TTA)\nresults = trainer.evaluate()\nresults","metadata":{"execution":{"iopub.status.busy":"2025-04-24T04:53:44.841365Z","iopub.execute_input":"2025-04-24T04:53:44.841654Z","iopub.status.idle":"2025-04-24T04:53:53.725808Z","shell.execute_reply.started":"2025-04-24T04:53:44.841630Z","shell.execute_reply":"2025-04-24T04:53:53.725183Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Evaluate the final model on the validation set (with TTA)\ndef tta_predict(trainer, dataset, n=5):\n    \"\"\"\n    Runs TTA by averaging predictions over n augmented samples per item.\n    \"\"\"\n    all_logits = []\n    dataset.training = True\n    for tta_round in range(n):\n        # Set dataset to training mode for augmentation\n        preds = trainer.predict(dataset)\n        all_logits.append(preds.predictions)\n    # Average logits\n    avg_logits = np.mean(np.stack(all_logits, axis=0), axis=0)\n    return avg_logits\n\ntta_logits = tta_predict(trainer, val_ds, n=5)\npredictions = np.argmax(tta_logits, axis=1)\nval_f1_tta = f1_score(val_ds.df['disease'], predictions, average='macro')\nval_f1 = results['eval_f1_macro']\nprint(f\"Validation F1 score (TTA): {val_f1_tta:.4f}\")\nprint(f\"Validation F1 score (no TTA): {val_f1:.4f}\")\nprint(f\"F1 score improvement from TTA: {val_f1_tta - val_f1:.4f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-24T04:53:53.727368Z","iopub.execute_input":"2025-04-24T04:53:53.727592Z","iopub.status.idle":"2025-04-24T04:54:28.636663Z","shell.execute_reply.started":"2025-04-24T04:53:53.727571Z","shell.execute_reply":"2025-04-24T04:54:28.635815Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Inference","metadata":{}},{"cell_type":"code","source":"# Load the test dataset and make predictions\ntest_ds = CoughASTDataset(\n    df_test, AUDIO_DIR, feature_extractor, training=False,\n    sample_rate=CONFIG[\"sample_rate\"],\n    do_random_crop=CONFIG[\"do_random_crop\"],\n    time_shift_pct=CONFIG[\"time_shift_pct\"],\n    add_noise=CONFIG[\"add_noise\"],\n    noise_std=CONFIG[\"noise_std\"],\n    freq_mask_param=CONFIG[\"freq_mask_param\"],\n    time_mask_param=CONFIG[\"time_mask_param\"],\n)\n\ntest_pred = trainer.predict(test_ds)\npredictions = np.argmax(test_pred.predictions, axis=1)","metadata":{"execution":{"iopub.status.busy":"2025-04-24T04:54:28.637937Z","iopub.execute_input":"2025-04-24T04:54:28.638156Z","iopub.status.idle":"2025-04-24T04:55:08.068185Z","shell.execute_reply.started":"2025-04-24T04:54:28.638136Z","shell.execute_reply":"2025-04-24T04:55:08.067305Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Create submission file\nsubmission_df = pd.DataFrame({\n    'candidateID': df_test['candidateID'],\n    'disease': predictions\n})\nsubmission_df.to_csv(\"submission.csv\", index=False)","metadata":{"execution":{"iopub.status.busy":"2025-04-24T04:55:08.068956Z","iopub.execute_input":"2025-04-24T04:55:08.069168Z","iopub.status.idle":"2025-04-24T04:55:08.089252Z","shell.execute_reply.started":"2025-04-24T04:55:08.069148Z","shell.execute_reply":"2025-04-24T04:55:08.088665Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Try TTA on the test set\ntta_logits = tta_predict(trainer, test_ds, n=5)\npredictions = np.argmax(tta_logits, axis=1)\nsubmission_df = pd.DataFrame({\n    'candidateID': df_test['candidateID'],\n    'disease': predictions\n})\nsubmission_df.to_csv(\"submission_tta.csv\", index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-24T04:55:08.090079Z","iopub.execute_input":"2025-04-24T04:55:08.090296Z","iopub.status.idle":"2025-04-24T04:56:57.183849Z","shell.execute_reply.started":"2025-04-24T04:55:08.090274Z","shell.execute_reply":"2025-04-24T04:56:57.183231Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"wandb.finish()","metadata":{"execution":{"iopub.status.busy":"2025-04-24T04:56:57.184592Z","iopub.execute_input":"2025-04-24T04:56:57.184904Z","iopub.status.idle":"2025-04-24T04:56:58.893083Z","shell.execute_reply.started":"2025-04-24T04:56:57.184865Z","shell.execute_reply":"2025-04-24T04:56:58.892442Z"},"trusted":true},"outputs":[],"execution_count":null}]}