{"cells":[{"cell_type":"code","execution_count":null,"metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","execution":{"iopub.execute_input":"2026-02-15T15:08:13.141162Z","iopub.status.busy":"2026-02-15T15:08:13.140855Z","iopub.status.idle":"2026-02-15T15:08:53.048691Z","shell.execute_reply":"2026-02-15T15:08:53.047823Z","shell.execute_reply.started":"2026-02-15T15:08:13.141132Z"},"trusted":true},"outputs":[],"source":"!pip install -q --no-cache-dir pyannote.audio==3.3.1\n# If pip complains about \"dependency conflicts\", IGNORE IT. \n# Conflict warnings are better than an ImportError crash.\n# !pip uninstall -y -q torchcodec\n!pip uninstall -y -q numpy pandas scipy torchcodec\n!pip install -q numpy==2.0.2 pandas scipy torchcodec==0.8.0\n\nprint(\"SUCCESS: Tools installed without breaking the foundation.\")"},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2026-02-15T15:08:53.050831Z","iopub.status.busy":"2026-02-15T15:08:53.050564Z","iopub.status.idle":"2026-02-15T15:09:56.506692Z","shell.execute_reply":"2026-02-15T15:09:56.50601Z","shell.execute_reply.started":"2026-02-15T15:08:53.050802Z"},"trusted":true},"outputs":[],"source":"import os\nimport pandas as pd\nfrom pathlib import Path\nimport wave\nimport contextlib\n\nWORKING_DIR = Path(\"/kaggle/working\")\n# Ensure the directory for checkpoints exists\n(WORKING_DIR / \"checkpoints\").mkdir(parents=True, exist_ok=True)\n\n# --- 1. Auto-Locate directories ---\naudio_dir = None\ncsv_dir = None\n\nfor root, dirs, files in os.walk(\"/kaggle/input\"):\n    if \"audio\" in root.lower() and any(f.endswith(\".wav\") for f in files):\n        audio_dir = Path(root)\n    if \"annotation\" in root.lower() and any(f.endswith(\".csv\") for f in files):\n        csv_dir = Path(root)\n\n# --- 2. Robust Time Conversion ---\ndef time_to_sec(value):\n    if isinstance(value, str) and \":\" in value:\n        parts = value.split(':')\n        if len(parts) == 3:\n            h, m, s = parts\n            return int(h) * 3600 + int(m) * 60 + float(s)\n        elif len(parts) == 2:\n            m, s = parts\n            return int(m) * 60 + float(s)\n    return float(value)\n\n# --- 3. Create Files ---\nrttm_lines = []\nuem_lines = []\ntrain_list = []\nwav_files = sorted(audio_dir.glob(\"*.wav\"))\n\nfor audio_file in wav_files:\n    uri = audio_file.stem  \n    train_list.append(uri)\n\n    with contextlib.closing(wave.open(str(audio_file), 'r')) as f:\n        frames = f.getnframes()\n        rate = f.getframerate()\n        duration_full = frames / float(rate)\n    \n    # --- 4. FORMAT UEM (filename NA start end) ---\n    # As per your requirement: filename NA 0.000 {duration}\n    uem_line = f\"{uri} NA 0.000 {duration_full:.3f}\"\n    uem_lines.append(uem_line)\n\n    \n    csv_path = csv_dir / f\"{uri}.csv\"\n    \n    if csv_path.exists():\n        df = pd.read_csv(csv_path)\n        start_col = next((c for c in df.columns if 'start' in c.lower()), None)\n        end_col = next((c for c in df.columns if 'end' in c.lower()), None)\n        spk_col = next((c for c in df.columns if 'speaker' in c.lower()), None)\n        \n        if not start_col or not end_col:\n            continue\n\n        for _, row in df.iterrows():\n            start = time_to_sec(row[start_col])\n            end = time_to_sec(row[end_col])\n            duration = max(0.1, end - start)\n            spk_id = str(row[spk_col]).replace(\"SPEAKER_\", \"\").strip() if spk_col else \"1\"\n            # RTTM Format: SPEAKER {uri} 1 {start} {duration} <NA> <NA> {label} <NA> <NA>\n            rttm_line = f\"SPEAKER {uri} 1 {start:.3f} {duration:.3f} <NA> <NA> {spk_id} <NA> <NA>\"\n            rttm_lines.append(rttm_line)\n\n# Write output files\nwith open(WORKING_DIR / \"train_list.lst\", \"w\") as f: f.write(\"\\n\".join(train_list))\nwith open(WORKING_DIR / \"train.rttm\", \"w\") as f: f.write(\"\\n\".join(rttm_lines))\nwith open(WORKING_DIR / \"train.uem\", \"w\") as f: \n    f.write(\"\\n\".join(uem_lines))\n\n# --- 4. Create database.yml (Corrected for FileFinder) ---\ndatabase_yaml = f\"\"\"\nDatabases:\n  BengaliData: {audio_dir}/{{uri}}.wav\n\nProtocols:\n  BengaliData:\n    SpeakerDiarization:\n      Train:\n        scope: file\n        train:\n          uri: {WORKING_DIR}/train_list.lst\n          annotation: {WORKING_DIR}/train.rttm\n          annotated: {WORKING_DIR}/train.uem\n\"\"\"\n# annotated: {WORKING_DIR}/train_list.lst\nwith open(WORKING_DIR / \"database.yml\", \"w\") as f: f.write(database_yaml.strip())\n\nprint(f\"✅ Success! Processed {len(train_list)} files. YAML created at {WORKING_DIR}/database.yml\")"},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":""},{"cell_type":"code","execution_count":null,"metadata":{"collapsed":true,"execution":{"iopub.execute_input":"2026-02-15T15:09:56.508019Z","iopub.status.busy":"2026-02-15T15:09:56.5077Z","iopub.status.idle":"2026-02-15T15:10:19.194059Z","shell.execute_reply":"2026-02-15T15:10:19.193344Z","shell.execute_reply.started":"2026-02-15T15:09:56.507985Z"},"jupyter":{"outputs_hidden":true},"trusted":true},"outputs":[],"source":"!ls -R /kaggle"},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2026-02-15T15:10:19.196697Z","iopub.status.busy":"2026-02-15T15:10:19.196093Z","iopub.status.idle":"2026-02-15T15:10:34.439273Z","shell.execute_reply":"2026-02-15T15:10:34.438494Z","shell.execute_reply.started":"2026-02-15T15:10:19.196656Z"},"trusted":true},"outputs":[],"source":"import os\nimport torch\nfrom kaggle_secrets import UserSecretsClient\nfrom huggingface_hub import login, HfApi\nfrom pyannote.audio import Model\nfrom pyannote.audio.tasks import Segmentation \nfrom pyannote.database import FileFinder, get_protocol, registry\nfrom pytorch_lightning import Trainer, Callback\nfrom pytorch_lightning.callbacks import ModelCheckpoint\n\n# --- 1. Authentication ---\nuser_secrets = UserSecretsClient()\nhf_token = \"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\"\nlogin(token=hf_token)\n\n# --- 2. Setup Registry ---\nYAML_PATH = \"/kaggle/working/database.yml\"\nos.environ[\"PYANNOTE_DATABASE_CONFIG\"] = YAML_PATH\nregistry.load_database(YAML_PATH)\n]\n# Use FileFinder to connect URIs in the RTTM to the actual .wav files\npreprocessors = {\"audio\": FileFinder()}\nprotocol = registry.get_protocol(\"BengaliData.SpeakerDiarization.Train\", preprocessors=preprocessors)\nfor file in protocol.train():\n   print(file[\"uri\"])\n   assert \"annotation\" in file\n\n\n# --- 3. Custom Hub Sync Callback ---\nclass PushToHubCallback(Callback):\n    def __init__(self, repo_id, token):\n        super().__init__()\n        self.repo_id = repo_id\n        self.token = token\n        self.api = HfApi()\n\n    def on_train_epoch_end(self, trainer, pl_module):\n        epoch = trainer.current_epoch\n        checkpoint_path = f\"model_epoch_{epoch}.ckpt\"\n        trainer.save_checkpoint(checkpoint_path)\n        print(f\"☁️ Uploading epoch {epoch} to Hugging Face...\")\n        try:\n            self.api.upload_file(\n                path_or_fileobj=checkpoint_path,\n                path_in_repo=f\"checkpoint_epoch_{epoch}.ckpt\",\n                repo_id=self.repo_id,\n                token=self.token\n            )\n            print(\"✅ Upload successful.\")\n        except Exception as e:\n            print(f\"⚠️ Upload failed: {e}\")"},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2026-02-15T15:10:34.440322Z","iopub.status.busy":"2026-02-15T15:10:34.440094Z","iopub.status.idle":"2026-02-15T15:11:04.652617Z","shell.execute_reply":"2026-02-15T15:11:04.651841Z","shell.execute_reply.started":"2026-02-15T15:10:34.440298Z"},"trusted":true},"outputs":[],"source":"# --- 3.5 The SOTA Trinity Augmentation ---\nfrom torch_audiomentations import Compose, AddBackgroundNoise, ApplyImpulseResponse, PitchShift, Gain, OneOf\nimport os\n\nSAMPLE_RATE = 16000  # pyannote uses 16kHz\n\ndef get_augmentations(p=0.5):\n    \"\"\"\n    Defines the 'Trinity' of augmentations for Speaker Diarization.\n    \"\"\"\n    # 1. Locate Datasets (corrected for Kaggle dataset paths)\n    musan_paths = [\n        \"/kaggle/input/datasets/dogrose/musan-dataset/musan/noise\", \n        \"/kaggle/input/datasets/dogrose/musan-dataset/musan/music\", \n        \"/kaggle/input/datasets/dogrose/musan-dataset/musan/speech\"\n    ]\n    rir_path = \"/kaggle/input/datasets/nhattruongdev/rirs-noises/RIRS_NOISES/simulated_rirs\"\n    \n    # Verify paths exist\n    valid_musan = [p for p in musan_paths if os.path.exists(p)]\n    valid_rir = rir_path if os.path.exists(rir_path) else None\n    \n    transforms = []\n    \n    # --- A. Gain (Robustness to volume changes) ---\n    transforms.append(\n        Gain(\n            min_gain_in_db=-15.0,\n            max_gain_in_db=5.0,\n            p=0.5,\n            output_type=\"dict\",\n        )\n    )\n\n    # --- B. Noise (MUSAN) ---\n    if valid_musan:\n        transforms.append(\n            AddBackgroundNoise(\n                background_paths=valid_musan,\n                min_snr_in_db=5.0,\n                max_snr_in_db=15.0,\n                p=p,\n                sample_rate=SAMPLE_RATE,\n                output_type=\"dict\",\n            )\n        )\n        print(f\"✅ MUSAN Noise Augmentation enabled: {len(valid_musan)} sources.\")\n    else:\n        print(\"⚠️ MUSAN dataset not found. Skipping Noise Augmentation.\")\n\n    # --- C. Reverb (RIR) ---\n    if valid_rir:\n        transforms.append(\n            ApplyImpulseResponse(\n                ir_paths=valid_rir,\n                p=p,\n                sample_rate=SAMPLE_RATE,\n                output_type=\"dict\",\n            )\n        )\n        print(\"✅ RIR Reverb Augmentation enabled.\")\n    else:\n        print(\"⚠️ RIR dataset not found. Skipping Reverb Augmentation.\")\n\n    # --- D. Pitch Shift (The On-the-Fly Speed Proxy) ---\n    # This simulates different vocal tracts without changing duration\n    transforms.append(\n        PitchShift(\n            min_transpose_semitones=-2.0,\n            max_transpose_semitones=2.0,\n            sample_rate=SAMPLE_RATE,\n            p=p,\n            output_type=\"dict\",\n        )\n    )\n    \n    return Compose(transforms=transforms, output_type=\"dict\")\n\n# Initialize the augmentation wrapper\naugmentation_pipeline = get_augmentations(p=0.5)"},{"cell_type":"code","execution_count":null,"metadata":{"execution":{"iopub.execute_input":"2026-02-15T15:11:04.653991Z","iopub.status.busy":"2026-02-15T15:11:04.653632Z","iopub.status.idle":"2026-02-15T16:24:31.977447Z","shell.execute_reply":"2026-02-15T16:24:31.976823Z","shell.execute_reply.started":"2026-02-15T15:11:04.653963Z"},"trusted":true},"outputs":[],"source":"# --- 4. Load Model (with PyTorch 2.6 fix) ---\n# Allow pyannote-specific classes to be unpickled\nfrom torch.torch_version import TorchVersion\ntry:\n    from pyannote.audio.core.task import Specifications\n    torch.serialization.add_safe_globals([TorchVersion, Specifications])\nexcept:\n    # Fallback for different pyannote versions\n    torch.serialization.add_safe_globals([TorchVersion])\n\n\n# Force full checkpoint loading (trusted source)\n\n\nfrom functools import partial\n\n# --- FIX FOR RECURSION ERROR ---\n# Only wrap torch.load if we haven't already wrapped it\nif not hasattr(torch, \"_is_wrapped\"):\n    _original_torch_load = torch.load\n\n    def torch_load_no_safe(*args, **kwargs):\n        kwargs[\"weights_only\"] = False\n        return _original_torch_load(*args, **kwargs)\n    \n    torch.load = torch_load_no_safe\n    \n    torch._is_wrapped = True \n\n\nprint(\"📂 Loading pretrained model...\")\nmodel = Model.from_pretrained(\"pyannote/segmentation-3.0\", use_auth_token=hf_token)\n\n# Initialize Task with Trinity Augmentation\ntask = Segmentation(\n    protocol, \n    duration=5.0, \n    batch_size=32, \n    augmentation=augmentation_pipeline,  # <--- INJECT TRINITY HERE\n    num_workers=2  # Need workers to load MUSAN/RIR files in background\n)\nmodel.task = task\n# model.setup(stage=\"fit\") # Usually handled by trainer.fit\n\n# --- 5. Training ---\nREPO_ID = \"ashque12/diarize\" \n\ntrainer = Trainer(\n    accelerator=\"gpu\",\n    devices=1,\n    max_epochs=10,\n    num_sanity_val_steps=0,  # <--- Add this to skip the sanity check\n    limit_val_batches=0.0,   # <--- Add this to disable validation entirely\n    callbacks=[\n        PushToHubCallback(repo_id=REPO_ID, token=hf_token),\n        ModelCheckpoint(dirpath=\"/kaggle/working/checkpoints\", monitor=\"train_loss\", save_top_k=2)\n    ],\n    default_root_dir=\"/kaggle/working/\"\n)\n\nprint(\"🚀 Starting fine-tuning...\")\ntrainer.fit(model)"},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":""}],"metadata":{"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"databundleVersionId":15506307,"sourceId":129271,"sourceType":"competition"},{"datasetId":975000,"sourceId":1648809,"sourceType":"datasetVersion"},{"datasetId":3192482,"sourceId":5539168,"sourceType":"datasetVersion"}],"dockerImageVersionId":31260,"isGpuEnabled":true,"isInternetEnabled":true,"language":"python","sourceType":"notebook"},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.12.12"}},"nbformat":4,"nbformat_minor":4}