{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":129271,"databundleVersionId":15506307},{"sourceType":"datasetVersion","sourceId":14918259,"datasetId":9545488,"databundleVersionId":15784618}],"dockerImageVersionId":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip -q install -U \"pyannote.audio>=3.1\" soundfile tqdm","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# 2-GPU parallel diarization inference (community-1 + finetuned segmentation best.ckpt)\n# - community-1 pipeline returns DiarizeOutput in pyannote.audio 4.x\n#   -> Annotation is at output.speaker_diarization\n# - Uses in-memory waveform dict (avoids torchcodec decoding)\n# - Fixes PyTorch 2.6+ \"weights_only\" unpickling via safe_globals + weights_only=False\n# - Saves submission CSV: filename, diarization(JSON list)\n# ============================================================\n\nimport os, glob, json, warnings\nimport numpy as np\nimport pandas as pd\n\nos.environ[\"PYANNOTE_METRICS_ENABLED\"] = \"0\"\nwarnings.filterwarnings(\"ignore\", message=\".*torchcodec is not installed correctly.*\")\n\nimport torch\nimport torch.nn.functional as F\nimport soundfile as sf\nfrom tqdm.auto import tqdm\nfrom concurrent.futures import ThreadPoolExecutor\n\nfrom pyannote.audio import Pipeline, Inference\nfrom pyannote.audio.core.io import Audio as TorchcodecAudio\n\n# ----------------------------\n# CONFIG\n# ----------------------------\nPIPE_REPO = \"pyannote/speaker-diarization-community-1\"\nFINETUNED_SEG_CKPT = \"/kaggle/input/best-version-community1/best (1).ckpt\"\n\nAUDIO_DIR = \"/kaggle/input/competitions/dl-sprint-4-0-bengali-speaker-diarization-challenge/diarization/diarization/test/audio\"\nOUT_CSV   = \"/kaggle/working/submission_seg7.csv\"\n\nHF_TOKEN =\"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\"\n\nMERGE_GAP_SEC = 0.20\nMIN_SEG_SEC   = 0.10\n\nMIN_SPEAKERS = None\nMAX_SPEAKERS = None\n\n# ----------------------------\n# Sanity checks\n# ----------------------------\nif not os.path.exists(AUDIO_DIR):\n    raise FileNotFoundError(f\"AUDIO_DIR not found: {AUDIO_DIR}\")\n\nwav_files = sorted(glob.glob(os.path.join(AUDIO_DIR, \"*.wav\")))\nif not wav_files:\n    raise FileNotFoundError(f\"No .wav files found in {AUDIO_DIR}\")\n\nif not os.path.exists(FINETUNED_SEG_CKPT):\n    raise FileNotFoundError(f\"FINETUNED_SEG_CKPT not found: {FINETUNED_SEG_CKPT}\")\n\nif not HF_TOKEN:\n    raise RuntimeError(\"HF_TOKEN not found. Add it in Kaggle Secrets as HF_TOKEN=hf_...\")\n\n# ----------------------------\n# PyTorch 2.6+ SAFE GLOBALS allowlist\n# ----------------------------\ntry:\n    import torch.serialization\n    from pyannote.audio.core.task import Specifications\n    torch.serialization.add_safe_globals([Specifications])\n    print(\"✅ Added safe global: pyannote.audio.core.task.Specifications\")\nexcept Exception as e:\n    print(\"⚠️ Could not add safe globals (will still try weights_only=False):\", repr(e))\n\n# ----------------------------\n# SoundFile-based Audio backend (bypasses torchcodec if pipeline crops from paths)\n# ----------------------------\nclass SoundFileAudio(TorchcodecAudio):\n    def get_duration(self, file):\n        file = self.validate_file(file)\n        if \"waveform\" in file:\n            frames = int(file[\"waveform\"].shape[-1])\n            return frames / int(file[\"sample_rate\"])\n        info = sf.info(file[\"audio\"])\n        return float(info.frames) / float(info.samplerate)\n\n    def __call__(self, file):\n        file = self.validate_file(file)\n        channel = file.get(\"channel\", None)\n        if \"waveform\" in file:\n            return super().__call__(file)\n        wav, sr = sf.read(file[\"audio\"], dtype=\"float32\", always_2d=True)  # (T,C)\n        waveform = torch.from_numpy(wav.T)  # (C,T)\n        return self.downmix_and_resample(waveform, int(sr), channel=channel)\n\n    def crop(self, file, segment, mode=\"raise\"):\n        file = self.validate_file(file)\n        channel = file.get(\"channel\", None)\n        if \"waveform\" in file:\n            return super().crop(file, segment, mode=mode)\n\n        path = file[\"audio\"]\n        with sf.SoundFile(path) as f:\n            sr = int(f.samplerate)\n            num_frames = int(len(f))\n            duration = num_frames / sr\n\n            start = float(segment.start)\n            end = float(segment.end)\n\n            start_sample = self.get_num_samples(start, sr)\n            pad_start = max(0, -start_sample)\n            if start < 0:\n                if mode == \"raise\":\n                    raise ValueError(f\"negative start time {start:.3f}s\")\n                start = 0.0\n                start_sample = 0\n\n            end_sample = self.get_num_samples(end, sr)\n            pad_end = max(end_sample, num_frames) - num_frames\n            if end > duration:\n                if mode == \"raise\":\n                    raise ValueError(f\"end time {end:.3f}s > duration {duration:.3f}s\")\n                end = duration\n                end_sample = num_frames\n                pad_end = 0\n\n            read_start = max(0, start_sample)\n            read_end = min(num_frames, end_sample)\n            frames_to_read = max(0, read_end - read_start)\n\n            f.seek(read_start)\n            data = f.read(frames_to_read, dtype=\"float32\", always_2d=True)  # (T,C)\n            waveform = torch.from_numpy(data.T)  # (C,T)\n\n        if pad_start or pad_end:\n            waveform = F.pad(waveform, (pad_start, pad_end))\n\n        return self.downmix_and_resample(waveform, sr, channel=channel)\n\n# ----------------------------\n# Helpers\n# ----------------------------\ndef merge_and_filter(segments, merge_gap=0.2, min_dur=0.1):\n    if not segments:\n        return []\n    segments = sorted(segments, key=lambda x: (x[\"start_time\"], x[\"end_time\"]))\n    out = []\n    cur = dict(segments[0])\n\n    for s in segments[1:]:\n        gap = s[\"start_time\"] - cur[\"end_time\"]\n        if s[\"speaker_id\"] == cur[\"speaker_id\"] and gap <= merge_gap:\n            cur[\"end_time\"] = max(cur[\"end_time\"], s[\"end_time\"])\n        else:\n            if (cur[\"end_time\"] - cur[\"start_time\"]) >= min_dur:\n                out.append(cur)\n            cur = dict(s)\n\n    if (cur[\"end_time\"] - cur[\"start_time\"]) >= min_dur:\n        out.append(cur)\n    return out\n\ndef get_annotation_from_output(diar_output):\n    # pyannote.audio 4.x: DiarizeOutput -> diar_output.speaker_diarization is Annotation\n    cand = getattr(diar_output, \"speaker_diarization\", None)\n    if cand is not None:\n        if callable(cand):\n            cand = cand()\n        if hasattr(cand, \"itertracks\"):\n            return cand\n\n    # older / other variants\n    for name in [\"diarization\", \"annotation\", \"hypothesis\", \"predicted_annotation\"]:\n        cand = getattr(diar_output, name, None)\n        if cand is not None:\n            if callable(cand):\n                cand = cand()\n            if hasattr(cand, \"itertracks\"):\n                return cand\n\n    # sometimes dict-like\n    for key in [\"speaker_diarization\", \"diarization\", \"annotation\", \"hypothesis\"]:\n        try:\n            cand = diar_output[key]\n            if hasattr(cand, \"itertracks\"):\n                return cand\n        except Exception:\n            pass\n\n    raise TypeError(f\"Could not extract Annotation from diarization output type={type(diar_output)}\")\n\ndef diarize_one_wav(pipe, wav_path: str):\n    # In-memory audio dict (avoids torchcodec decoding)\n    y, sr = sf.read(wav_path, always_2d=False)\n    if y.ndim > 1:\n        y = np.mean(y, axis=1)\n    y = y.astype(np.float32)\n\n    waveform = torch.from_numpy(y).unsqueeze(0)  # (1, T)\n    audio = {\"waveform\": waveform, \"sample_rate\": int(sr)}\n\n    kwargs = {}\n    if MIN_SPEAKERS is not None: kwargs[\"min_speakers\"] = int(MIN_SPEAKERS)\n    if MAX_SPEAKERS is not None: kwargs[\"max_speakers\"] = int(MAX_SPEAKERS)\n\n    with torch.no_grad():\n        out = pipe(audio, **kwargs)\n\n    ann = get_annotation_from_output(out)\n\n    segs = []\n    for turn, _, spk in ann.itertracks(yield_label=True):\n        segs.append({\"start_time\": float(turn.start), \"end_time\": float(turn.end), \"speaker_id\": str(spk)})\n\n    return merge_and_filter(segs, MERGE_GAP_SEC, MIN_SEG_SEC)\n\n# ----------------------------\n# Load community-1 + swap segmentation using best.ckpt\n# ----------------------------\ndef load_pipe_with_finetuned_ckpt(device_str: str):\n    device = torch.device(device_str)\n\n    # Base pipeline\n    p = Pipeline.from_pretrained(PIPE_REPO, token=HF_TOKEN)\n\n    # Preserve original segmentation inference settings\n    orig_bs = getattr(getattr(p, \"_segmentation\", None), \"batch_size\", 32)\n    seg_step = float(getattr(p, \"segmentation_step\", 0.1))\n\n    # Determine segmentation class from base pipeline\n    if not hasattr(p, \"_segmentation\") or not hasattr(p._segmentation, \"model\"):\n        raise RuntimeError(\"Could not find p._segmentation.model (pipeline internals changed).\")\n    SegClass = p._segmentation.model.__class__\n\n    # Load finetuned segmentation checkpoint safely\n    try:\n        seg_model = SegClass.load_from_checkpoint(\n            FINETUNED_SEG_CKPT,\n            map_location=\"cpu\",\n            strict=False,\n            weights_only=False,   # critical for PyTorch 2.6+ default behavior\n        )\n    except TypeError:\n        seg_model = SegClass.load_from_checkpoint(\n            FINETUNED_SEG_CKPT,\n            map_location=\"cpu\",\n            strict=False,\n        )\n\n    seg_model.eval()\n    seg_model.freeze()\n\n    # SoundFile backend for safety\n    sr = int(getattr(getattr(seg_model, \"hparams\", None), \"sample_rate\", 16000))\n    seg_model.audio = SoundFileAudio(sample_rate=sr, mono=\"downmix\")\n\n    seg_dur = float(getattr(getattr(seg_model, \"specifications\", None), \"duration\", 10.0))\n\n    # Move to device\n    p.to(device)\n    seg_model.to(device)\n\n    # Swap segmentation inference\n    p._segmentation = Inference(\n        seg_model,\n        duration=seg_dur,\n        step=seg_step * seg_dur,\n        skip_aggregation=True,\n        batch_size=orig_bs,\n    )\n\n    # Also set pipeline audio backend (path-based ops fallback)\n    p._audio = SoundFileAudio(sample_rate=16000, mono=\"downmix\")\n\n    return p\n\n# ----------------------------\n# Multi-GPU runner\n# ----------------------------\nngpu = torch.cuda.device_count()\nprint(\"Detected GPUs:\", ngpu)\nprint(\"Total wavs:\", len(wav_files))\n\nmid = len(wav_files) // 2\ngroup0 = wav_files[:mid]\ngroup1 = wav_files[mid:]\n\ndef process_group(wavs, device_str: str, position: int):\n    pipe = load_pipe_with_finetuned_ckpt(device_str)\n    out = {}\n    for wav in tqdm(wavs, desc=f\"{device_str}\", position=position, leave=True):\n        base = os.path.splitext(os.path.basename(wav))[0]\n        segs = diarize_one_wav(pipe, wav)\n        out[base] = json.dumps(segs, ensure_ascii=False)\n    return out\n\nresults = {}\nif ngpu >= 2:\n    with ThreadPoolExecutor(max_workers=2) as ex:\n        f0 = ex.submit(process_group, group0, \"cuda:0\", 0)\n        f1 = ex.submit(process_group, group1, \"cuda:1\", 1)\n        results.update(f0.result())\n        results.update(f1.result())\nelif ngpu == 1:\n    results.update(process_group(wav_files, \"cuda:0\", 0))\nelse:\n    results.update(process_group(wav_files, \"cpu\", 0))\n\n# ----------------------------\n# Save submission\n# ----------------------------\nrows = []\nfor wav in wav_files:\n    base = os.path.splitext(os.path.basename(wav))[0]\n    rows.append({\"filename\": base, \"diarization\": results[base]})\n\nsub = pd.DataFrame(rows, columns=[\"filename\", \"diarization\"])\nsub.to_csv(OUT_CSV, index=False)\n\nprint(\"✅ Saved:\", OUT_CSV)\nprint(sub.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-22T10:16:26.639152Z","iopub.execute_input":"2026-02-22T10:16:26.639763Z","iopub.status.idle":"2026-02-22T10:33:39.871683Z","shell.execute_reply.started":"2026-02-22T10:16:26.639734Z","shell.execute_reply":"2026-02-22T10:33:39.870843Z"}},"outputs":[],"execution_count":null}]}