{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# VMC2026 Track 2 — Baseline Pipeline (Kaggle)\n\nQMOS (SpeechMOS) + EmoCat (emotion2vec) + **EMOS (emotion2vec target-prob, mặc định offline)** → gộp `answer.txt`.\n\n**Trước khi chạy:** Accelerator = **GPU T4**, Internet = **On**.\n- **+ Add Input** → tab Datasets → dataset Track 2 đã upload (Kaggle tự giải nén → có thư mục `vmc2026-track2/`).\n- Với mặc định `EMOS_METHOD='emotion2vec'`: **KHÔNG cần** `GEMINI_API_KEY`. Chỉ cần Secrets khi đổi sang `'gemini'` (để có thêm VAD).\n\nChạy được ngay: **QMOS + EmoCat + EMOS** (chỉ cần wav + `metadata.csv` chứa cảm xúc target).\n\n> ⚠️ Train phase: dự đoán tập **DEV** (`sets/dev.scp`, ~2730 mẫu). Thư mục `wav/` có cả train+dev nên KHÔNG glob hết — chỉ lấy đúng dev.scp."
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Config — SỬA Ở ĐÂY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "import os, glob\n\n# ── Data Track 2 trên Kaggle (dataset đã upload, KHÔNG có thư mục con lồng) ──\nDATA_ROOT    = '/kaggle/input/vmc2026-track2-full'   # << slug dataset bạn upload\nWAV_DIR      = f'{DATA_ROOT}/wav'\nMETADATA_CSV = f'{DATA_ROOT}/metadata.csv'      # wavID|emotion|transcript (KHÔNG header)\nDEV_SCP      = f'{DATA_ROOT}/sets/dev.scp'      # danh sách wav tập DEV (tập cần nộp ở train phase)\n\n# Test nhanh trên ESD: trỏ WAV_DIR vào ESD, đặt DEV_SCP=None và METADATA_CSV=None.\n# WAV_DIR = '/kaggle/input/datasets/nguyenthanhlim/emotional-speech-dataset-esd/Emotion Speech Dataset'\n# DEV_SCP = None; METADATA_CSV = None\n\nLIMIT = 20   # << 20 = chạy THỬ nhanh. Đổi None để chạy TOÀN BỘ DEV rồi nộp.\n\n# ── Cách tính EMOS ──────────────────────────────────────────────────────────\n# 'emotion2vec': OFFLINE, MIỄN PHÍ (exp01, khuyến nghị) — P(cảm xúc target) từ emotion2vec → scale 1–5.\n# 'gemini'     : LLM-as-judge qua Gemini API (cần GEMINI_API_KEY, tốn phí). Chỉ cách này có VAD.\nEMOS_METHOD = 'emotion2vec'\n\nOUT_DIR = '/kaggle/working'\nRUN_QMOS, RUN_EMOCAT = True, True\n_have_meta = bool(METADATA_CSV) and os.path.exists(METADATA_CSV)\nRUN_EMOS = _have_meta                                # cả 2 cách đều cần target từ metadata\nRUN_VAD  = _have_meta and EMOS_METHOD == 'gemini'    # VAD chỉ có ở Gemini\nEMOTIONS5 = ['angry', 'happy', 'neutral', 'sad', 'surprised']\n\n# Chuẩn hóa nhãn cảm xúc target (metadata) → đúng 1 trong 5 lớp của emotion2vec.\n_EMO_ALIAS = {'angry':'angry','anger':'angry','happy':'happy','happiness':'happy','joy':'happy',\n              'neutral':'neutral','calm':'neutral','sad':'sad','sadness':'sad',\n              'surprise':'surprised','surprised':'surprised','surprising':'surprised'}\ndef norm_emotion(label):\n    key = str(label).strip().lower()\n    return _EMO_ALIAS.get(key, key if key in EMOTIONS5 else None)\n\ndef list_wavs(d):\n    # Có DEV_SCP → đọc danh sách tên file tập DEV (wav nằm phẳng trong wav/).\n    # Không có   → quét đệ quy mọi .wav (chế độ test ESD, lồng speaker/emotion).\n    if DEV_SCP and os.path.exists(DEV_SCP):\n        with open(DEV_SCP) as f:\n            names = [ln.strip() for ln in f if ln.strip()]\n        wavs = [os.path.join(d, n) for n in names]\n    else:\n        wavs = sorted(glob.glob(os.path.join(d, '**', '*.wav'), recursive=True))\n    return wavs[:LIMIT] if LIMIT else wavs\n\nprint('WAV_DIR:', WAV_DIR, '| EMOS_METHOD:', EMOS_METHOD)\nprint('Số wav:', len(list_wavs(WAV_DIR)) if os.path.isdir(WAV_DIR) else '(chưa thấy thư mục)')\nprint('Chế độ DEV (dev.scp):', bool(DEV_SCP and os.path.exists(DEV_SCP)))\nif METADATA_CSV and os.path.exists(METADATA_CSV) and DEV_SCP and os.path.exists(DEV_SCP):\n    n_meta = sum(1 for _ in open(METADATA_CSV))\n    n_dev  = sum(1 for _ in open(DEV_SCP))\n    print(f'metadata.csv: {n_meta} dòng | dev.scp: {n_dev} dòng')"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Cài đặt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q speechmos funasr librosa soundfile pandas google-genai loguru tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. QMOS — SpeechMOS (UTMOS, không cần fairseq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "def run_qmos(wav_dir):\n    import torch, librosa\n    dev = 'cuda' if torch.cuda.is_available() else 'cpu'\n    predictor = torch.hub.load('tarepan/SpeechMOS:v1.2.0', 'utmos22_strong', trust_repo=True).to(dev)  # << GPU\n    print('QMOS device:', dev)\n    scores, missing = {}, 0\n    for w in list_wavs(wav_dir):              # w là đường dẫn đầy đủ\n        if not os.path.exists(w):             # mẫu ESD/DailyTalk chưa lấy ngoài → bỏ qua, không crash\n            missing += 1; continue\n        wave, _ = librosa.load(w, sr=16000, mono=True)\n        wave_t = torch.from_numpy(wave).unsqueeze(0).to(dev)   # << đưa input lên GPU\n        scores[w] = float(predictor(wave_t, sr=16000).mean().item())\n    if missing: print(f'[QMOS] Bỏ qua {missing} file thiếu (chưa có ESD/DailyTalk) → điểm mặc định.')\n    return scores\n\nqmos_scores = run_qmos(WAV_DIR) if RUN_QMOS else {}\nprint('QMOS xong:', len(qmos_scores))\nlist(qmos_scores.items())[:3]"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. EmoCat — emotion2vec+ large\n",
    "Đã sửa bug bản gốc + lọc 5 lớp + chuẩn hóa tổng = 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "def run_emocat(wav_dir):\n    import torch\n    from funasr import AutoModel\n    dev = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n    model = AutoModel(model='iic/emotion2vec_plus_large', hub='hf', device=dev)  # << chạy GPU\n    print('EmoCat device:', dev)\n    results, missing = {}, 0\n    for w in list_wavs(wav_dir):              # w là đường dẫn đầy đủ\n        if not os.path.exists(w):             # mẫu ESD/DailyTalk chưa lấy ngoài → bỏ qua\n            missing += 1; continue\n        rec = model.generate(w, granularity='utterance', extract_embedding=False)\n        probs = {e: 0.0 for e in EMOTIONS5}\n        for lab, sc in zip(rec[0]['labels'], rec[0]['scores']):\n            name = lab.split('/')[-1]\n            if name in probs:\n                probs[name] = float(sc)\n        total = sum(probs.values())\n        if total > 0:\n            probs = {k: v / total for k, v in probs.items()}\n        results[w] = probs\n    if missing: print(f'[EmoCat] Bỏ qua {missing} file thiếu (chưa có ESD/DailyTalk) → phân bố mặc định.')\n    return results\n\nemocat_probs = run_emocat(WAV_DIR) if RUN_EMOCAT else {}\nprint('EmoCat xong:', len(emocat_probs))\nlist(emocat_probs.items())[:2]"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 4. EMOS — emotion2vec target-prob (mặc định) hoặc Gemini\n**emotion2vec (exp01, offline):** lấy P(cảm xúc target) từ emotion2vec (đã tính ở cell EmoCat), scale [0,1]→[1,5]. Chấm đủ 2.730 mẫu, KHÔNG cần API. SRCC chỉ quan tâm thứ hạng nên scale tuyến tính không đổi tương quan.\n\n**Gemini (`EMOS_METHOD='gemini'`):** LLM-as-judge, cần `GEMINI_API_KEY` + credit; tự lọc metadata về DEV để đỡ tốn. Chỉ cách này có VAD."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "emos_scores, vad_scores = {}, {}   # key = TÊN FILE wav (uttID, có .wav)\n\n# Đọc cảm xúc target từ metadata.csv → {stem: emotion_chuẩn}\ndef load_target_emotions():\n    tgt = {}\n    if not (METADATA_CSV and os.path.exists(METADATA_CSV)):\n        return tgt\n    with open(METADATA_CSV, encoding='utf-8') as f:\n        for ln in f:\n            parts = ln.strip().split('|')\n            if len(parts) >= 2:\n                stem = os.path.splitext(os.path.basename(parts[0]))[0]\n                tgt[stem] = norm_emotion(parts[1])\n    return tgt\n\ntarget_map = load_target_emotions()\nprint('Nhãn cảm xúc target đọc được:', len(target_map))\n\nif RUN_EMOS and EMOS_METHOD == 'emotion2vec':\n    # ── EMOS OFFLINE: P(cảm xúc target) từ emotion2vec (cell EmoCat), scale [0,1]→[1,5] ──\n    assert RUN_EMOCAT and emocat_probs, 'EMOS theo emotion2vec cần chạy cell EmoCat (mục 3) trước.'\n    miss_t = miss_p = 0\n    for w in list_wavs(WAV_DIR):\n        name = os.path.basename(w)\n        tgt = target_map.get(os.path.splitext(name)[0])\n        probs = emocat_probs.get(w)\n        if tgt is None:\n            miss_t += 1; continue\n        if not probs:\n            miss_p += 1; continue\n        emos_scores[name] = 1.0 + 4.0 * probs.get(tgt, 0.0)   # p=0→1 điểm, p=1→5 điểm\n    if miss_t: print(f'[EMOS-e2v] {miss_t} mẫu thiếu nhãn target → mặc định 3.')\n    if miss_p: print(f'[EMOS-e2v] {miss_p} mẫu thiếu prob emotion2vec → mặc định 3.')\n    print(f'✅ EMOS (emotion2vec) cho {len(emos_scores)} mẫu — không cần API.')\n\nelif RUN_EMOS or RUN_VAD:   # EMOS_METHOD == 'gemini'\n    try:\n        from kaggle_secrets import UserSecretsClient\n        os.environ['GEMINI_API_KEY'] = UserSecretsClient().get_secret('GEMINI_API_KEY')\n        print('Đã nạp GEMINI_API_KEY từ Secrets')\n    except Exception as e:\n        print('Chưa nạp được key:', e)\n\n    # ── Lọc metadata.csv → CHỈ giữ mẫu thuộc DEV (tránh trả tiền Gemini cho mẫu train) ──\n    dev_stems = {os.path.splitext(n.strip())[0] for n in open(DEV_SCP) if n.strip()}\n    META_DEV = '/kaggle/working/metadata_dev.csv'\n    kept = 0\n    with open(METADATA_CSV) as fin, open(META_DEV, 'w') as fout:\n        for line in fin:\n            if not line.strip():\n                continue\n            stem = os.path.splitext(os.path.basename(line.split('|')[0].strip()))[0]\n            if stem in dev_stems:\n                fout.write(line); kept += 1\n    print(f'metadata_dev.csv: {kept} dòng (kỳ vọng ~{len(dev_stems)})')\n\n    GEMINI_ROWS = f'--end-row {LIMIT}' if LIMIT else ''\n    !git clone -q https://github.com/voicemos-challenge/vmc2026-baselines.git /kaggle/working/vmc2026-baselines\n    !cd /kaggle/working/vmc2026-baselines/track2/EMOS && python Gemini_EMOS.py --metadata-path $META_DEV --base-path $WAV_DIR --output-file /kaggle/working/emos.csv --workers 4 --resume $GEMINI_ROWS\n    !cd /kaggle/working/vmc2026-baselines/track2/VAD && python Gemini_VAD.py --metadata-path $META_DEV --base-path $WAV_DIR --output-file /kaggle/working/vad.csv --workers 4 --resume $GEMINI_ROWS\n\n    import pandas as pd\n    if os.path.exists('/kaggle/working/emos.csv'):\n        d = pd.read_csv('/kaggle/working/emos.csv'); emos_scores = dict(zip(d['uttID'], d['emos']))\n    if os.path.exists('/kaggle/working/vad.csv'):\n        d = pd.read_csv('/kaggle/working/vad.csv')   # cột chuẩn: uttID, val, aro, dom\n        for _, r in d.iterrows():\n            vad_scores[r['uttID']] = (r['val'], r['aro'], r['dom'])\n\n    if emos_scores:\n        dev_bases = {os.path.basename(w) for w in list_wavs(WAV_DIR)}\n        if not (set(emos_scores) & dev_bases):\n            print('⚠️ KEY LỆCH: uttID không khớp tên file dev → EMOS/VAD sẽ về mặc định!')\n        else:\n            print('✅ Key khớp — EMOS/VAD sẽ gộp đúng.')\n\nprint('EMOS:', len(emos_scores), '| VAD:', len(vad_scores))"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Gộp answer.txt (tự bỏ cột thiếu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "def fmt_cat(p):\n    return '|'.join(f'{e}:{p[e]:.6g}' for e in EMOTIONS5)\n\ndef build_answer(out_path):\n    wavs = list_wavs(WAV_DIR)\n    have_cat = RUN_EMOCAT and len(emocat_probs) > 0\n    have_vad = RUN_VAD and len(vad_scores) > 0\n    cols = ['wav', 'QMOS', 'EMOS']\n    if have_cat: cols.append('CAT')\n    if have_vad: cols += ['VAL', 'ARO', 'DOM']\n    with open(out_path, 'w') as f:\n        f.write(','.join(cols) + '\\n')\n        for w in wavs:\n            name = os.path.basename(w)            # tên file = cột wav & key của emos/vad\n            row = [name, f\"{qmos_scores.get(w, 3.0):.6g}\", str(emos_scores.get(name, 3))]\n            if have_cat: row.append(fmt_cat(emocat_probs.get(w, {e: 0.2 for e in EMOTIONS5})))\n            if have_vad:\n                v = vad_scores.get(name, (3, 3, 3)); row += [str(v[0]), str(v[1]), str(v[2])]\n            f.write(','.join(row) + '\\n')\n    print(f'Ghi {len(wavs)} dòng → {out_path} | cột: {cols}')\n\nanswer_path = os.path.join(OUT_DIR, 'answer.txt')\nbuild_answer(answer_path)\n!head -3 {answer_path}"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Validate + zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import csv\n",
    "with open(answer_path) as f:\n",
    "    rows = list(csv.reader(f))\n",
    "header = rows[0]\n",
    "assert header[0] == 'wav' and 'QMOS' in header and 'EMOS' in header, 'Header sai'\n",
    "for i, r in enumerate(rows[1:], 2):\n",
    "    assert len(r) == len(header), f'Dòng {i} sai số cột'\n",
    "print(f'OK: {len(rows)-1} dòng, header = {header}')\n",
    "!cd /kaggle/working && zip -j submission_track2.zip answer.txt && unzip -l submission_track2.zip"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}