{"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":[{"sourceId":129271,"databundleVersionId":15506307,"sourceType":"competition"}],"dockerImageVersionId":31260,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"**Install Dependencies**","metadata":{}},{"cell_type":"code","source":"!pip install --quiet speechbrain==1.0.3 soundfile scikit-learn tqdm matplotlib joblib\n!pip install --quiet git+https://github.com/snakers4/silero-vad.git","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-10T12:06:21.265805Z","iopub.execute_input":"2026-02-10T12:06:21.266578Z","iopub.status.idle":"2026-02-10T12:06:33.925733Z","shell.execute_reply.started":"2026-02-10T12:06:21.266548Z","shell.execute_reply":"2026-02-10T12:06:33.925017Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Imports**","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport json\nimport gc\nimport random\nfrom pathlib import Path\nimport soundfile as sf\nimport torch\nimport torch.nn.functional as F\nimport torchaudio\nfrom torch.utils.data import Dataset, DataLoader\nfrom tqdm.auto import tqdm\nfrom sklearn.cluster import AgglomerativeClustering\nfrom sklearn.metrics import silhouette_score\nfrom silero_vad import load_silero_vad, get_speech_timestamps\nfrom collections import defaultdict\nfrom speechbrain.inference import EncoderClassifier\nfrom scipy.signal import medfilt\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\nfrom functools import lru_cache\nimport warnings\nwarnings.filterwarnings('ignore')\n\nprint(f\"NumPy: {np.__version__}\")\nprint(f\"Pandas: {pd.__version__}\")\nprint(f\"PyTorch: {torch.__version__}\")\nprint(f\"CUDA Available: {torch.cuda.is_available()}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-10T12:06:33.927410Z","iopub.execute_input":"2026-02-10T12:06:33.927646Z","iopub.status.idle":"2026-02-10T12:06:33.934866Z","shell.execute_reply.started":"2026-02-10T12:06:33.927622Z","shell.execute_reply":"2026-02-10T12:06:33.934131Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Configuration**","metadata":{}},{"cell_type":"code","source":"class Config:\n\n    TRAIN_AUDIO = Path(\"/kaggle/input/competitions/dl-sprint-4-0-bengali-speaker-diarization-challenge/diarization/diarization/train/audio\")\n    TEST_AUDIO = Path(\"/kaggle/input/competitions/dl-sprint-4-0-bengali-speaker-diarization-challenge/diarization/diarization/test/audio\")\n    TRAIN_CSV = Path(\"/kaggle/input/competitions/dl-sprint-4-0-bengali-speaker-diarization-challenge/diarization/diarization/train/annotation\")\n    OUT_PATH = Path(\"/kaggle/working/submission.csv\")\n    MODEL_PATH = Path(\"/kaggle/working/finetuned_ecapa.pth\")\n\n    DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n    VAD_DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n    \n\n    SR = 16000\n    \n\n    VAD_THRESHOLD = 0.35\n    VAD_PAD_MS = 200\n    VAD_MIN_SPEECH_MS = 250\n    VAD_MIN_SILENCE_MS = 100\n\n    WIN_SHORT = {\"win_sec\": 1.5, \"hop_sec\": 0.5}\n    WIN_MEDIUM = {\"win_sec\": 2.0, \"hop_sec\": 0.75}\n    WIN_LONG = {\"win_sec\": 2.5, \"hop_sec\": 1.0}\n    \n\n    BATCH_SIZE = 32  \n    MAX_SPK = 8\n    MERGE_GAP_SEC = 0.3\n    MIN_TURN_SEC = 0.5\n    MEDIAN_FILTER_SIZE = 5\n    \n\n    USE_FINETUNING = False\n    EPOCHS = 3\n    LR = 1e-4\n    TRAIN_BATCH = 8\n    \n\n    MAX_WORKERS = 4  \n    \n\n    USE_AMP = True  \n    PIN_MEMORY = True\n    \nconfig = Config()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-10T12:06:33.935728Z","iopub.execute_input":"2026-02-10T12:06:33.936027Z","iopub.status.idle":"2026-02-10T12:06:33.951997Z","shell.execute_reply.started":"2026-02-10T12:06:33.935997Z","shell.execute_reply":"2026-02-10T12:06:33.951090Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Optimized Audio Loading**","metadata":{}},{"cell_type":"code","source":"@lru_cache(maxsize=128)\ndef load_audio_cached(wav_path):\n\n    wav, sr = sf.read(str(wav_path), dtype=\"float32\")\n    if sr != config.SR:\n        wav = torchaudio.functional.resample(\n            torch.from_numpy(wav), sr, config.SR\n        ).numpy()\n    return wav.astype(np.float32)\n\ndef load_audio(wav_path):\n\n    wav, sr = sf.read(str(wav_path), dtype=\"float32\")\n    if sr != config.SR:\n        wav = torchaudio.functional.resample(\n            torch.from_numpy(wav), sr, config.SR\n        ).numpy()\n    return wav.astype(np.float32)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-10T12:06:33.952982Z","iopub.execute_input":"2026-02-10T12:06:33.953486Z","iopub.status.idle":"2026-02-10T12:06:33.966133Z","shell.execute_reply.started":"2026-02-10T12:06:33.953465Z","shell.execute_reply":"2026-02-10T12:06:33.965485Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Cached VAD Model**","metadata":{}},{"cell_type":"code","source":"\n_vad_model = None\n\ndef get_vad_model():\n\n    global _vad_model\n    if _vad_model is None:\n        _vad_model = load_silero_vad()\n        _vad_model.to(config.VAD_DEVICE)\n        _vad_model.eval()\n    return _vad_model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-10T12:06:33.967593Z","iopub.execute_input":"2026-02-10T12:06:33.967933Z","iopub.status.idle":"2026-02-10T12:06:33.987752Z","shell.execute_reply.started":"2026-02-10T12:06:33.967901Z","shell.execute_reply":"2026-02-10T12:06:33.987069Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Optimized VAD Segments**","metadata":{}},{"cell_type":"code","source":"def get_vad_segments(wav):\n\n    vad_model = get_vad_model()\n    device = config.VAD_DEVICE\n    \n    with torch.no_grad():\n\n        wav_tensor = torch.tensor(wav, dtype=torch.float32, device= device)\n        segments = get_speech_timestamps(\n            wav_tensor,\n            vad_model,\n            sampling_rate=config.SR,\n            threshold=config.VAD_THRESHOLD,\n            speech_pad_ms=config.VAD_PAD_MS,\n            min_speech_duration_ms=config.VAD_MIN_SPEECH_MS,\n            min_silence_duration_ms=config.VAD_MIN_SILENCE_MS\n        )\n    \n    if not segments:\n        return []\n    \n\n    segments_sec = [(s[\"start\"] / config.SR, s[\"end\"] / config.SR) for s in segments]\n\n    if len(segments_sec) > 1:\n        starts = np.array([s for s, _ in segments_sec])\n        ends = np.array([e for _, e in segments_sec])\n        \n        merged = []\n        current_start = starts[0]\n        current_end = ends[0]\n        \n        for i in range(1, len(starts)):\n            if starts[i] - current_end < config.MERGE_GAP_SEC:\n                current_end = max(current_end, ends[i])\n            else:\n                merged.append((current_start, current_end))\n                current_start = starts[i]\n                current_end = ends[i]\n        \n        merged.append((current_start, current_end))\n        return merged\n    \n    return segments_sec","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-10T12:06:33.988624Z","iopub.execute_input":"2026-02-10T12:06:33.988892Z","iopub.status.idle":"2026-02-10T12:06:34.000201Z","shell.execute_reply.started":"2026-02-10T12:06:33.988866Z","shell.execute_reply":"2026-02-10T12:06:33.999413Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Optimized Window Generation**","metadata":{}},{"cell_type":"code","source":"def generate_windows_optimized(segments, total_dur):\n\n\n    if total_dur < 30:\n        win_cfg = config.WIN_SHORT\n    elif total_dur < 120:\n        win_cfg = config.WIN_MEDIUM\n    else:\n        win_cfg = config.WIN_LONG\n    \n    win = win_cfg[\"win_sec\"]\n    hop = win_cfg[\"hop_sec\"]\n    \n    windows = []\n    for s, e in segments:\n        seg_dur = e - s\n        if seg_dur < 0.3:\n            continue\n        \n\n        num_windows = max(1, int((seg_dur - win) / hop) + 1)\n        \n        for i in range(num_windows):\n            start = s + i * hop\n            end = start + win\n            \n            if end > e:\n\n                end = e\n                start = max(s, e - win)\n            \n            if end - start >= 0.3:\n                windows.append((start, end))\n            \n            if end >= e:\n                break\n    \n\n    if not windows:\n        return []\n    \n    windows_arr = np.array(windows)\n\n    windows_rounded = np.round(windows_arr, 3)\n    _, unique_idx = np.unique(windows_rounded, axis=0, return_index=True)\n    \n    return windows_arr[np.sort(unique_idx)].tolist()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-10T12:06:34.001210Z","iopub.execute_input":"2026-02-10T12:06:34.001431Z","iopub.status.idle":"2026-02-10T12:06:34.017218Z","shell.execute_reply.started":"2026-02-10T12:06:34.001397Z","shell.execute_reply":"2026-02-10T12:06:34.016658Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Optimized Speaker Estimation**","metadata":{}},{"cell_type":"code","source":"def estimate_speakers(embeddings, max_k=None):\n\n    max_k = max_k or config.MAX_SPK\n    \n    n_samples = len(embeddings)\n    if n_samples < 3:\n        return 1\n    \n    max_k = min(max_k, n_samples - 1, 6)  \n    \n    best_k = 2\n    best_score = -1\n    \n\n    score_history = []\n    \n    for k in range(2, max_k + 1):\n        labels = AgglomerativeClustering(\n            n_clusters=k, \n            metric=\"cosine\", \n            linkage=\"average\"\n        ).fit_predict(embeddings)\n        \n        if len(set(labels)) < 2:\n            continue\n        \n        score = silhouette_score(embeddings, labels, metric=\"cosine\")\n        score_history.append(score)\n        \n        if score > best_score:\n            best_score = score\n            best_k = k\n\n        if len(score_history) >= 3:\n            if score_history[-1] < score_history[-2] < score_history[-3]:\n                break\n    \n\n    if best_score < 0.05:\n        return 2\n    \n    return best_k","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-10T12:06:34.018235Z","iopub.execute_input":"2026-02-10T12:06:34.018513Z","iopub.status.idle":"2026-02-10T12:06:34.032811Z","shell.execute_reply.started":"2026-02-10T12:06:34.018484Z","shell.execute_reply":"2026-02-10T12:06:34.032143Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Load Model**","metadata":{}},{"cell_type":"code","source":"def load_model():\n\n    model = EncoderClassifier.from_hparams(\n        source=\"speechbrain/spkrec-ecapa-voxceleb\",\n        run_opts={\"device\": config.DEVICE}\n    )\n    \n    if config.MODEL_PATH.exists():\n        print(f\"Loading fine-tuned model from {config.MODEL_PATH}\")\n        model.mods.embedding_model.load_state_dict(\n            torch.load(config.MODEL_PATH, map_location=config.DEVICE, weights_only=True)\n        )\n    else:\n        print(\"Using pretrained model (no fine-tuned weights found)\")\n    \n    model.eval()\n\n    for param in model.parameters():\n        param.requires_grad = False\n    \n    return model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-10T12:06:34.033638Z","iopub.execute_input":"2026-02-10T12:06:34.034008Z","iopub.status.idle":"2026-02-10T12:06:34.046434Z","shell.execute_reply.started":"2026-02-10T12:06:34.033977Z","shell.execute_reply":"2026-02-10T12:06:34.045759Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Optimized Diarization**","metadata":{}},{"cell_type":"code","source":"def diarize(wav_path, model):\n\n    wav = load_audio(wav_path)\n    total_dur = len(wav) / config.SR\n    \n\n    segments = get_vad_segments(wav)\n    if not segments:\n        return []\n    \n\n    windows = generate_windows_optimized(segments, total_dur)\n    if not windows:\n        return []\n\n    chunks = []\n    valid_windows = []\n    \n    for s, e in windows:\n        start_idx = int(s * config.SR)\n        end_idx = int(e * config.SR)\n        chunk = wav[start_idx:end_idx]\n        \n        if len(chunk) >= 0.3 * config.SR:\n            chunks.append(torch.tensor(chunk, dtype=torch.float32))\n            valid_windows.append((s, e))\n    \n    if not chunks:\n        return []\n    \n\n    emb_list = []\n    \n    with torch.no_grad():\n        for i in range(0, len(chunks), config.BATCH_SIZE):\n            batch = chunks[i:i + config.BATCH_SIZE]\n            \n            max_len = max(c.shape[0] for c in batch)\n            padded = torch.zeros(len(batch), max_len, dtype=torch.float32)\n            \n            for j, c in enumerate(batch):\n                padded[j, :len(c)] = c\n            \n\n            batch_tensor = padded.to(config.DEVICE)\n\n            if config.USE_AMP and config.DEVICE == \"cuda\":\n                with torch.cuda.amp.autocast():\n                    emb = model.encode_batch(batch_tensor).squeeze(1)\n            else:\n                emb = model.encode_batch(batch_tensor).squeeze(1)\n            \n            emb_list.append(emb.cpu())  \n\n            if config.DEVICE == \"cuda\" and i % (config.BATCH_SIZE * 4) == 0:\n                torch.cuda.empty_cache()\n    \n\n    embeddings = torch.cat(emb_list, dim=0)\n    embeddings = F.normalize(embeddings, dim=1)\n    embeddings_np = embeddings.numpy()\n    \n\n    del emb_list, chunks\n    gc.collect()\n    \n\n    n_spk = estimate_speakers(embeddings_np)\n    \n\n    labels = AgglomerativeClustering(\n        n_clusters=n_spk,\n        metric=\"cosine\",\n        linkage=\"average\"\n    ).fit_predict(embeddings_np)\n    \n\n    if config.MEDIAN_FILTER_SIZE > 1 and len(labels) >= config.MEDIAN_FILTER_SIZE:\n        kernel = config.MEDIAN_FILTER_SIZE\n        if kernel % 2 == 0:\n            kernel += 1\n        labels = medfilt(labels.astype(float), kernel_size=kernel).astype(int)\n    \n\n    items = sorted(zip(valid_windows, labels), key=lambda x: x[0][0])\n    \n    turns = []\n    if items:\n        cur_s, cur_e = items[0][0]\n        cur_lab = int(items[0][1])\n        \n        for (ws, we), lab in items[1:]:\n            lab = int(lab)\n            if lab == cur_lab and ws <= cur_e + config.MERGE_GAP_SEC:\n                cur_e = max(cur_e, we)\n            else:\n                if cur_e - cur_s >= config.MIN_TURN_SEC:\n                    turns.append((cur_s, cur_e, cur_lab))\n                cur_s, cur_e, cur_lab = ws, we, lab\n        \n        if cur_e - cur_s >= config.MIN_TURN_SEC:\n            turns.append((cur_s, cur_e, cur_lab))\n    \n    label_map = {}\n    spk_id = 1\n    output = []\n    \n    for s, e, lab in turns:\n        if lab not in label_map:\n            label_map[lab] = f\"SPEAKER_{spk_id}\"\n            spk_id += 1\n        \n        output.append({\n            \"start_time\": round(s, 2),\n            \"end_time\": round(e, 2),\n            \"speaker_id\": label_map[lab]\n        })\n    \n    return output","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-10T12:06:34.047238Z","iopub.execute_input":"2026-02-10T12:06:34.047462Z","iopub.status.idle":"2026-02-10T12:06:34.060789Z","shell.execute_reply.started":"2026-02-10T12:06:34.047433Z","shell.execute_reply":"2026-02-10T12:06:34.060081Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Optimized Inference with Parallel Processing**","metadata":{}},{"cell_type":"code","source":"def process_single_file(wav_file, model):\n\n    try:\n        diar = diarize(wav_file, model)\n        return {\n            \"filename\": wav_file.stem,\n            \"diarization\": json.dumps(diar, ensure_ascii=False)\n        }\n    except Exception as e:\n        print(f\"Error processing {wav_file.name}: {e}\")\n        return {\n            \"filename\": wav_file.stem,\n            \"diarization\": json.dumps([], ensure_ascii=False)\n        }\n\ndef run_inference(model):\n\n    print(\"\\n\" + \"=\"*50)\n    print(\"RUNNING INFERENCE\")\n    print(\"=\"*50)\n    \n    test_files = sorted(config.TEST_AUDIO.glob(\"*.wav\"))\n    print(f\"Found {len(test_files)} test files\")\n    \n    rows = []\n    \n\n    for wav_file in tqdm(test_files, desc=\"Processing audio files\"):\n        result = process_single_file(wav_file, model)\n        rows.append(result)\n\n        if len(rows) % 50 == 0:\n            gc.collect()\n            if config.DEVICE == \"cuda\":\n                torch.cuda.empty_cache()\n    \n\n    df = pd.DataFrame(rows)\n    df.to_csv(config.OUT_PATH, index=False)\n    \n    print(f\"\\nSaved to {config.OUT_PATH}\")\n    print(f\"Processed {len(df)} files\")\n    print(\"\\nFirst few rows:\")\n    print(df.head())\n    \n    return df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-10T12:06:34.061651Z","iopub.execute_input":"2026-02-10T12:06:34.061924Z","iopub.status.idle":"2026-02-10T12:06:34.077182Z","shell.execute_reply.started":"2026-02-10T12:06:34.061899Z","shell.execute_reply":"2026-02-10T12:06:34.076525Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Training**","metadata":{}},{"cell_type":"code","source":"def train():\n    print(\"Training is disabled in this optimized version.\")\n    print(\"Set config.USE_FINETUNING = True to enable.\")\n    pass","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-10T12:06:34.078103Z","iopub.execute_input":"2026-02-10T12:06:34.078510Z","iopub.status.idle":"2026-02-10T12:06:34.091940Z","shell.execute_reply.started":"2026-02-10T12:06:34.078472Z","shell.execute_reply":"2026-02-10T12:06:34.091265Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Main Execution**","metadata":{}},{"cell_type":"code","source":"\nprint(\"=\"*50)\nprint(\"OPTIMIZED BENGALI SPEAKER DIARIZATION\")\nprint(\"=\"*50)\nprint(f\"Device: {config.DEVICE}\")\nprint(f\"Batch size: {config.BATCH_SIZE}\")\nprint(f\"AMP enabled: {config.USE_AMP}\")\nprint(\"=\"*50)\n\n\nif config.USE_FINETUNING:\n    train()\n    gc.collect()\n    if config.DEVICE == \"cuda\":\n        torch.cuda.empty_cache()\n\n\nprint(\"\\nLoading model...\")\nmodel = load_model()\n\n\nsubmission = run_inference(model)\n\nprint(\"\\n\" + \"=\"*50)\nprint(\"PIPELINE COMPLETE!\")\nprint(\"=\"*50)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-10T12:06:34.092827Z","iopub.execute_input":"2026-02-10T12:06:34.093054Z","iopub.status.idle":"2026-02-10T12:21:00.255950Z","shell.execute_reply.started":"2026-02-10T12:06:34.093026Z","shell.execute_reply":"2026-02-10T12:21:00.255052Z"}},"outputs":[],"execution_count":null}]}