{"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":30918,"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 shutil\nimport os\nimport gc\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-23T06:26:57.816881Z","iopub.execute_input":"2025-04-23T06:26:57.817737Z","iopub.status.idle":"2025-04-23T06:26:57.822530Z","shell.execute_reply.started":"2025-04-23T06:26:57.817704Z","shell.execute_reply":"2025-04-23T06:26:57.821725Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"CONFIG = {\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-linear-warmup-mix-up-crop-noise-mask-aug\", # for wandb logging\n\n    \"sample_rate\": 16000, # should match the pretrained model's sample rate\n    \"num_labels\": 3,\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,    # 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_param\": 48,  # Number of frequency bins to mask\n    \"time_mask_param\": 192, # 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\": \"linear\",\n    \"warmup_ratio\": 0.1,\n    \"learning_rate\": 5e-5,\n    \"weight_decay\": 0.01,\n    \"class_weights\": [1, 1, 1], # [1.2976, 0.7633, 1.0878] for weighted loss\n    \"label_smoothing_factor\": 0.0,\n    \"fp16\": False,\n    \"report_to\": \"wandb\",\n}\n\n# Set random seed for reproducibility\nset_seed(CONFIG[\"seed\"])","metadata":{"execution":{"iopub.status.busy":"2025-04-23T06:26:57.824000Z","iopub.execute_input":"2025-04-23T06:26:57.824327Z","iopub.status.idle":"2025-04-23T06:26:57.847482Z","shell.execute_reply.started":"2025-04-23T06:26:57.824297Z","shell.execute_reply":"2025-04-23T06:26:57.846750Z"},"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-23T06:26:57.848800Z","iopub.execute_input":"2025-04-23T06:26:57.849089Z","iopub.status.idle":"2025-04-23T06:27:06.662104Z","shell.execute_reply.started":"2025-04-23T06:26:57.849064Z","shell.execute_reply":"2025-04-23T06:27:06.661403Z"},"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-23T06:27:06.663175Z","iopub.execute_input":"2025-04-23T06:27:06.663752Z","iopub.status.idle":"2025-04-23T06:27:06.667620Z","shell.execute_reply.started":"2025-04-23T06:27:06.663727Z","shell.execute_reply":"2025-04-23T06:27:06.666813Z"},"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-23T06:27:06.668556Z","iopub.execute_input":"2025-04-23T06:27:06.668855Z","iopub.status.idle":"2025-04-23T06:27:06.716643Z","shell.execute_reply.started":"2025-04-23T06:27:06.668827Z","shell.execute_reply":"2025-04-23T06:27:06.715951Z"},"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-23T06:27:06.717484Z","iopub.execute_input":"2025-04-23T06:27:06.717698Z","iopub.status.idle":"2025-04-23T06:27:06.729118Z","shell.execute_reply.started":"2025-04-23T06:27:06.717679Z","shell.execute_reply":"2025-04-23T06:27:06.728115Z"},"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-23T06:27:06.731524Z","iopub.execute_input":"2025-04-23T06:27:06.731802Z","iopub.status.idle":"2025-04-23T06:27:06.983481Z","shell.execute_reply.started":"2025-04-23T06:27:06.731780Z","shell.execute_reply":"2025-04-23T06:27:06.982671Z"},"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-23T06:27:06.985225Z","iopub.execute_input":"2025-04-23T06:27:06.985555Z","iopub.status.idle":"2025-04-23T06:27:06.997926Z","shell.execute_reply.started":"2025-04-23T06:27:06.985524Z","shell.execute_reply":"2025-04-23T06:27:06.997119Z"},"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-23T06:27:06.998827Z","iopub.execute_input":"2025-04-23T06:27:06.999119Z","iopub.status.idle":"2025-04-23T06:27:08.969928Z","shell.execute_reply.started":"2025-04-23T06:27:06.999089Z","shell.execute_reply":"2025-04-23T06:27:08.968621Z"},"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-23T06:27:08.970937Z","iopub.execute_input":"2025-04-23T06:27:08.971204Z","iopub.status.idle":"2025-04-23T06:27:11.484147Z","shell.execute_reply.started":"2025-04-23T06:27:08.971180Z","shell.execute_reply":"2025-04-23T06:27:11.483321Z"},"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-23T06:27:11.485064Z","iopub.execute_input":"2025-04-23T06:27:11.485416Z","iopub.status.idle":"2025-04-23T06:27:11.489653Z","shell.execute_reply.started":"2025-04-23T06:27:11.485382Z","shell.execute_reply":"2025-04-23T06:27:11.488880Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn.functional as F\n\nclass 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-23T06:27:11.490438Z","iopub.execute_input":"2025-04-23T06:27:11.490740Z","iopub.status.idle":"2025-04-23T06:27:11.515115Z","shell.execute_reply.started":"2025-04-23T06:27:11.490711Z","shell.execute_reply":"2025-04-23T06:27:11.514291Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Hyperparameter Tuning using wandb sweep\n\n# ──────────────────────────────────────────────────────────────────────────────\n# 1) Define sweep config with all hyperparameters\n# ──────────────────────────────────────────────────────────────────────────────\nsweep_config = {\n    \"method\": \"bayes\",\n    \"metric\": {\n        \"name\": \"eval_f1_macro\", \n        \"goal\": \"maximize\"\n    },\n    \"parameters\": {\n        \"num_epochs\": {\n            # use 10 instead of 20 epochs to save time\n            \"values\": [10]   \n        },\n        \"fp16\": {\n            \"values\": [False]\n        },\n        \"lr_scheduler_type\": {\n            \"values\": [\"constant\", \"linear\", \"cosine_with_restarts\"]\n        },\n        \"num_cycles\": {\n            # for lr_scheduler_type == cosine_with_restarts\n            \"values\": [1, 2, 3]\n        },\n        \"learning_rate\": {\n            \"distribution\": \"log_uniform_values\",\n            \"min\": 1e-6,\n            \"max\": 1e-4\n        },\n        \"weight_decay\": {\n            \"distribution\": \"log_uniform_values\",\n            \"min\": 1e-4,\n            \"max\": 1e-1\n        },\n        \"warmup_ratio\": {\n            \"values\": [0.0, 0.1, 0.2]\n        },\n        \"class_weights\": {\n            \"values\": [[1.0, 1.0, 1.0], [1.2976, 0.7633, 1.0878]]\n        },\n        \"label_smoothing_factor\": {\n            \"values\": [0.0, 0.1, 0.2]\n        },\n        \"mixup_alpha\": {\n            \"distribution\": \"log_uniform_values\",\n            \"min\": 0.1,\n            \"max\": 20.0\n        },\n        \"mixup_rate\": {\n            \"values\": [0.0, 0.5, 1.0]\n        },\n        \"num_mel_bins\": {\n            \"values\": [64, 128, 256]\n        },\n        \"max_length\": {\n            \"values\": [512, 1024, 1536]\n        },\n        \"do_random_crop\": {\n            \"values\": [True, False]\n        },\n        \"time_shift_pct\": {\n            \"distribution\": \"uniform\",\n            \"min\": 0.0,\n            \"max\": 1.0\n        },\n        \"add_noise\": {\n            \"values\": [True, False]\n        },\n        \"noise_std\": {\n            \"distribution\": \"log_uniform_values\",\n            \"min\": 1e-5,\n            \"max\": 1e-2\n        },\n        \"freq_mask_frac\": {\n            \"distribution\": \"uniform\",\n            \"min\": 0.0,\n            \"max\": 0.9\n        },\n        \"time_mask_frac\": {\n            \"distribution\": \"uniform\",\n            \"min\": 0.0,\n            \"max\": 0.9\n        }\n    }\n}\n\n# sweep_id = wandb.sweep(sweep=sweep_config, project=\"comp562-cough-sound-classification\")\nsweep_id = \"mtu964xg\" # to resume a past sweep\n\n\n# ──────────────────────────────────────────────────────────────────────────────\n# 2) Define the sweep “agent” function\n# ──────────────────────────────────────────────────────────────────────────────\ndef sweep_train():\n    run = wandb.init()\n    cfg = run.config\n\n    gc.collect()\n    torch.cuda.empty_cache()\n\n    # 2.1) Re-create the feature extractor with the swept max_length\n    feature_extractor = ASTFeatureExtractor.from_pretrained(\n        CONFIG['hf_model_name'], \n        num_mel_bins=cfg.num_mel_bins,\n        max_length=cfg.max_length,\n    )\n\n    # 2.2) Re-create train/val datasets with swept augment params\n    freq_mask_param = int(cfg.freq_mask_frac * cfg.num_mel_bins)\n    time_mask_param = int(cfg.time_mask_frac * cfg.max_length)\n\n    train_ds = CoughASTDataset(\n        train_df, AUDIO_DIR, feature_extractor,\n        training=True,\n        do_random_crop=cfg.do_random_crop,\n        time_shift_pct=cfg.time_shift_pct,\n        add_noise=cfg.add_noise,\n        noise_std=cfg.noise_std,\n        freq_mask_param=freq_mask_param,\n        time_mask_param=time_mask_param,\n    )\n    val_ds = CoughASTDataset(\n        val_df, AUDIO_DIR, feature_extractor,\n        training=False,\n        do_random_crop=cfg.do_random_crop,\n        time_shift_pct=cfg.time_shift_pct,\n        add_noise=cfg.add_noise,\n        noise_std=cfg.noise_std,\n        freq_mask_param=freq_mask_param,\n        time_mask_param=time_mask_param,\n    )\n\n    # 2.3) Build TrainingArguments with swept hyperparams\n    lr_kwargs = {}\n    if cfg.lr_scheduler_type == \"cosine_with_restarts\":\n        lr_kwargs[\"num_cycles\"] = cfg.num_cycles\n\n    total_dim = cfg.max_length * cfg.num_mel_bins\n    if total_dim in [1536 * 256]:\n        batch_size = 6\n    elif total_dim in [1536 * 128, 1536 * 64, 1024 * 256]:\n        batch_size = 8\n    else:\n        batch_size = 16\n\n    training_args = TrainingArguments(\n        output_dir=\"./sweep_results\",\n        eval_strategy=\"epoch\",\n        save_strategy=\"epoch\",\n        learning_rate=cfg.learning_rate,\n        weight_decay=cfg.weight_decay,\n        per_device_train_batch_size=batch_size,\n        per_device_eval_batch_size=batch_size,\n        num_train_epochs=cfg.num_epochs,\n        warmup_ratio=cfg.warmup_ratio,\n        lr_scheduler_type=cfg.lr_scheduler_type,\n        lr_scheduler_kwargs=lr_kwargs,\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=cfg.fp16,\n        report_to=\"wandb\",\n        run_name=run.name,\n    )\n\n    # 2.4) Instantiate the model with swept spectrogram dims\n    model = ASTForAudioClassification.from_pretrained(\n        CONFIG[\"hf_model_name\"],\n        num_labels=CONFIG[\"num_labels\"],\n        num_mel_bins=cfg.num_mel_bins,\n        max_length=cfg.max_length,\n        ignore_mismatched_sizes=True\n    )\n\n    # 2.5) Choose class weights from the sweep\n    cw = cfg.class_weights\n\n    # 2.6) Create and run the MixupTrainer\n    trainer = 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=cw,\n        mixup_alpha=cfg.mixup_alpha,\n        mixup_rate=cfg.mixup_rate,\n    )\n\n    try:\n        trainer.train()\n        metrics = trainer.evaluate()\n        wandb.log(metrics)\n        \n    finally:\n        # ─── delete saved checkpoints ───────────────────────────\n        ckpt_dir = training_args.output_dir\n        shutil.rmtree(ckpt_dir, ignore_errors=True)\n\n        # ─── finish W&B run ─────────────────────────────────────\n        run.finish()\n\n        # ─── memory cleanup ─────────────────────────────────────\n        del trainer, model, train_ds, val_ds, feature_extractor\n        gc.collect()\n        torch.cuda.empty_cache()\n\n\n# ──────────────────────────────────────────────────────────────────────────────\n# 3) Launch the sweep agent for N trials\n# ──────────────────────────────────────────────────────────────────────────────\nwandb.agent(sweep_id, function=sweep_train, count=20)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-23T06:27:46.275998Z","iopub.execute_input":"2025-04-23T06:27:46.276368Z","iopub.status.idle":"2025-04-23T06:34:47.993751Z","shell.execute_reply.started":"2025-04-23T06:27:46.276339Z","shell.execute_reply":"2025-04-23T06:34:47.992824Z"},"collapsed":true,"jupyter":{"outputs_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# # Define training arguments and trainer\n# training_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\n# trainer = 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-23T06:27:11.537429Z","iopub.status.idle":"2025-04-23T06:27:11.537701Z","shell.execute_reply":"2025-04-23T06:27:11.537580Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# # Train the model\n# trainer.train()","metadata":{"execution":{"iopub.status.busy":"2025-04-23T06:27:11.538585Z","iopub.status.idle":"2025-04-23T06:27:11.539286Z","shell.execute_reply":"2025-04-23T06:27:11.539113Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# # Evaluate the final model on the validation set (without TTA)\n# results = trainer.evaluate()\n# results","metadata":{"execution":{"iopub.status.busy":"2025-04-23T06:27:11.539923Z","iopub.status.idle":"2025-04-23T06:27:11.540191Z","shell.execute_reply":"2025-04-23T06:27:11.540065Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# # Evaluate the final model on the validation set (with TTA)\n# def 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\n# tta_logits = tta_predict(trainer, val_ds, n=5)\n# predictions = np.argmax(tta_logits, axis=1)\n# val_f1_tta = f1_score(val_ds.df['disease'], predictions, average='macro')\n# val_f1 = results['eval_f1_macro']\n# print(f\"Validation F1 score (TTA): {val_f1_tta:.4f}\")\n# print(f\"Validation F1 score (no TTA): {val_f1:.4f}\")\n# print(f\"F1 score improvement from TTA: {val_f1_tta - val_f1:.4f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-23T06:27:11.540740Z","iopub.status.idle":"2025-04-23T06:27:11.541061Z","shell.execute_reply":"2025-04-23T06:27:11.540905Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Inference","metadata":{}},{"cell_type":"code","source":"# # Load the test dataset and make predictions\n# test_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\n# test_pred = trainer.predict(test_ds)\n# predictions = np.argmax(test_pred.predictions, axis=1)","metadata":{"execution":{"iopub.status.busy":"2025-04-23T06:27:11.541636Z","iopub.status.idle":"2025-04-23T06:27:11.541924Z","shell.execute_reply":"2025-04-23T06:27:11.541816Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# # Create submission file\n# submission_df = pd.DataFrame({\n#     'candidateID': df_test['candidateID'],\n#     'disease': predictions\n# })\n# submission_df.to_csv(\"submission.csv\", index=False)","metadata":{"execution":{"iopub.status.busy":"2025-04-23T06:27:11.542709Z","iopub.status.idle":"2025-04-23T06:27:11.543059Z","shell.execute_reply":"2025-04-23T06:27:11.542945Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# # Try TTA on the test set\n# tta_logits = tta_predict(trainer, test_ds, n=5)\n# predictions = np.argmax(tta_logits, axis=1)\n# submission_df = pd.DataFrame({\n#     'candidateID': df_test['candidateID'],\n#     'disease': predictions\n# })\n# submission_df.to_csv(\"submission_tta.csv\", index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-23T06:27:11.543904Z","iopub.status.idle":"2025-04-23T06:27:11.544229Z","shell.execute_reply":"2025-04-23T06:27:11.544072Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# wandb.finish()","metadata":{"execution":{"iopub.status.busy":"2025-04-23T06:27:11.545040Z","iopub.status.idle":"2025-04-23T06:27:11.545396Z","shell.execute_reply":"2025-04-23T06:27:11.545247Z"},"trusted":true},"outputs":[],"execution_count":null}]}