{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e2d94d72",
   "metadata": {},
   "source": [
    "# VMC2026 Track 2 — exp06 (TRAIN QMOS head) — Kaggle\n",
    "\n",
    "**Mục tiêu:** QMOS là cột **duy nhất chưa train** (đang dùng UTMOS zero-shot → SRCC kẹt 0.414).\n",
    "`train.csv` CÓ sẵn cột `qMOS` → ta train 1 **head hồi quy nhỏ** trên đặc trưng SSL (đã cache ở exp04)\n",
    "để vượt 0.414.\n",
    "\n",
    "## Ý tưởng (đọc 1 lần cho hiểu)\n",
    "- Tái dùng đặc trưng **emotion2vec + SAILER** đã trích & cache trong `fusion_cache/` (exp04) → KHÔNG trích lại.\n",
    "- Thêm **chính điểm UTMOS** (SpeechMOS) làm 1 đặc trưng đầu vào → head chỉ cần **học chỉnh sửa (residual)**\n",
    "  quanh 0.414 thay vì học lại từ đầu → an toàn, gần như chắc chắn ≥ UTMOS đơn lẻ.\n",
    "- Nhãn vàng QMOS = **TB `qMOS` theo wav** (gộp các listener trong `train.csv`).\n",
    "- Có **val nội bộ 10%** → đo SRCC, so thẳng với UTMOS trên CÙNG tập val → biết có cải thiện thật\n",
    "  **trước khi** tốn lượt nộp CodaBench.\n",
    "- Cuối cùng: **GIỮ NGUYÊN exp04** (5 cột cảm xúc đang thắng), chỉ **thay cột QMOS** trong `answer.txt`.\n",
    "\n",
    "```\n",
    " mỗi wav ─► [e2v_emb | e2v_probs5 | sailer_emb | sailer_probs9 | sailer_vad3 | UTMOS] ─► MLP ─► QMOS\n",
    "                                                                                 (head train)\n",
    "```\n",
    "\n",
    "**Cách chạy trên Kaggle:** Settings → Accelerator = **GPU T4**, Internet = **On**.\n",
    "+ Add Input: (1) dataset Track 2 (15.477 wav, có `sets/train.csv`) ; (2) — nếu có — dataset chứa\n",
    "`fusion_cache/*.npz` đã Save Version ở exp04 (đỡ ~15') ; (3) file `answer.txt` của exp04 để ghép cột.\n",
    "Lần đầu đặt `LIMIT_TRAIN=300`, `LIMIT_DEV=20` để bắt lỗi setup, OK rồi đặt `None`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b42d5d49",
   "metadata": {},
   "source": [
    "## 0. Cấu hình — SỬA Ở ĐÂY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93e29194",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# ── Data Track 2 ─────────────────────────────────────────────────────────────\n",
    "DATA_ROOT    = \"/kaggle/input/vmc2026-track2-full/vmc2026-track2\"   # << SỬA slug cho khớp Add Input\n",
    "WAV_DIR      = f\"{DATA_ROOT}/wav\"\n",
    "TRAIN_CSV    = f\"{DATA_ROOT}/sets/train.csv\"   # nhãn người nghe: lisID|wavID|qMOS|emoCat|eMOS|val|dom|aro\n",
    "DEV_SCP      = f\"{DATA_ROOT}/sets/dev.scp\"     # danh sách wav tập DEV\n",
    "\n",
    "OUT_DIR   = \"/kaggle/working\"\n",
    "# Dùng CHUNG cache với exp04. Nếu đã Save Version cache ở exp04, trỏ CACHE_DIR vào dataset đó\n",
    "# (vd \"/kaggle/input/<slug-cache>/fusion_cache\") để khỏi trích lại; nếu không, để mặc định sẽ tự trích.\n",
    "CACHE_DIR = \"/kaggle/working/fusion_cache\"\n",
    "os.makedirs(CACHE_DIR, exist_ok=True)\n",
    "\n",
    "# File answer.txt của exp04 (5 cột cảm xúc đang thắng) để GHÉP cột QMOS mới vào.\n",
    "# Trỏ tới nơi bạn đặt file exp04. Nếu không có, notebook vẫn xuất qmos_dev.csv riêng + cảnh báo.\n",
    "EXP04_ANSWER = \"/kaggle/input/exp04-answer/answer.txt\"   # << SỬA; hoặc \"/kaggle/working/answer.txt\"\n",
    "\n",
    "# ── Đặc trưng dùng cho QMOS ──────────────────────────────────────────────────\n",
    "USE_E2V        = True     # nối embedding emotion2vec\n",
    "USE_SAILER     = True     # nối embedding SAILER/WavLM\n",
    "USE_CLASSPROB  = True     # nối thêm xác suất lớp (e2v5 + sailer9 + vad3)\n",
    "USE_UTMOS_FEAT = True     # nối thêm điểm UTMOS làm 1 đặc trưng (neo residual quanh 0.414)\n",
    "\n",
    "# ── Siêu tham số train head ──────────────────────────────────────────────────\n",
    "DEVICE      = \"cuda\"\n",
    "HIDDEN      = 256\n",
    "DROPOUT     = 0.3\n",
    "LR          = 1e-3\n",
    "EPOCHS      = 120\n",
    "BATCH       = 64\n",
    "VAL_FRAC    = 0.10\n",
    "PATIENCE    = 20\n",
    "SEED        = 42\n",
    "RANK_LAMBDA = 0.0         # 0 = chỉ MSE. >0 (vd 0.2) = cộng thêm pairwise ranking loss (tối ưu thứ hạng=SRCC)\n",
    "\n",
    "LIMIT_TRAIN = None        # số nhỏ (vd 300) để chạy thử; None = full\n",
    "LIMIT_DEV   = None\n",
    "\n",
    "def stem(p):\n",
    "    return os.path.splitext(os.path.basename(str(p)))[0]\n",
    "\n",
    "assert USE_E2V or USE_SAILER or USE_UTMOS_FEAT, \"Phải bật ít nhất 1 nguồn đặc trưng.\"\n",
    "print(\"DATA_ROOT:\", DATA_ROOT)\n",
    "for p in [WAV_DIR, TRAIN_CSV, DEV_SCP]:\n",
    "    print((\"  ✅ \" if os.path.exists(p) else \"  ❌ THIẾU \") + p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47ac221d",
   "metadata": {},
   "source": [
    "## 1. Cài đặt + (nếu cần) tải code SAILER\n",
    "emotion2vec qua `funasr`; SAILER cần `WavLMWrapper` trong repo `vox-profile-release` (clone + sys.path).\n",
    "Nếu cache đã đủ thì các model này sẽ KHÔNG được nạp (chỉ nạp khi còn file phải trích)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99ba1947",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, subprocess\n",
    "\n",
    "def pip_install(*pkgs):\n",
    "    subprocess.run([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", *pkgs], check=True)\n",
    "\n",
    "pip_install(\"speechmos\", \"funasr\", \"librosa\", \"soundfile\", \"pandas\", \"scipy\", \"scikit-learn\", \"tqdm\")\n",
    "\n",
    "if USE_SAILER:\n",
    "    pip_install(\"loralib\", \"speechbrain\")\n",
    "    REPO_DIR = \"/kaggle/working/vox-profile-release\"\n",
    "    if not os.path.exists(REPO_DIR):\n",
    "        subprocess.run([\"git\", \"clone\", \"--depth\", \"1\",\n",
    "                        \"https://github.com/tiantiaf0627/vox-profile-release.git\", REPO_DIR], check=True)\n",
    "    if REPO_DIR not in sys.path:\n",
    "        sys.path.insert(0, REPO_DIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac9dcefc",
   "metadata": {},
   "source": [
    "## 2. Nhãn vàng QMOS (gộp `qMOS` theo wavID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db4a41a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "def load_qmos_labels():\n",
    "    \"\"\"train.csv (sep='|') → DataFrame [wavID, qmos] với qmos = TB theo wav.\"\"\"\n",
    "    df = pd.read_csv(TRAIN_CSV, sep=\"|\")\n",
    "    cols = {c.lower().strip(): c for c in df.columns}\n",
    "    wav_col  = cols.get(\"wavid\") or cols.get(\"wav\") or list(df.columns)[1]\n",
    "    qmos_col = cols.get(\"qmos\")  or cols.get(\"qMOS\".lower()) or cols.get(\"mos\")\n",
    "    assert qmos_col, f\"Không thấy cột qMOS trong train.csv (cột: {list(df.columns)})\"\n",
    "    df[\"_stem\"] = df[wav_col].map(stem)\n",
    "    g = df.groupby(\"_stem\")[qmos_col].mean().reset_index()\n",
    "    g.columns = [\"wavID\", \"qmos\"]\n",
    "    return g\n",
    "\n",
    "qmos_df = load_qmos_labels()\n",
    "print(f\"wav train (gộp): {len(qmos_df)}\")\n",
    "print(\"qMOS:\", qmos_df[\"qmos\"].describe()[[\"mean\", \"std\", \"min\", \"max\"]].to_dict())\n",
    "qmos_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfd7df0c",
   "metadata": {},
   "source": [
    "## 3. Trích / nạp đặc trưng (cache CHUNG với exp04) + điểm UTMOS\n",
    "- `extract_e2v` / `extract_sailer`: y hệt exp04, cache `e2v_<tag>.npz` / `sailer_<tag>.npz`.\n",
    "- `extract_utmos`: chấm UTMOS từng wav → cache `utmos_<tag>.npz` (dùng vừa làm đặc trưng, vừa làm baseline so sánh)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec1e63a1",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "device = DEVICE if torch.cuda.is_available() else \"cpu\"\n",
    "print(\"Device:\", device, (\"✅ \" + torch.cuda.get_device_name(0)) if device == \"cuda\" else \"⚠️ CPU\")\n",
    "\n",
    "EMOTIONS5 = [\"angry\", \"happy\", \"neutral\", \"sad\", \"surprised\"]\n",
    "\n",
    "def extract_e2v(stems, tag):\n",
    "    \"\"\"→ dict {stem: emb_full[D1+5]}. Cache CACHE_DIR/e2v_<tag>.npz (giống exp04).\"\"\"\n",
    "    from tqdm.auto import tqdm\n",
    "    cache_path = os.path.join(CACHE_DIR, f\"e2v_{tag}.npz\")\n",
    "    store = {}\n",
    "    if os.path.exists(cache_path):\n",
    "        z = np.load(cache_path, allow_pickle=True)\n",
    "        store = {k: z[k] for k in z.files}\n",
    "        print(f\"[e2v/{tag}] nạp cache: {len(store)}\")\n",
    "    todo = [s for s in stems if s not in store]\n",
    "    if todo:\n",
    "        from funasr import AutoModel\n",
    "        m = AutoModel(model=\"iic/emotion2vec_plus_large\", hub=\"hf\", device=device)\n",
    "        for i, s in enumerate(tqdm(todo, desc=f\"e2v {tag}\")):\n",
    "            wav = os.path.join(WAV_DIR, s + \".wav\")\n",
    "            if not os.path.exists(wav):\n",
    "                continue\n",
    "            r = m.generate(wav, granularity=\"utterance\", extract_embedding=True)[0]\n",
    "            emb = np.asarray(r[\"feats\"], dtype=np.float32).reshape(-1)\n",
    "            probs = {e: 0.0 for e in EMOTIONS5}\n",
    "            for lab, sc in zip(r[\"labels\"], r[\"scores\"]):\n",
    "                name = lab.split(\"/\")[-1]\n",
    "                if name in probs:\n",
    "                    probs[name] = float(sc)\n",
    "            tot = sum(probs.values())\n",
    "            p5 = np.array([probs[e] / tot if tot > 0 else 0.2 for e in EMOTIONS5], dtype=np.float32)\n",
    "            store[s] = np.concatenate([emb, p5]).astype(np.float32)\n",
    "            if (i + 1) % 500 == 0:\n",
    "                np.savez(cache_path, **store)\n",
    "        np.savez(cache_path, **store)\n",
    "        del m\n",
    "        torch.cuda.empty_cache() if device == \"cuda\" else None\n",
    "    return store   # mỗi value = [D1 | 5]\n",
    "\n",
    "def _pool_feat(features):\n",
    "    f = features.detach().cpu().numpy()\n",
    "    if f.ndim <= 1:\n",
    "        return f.reshape(-1).astype(np.float32)\n",
    "    return f.mean(axis=tuple(range(f.ndim - 1))).reshape(-1).astype(np.float32)\n",
    "\n",
    "def extract_sailer(stems, tag):\n",
    "    \"\"\"→ dict {stem: vec[D2+9+3]}. Cache CACHE_DIR/sailer_<tag>.npz (giống exp04).\"\"\"\n",
    "    import librosa\n",
    "    from tqdm.auto import tqdm\n",
    "    cache_path = os.path.join(CACHE_DIR, f\"sailer_{tag}.npz\")\n",
    "    store = {}\n",
    "    if os.path.exists(cache_path):\n",
    "        z = np.load(cache_path, allow_pickle=True)\n",
    "        store = {k: z[k] for k in z.files}\n",
    "        print(f\"[sailer/{tag}] nạp cache: {len(store)}\")\n",
    "    todo = [s for s in stems if s not in store]\n",
    "    if todo:\n",
    "        from src.model.emotion.wavlm_emotion import WavLMWrapper\n",
    "        sailer = WavLMWrapper.from_pretrained(\"tiantiaf/wavlm-large-categorical-emotion\").to(device).eval()\n",
    "        with torch.no_grad():\n",
    "            for i, s in enumerate(tqdm(todo, desc=f\"sailer {tag}\")):\n",
    "                wav = os.path.join(WAV_DIR, s + \".wav\")\n",
    "                if not os.path.exists(wav):\n",
    "                    continue\n",
    "                wave, _ = librosa.load(wav, sr=16000, mono=True)\n",
    "                wave = wave[: 15 * 16000]\n",
    "                data = torch.from_numpy(wave).float().unsqueeze(0).to(device)\n",
    "                logits, feat, _det, arousal, valence, dominance = sailer(data, return_feature=True)\n",
    "                emb = _pool_feat(feat)\n",
    "                p9 = F.softmax(logits, dim=1)[0].detach().cpu().numpy().astype(np.float32)\n",
    "                vad3 = np.array([1 + 4 * float(valence.item()),\n",
    "                                 1 + 4 * float(arousal.item()),\n",
    "                                 1 + 4 * float(dominance.item())], dtype=np.float32)\n",
    "                store[s] = np.concatenate([emb, p9, vad3]).astype(np.float32)\n",
    "                if (i + 1) % 500 == 0:\n",
    "                    np.savez(cache_path, **store)\n",
    "        np.savez(cache_path, **store)\n",
    "        del sailer\n",
    "        torch.cuda.empty_cache() if device == \"cuda\" else None\n",
    "    return store   # mỗi value = [D2 | 9 | 3]\n",
    "\n",
    "def extract_utmos(names, tag):\n",
    "    \"\"\"Chấm UTMOS từng wav (theo TÊN file, vì DEV gọi .wav theo tên). → dict {stem: score}.\n",
    "    Cache CACHE_DIR/utmos_<tag>.npz. Dùng vừa làm đặc trưng vừa làm baseline so sánh.\"\"\"\n",
    "    import librosa\n",
    "    from tqdm.auto import tqdm\n",
    "    cache_path = os.path.join(CACHE_DIR, f\"utmos_{tag}.npz\")\n",
    "    store = {}\n",
    "    if os.path.exists(cache_path):\n",
    "        z = np.load(cache_path, allow_pickle=True)\n",
    "        store = {k: float(z[k]) for k in z.files}\n",
    "        print(f\"[utmos/{tag}] nạp cache: {len(store)}\")\n",
    "    todo = [n for n in names if stem(n) not in store]\n",
    "    if todo:\n",
    "        predictor = torch.hub.load(\"tarepan/SpeechMOS:v1.2.0\", \"utmos22_strong\",\n",
    "                                   trust_repo=True).to(device).eval()\n",
    "        with torch.no_grad():\n",
    "            for i, n in enumerate(tqdm(todo, desc=f\"utmos {tag}\")):\n",
    "                wav = os.path.join(WAV_DIR, n if n.endswith(\".wav\") else n + \".wav\")\n",
    "                if not os.path.exists(wav):\n",
    "                    continue\n",
    "                wave, _ = librosa.load(wav, sr=16000, mono=True)\n",
    "                sc = float(predictor(torch.from_numpy(wave).unsqueeze(0).to(device), sr=16000).mean().item())\n",
    "                store[stem(n)] = sc\n",
    "                if (i + 1) % 500 == 0:\n",
    "                    np.savez(cache_path, **{k: np.float32(v) for k, v in store.items()})\n",
    "        np.savez(cache_path, **{k: np.float32(v) for k, v in store.items()})\n",
    "        del predictor\n",
    "        torch.cuda.empty_cache() if device == \"cuda\" else None\n",
    "    return store"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aed7338b",
   "metadata": {},
   "source": [
    "## 4. Dựng feature + nhãn cho train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c09bb508",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_stems = list(qmos_df[\"wavID\"])\n",
    "if LIMIT_TRAIN:\n",
    "    train_stems = train_stems[:LIMIT_TRAIN]\n",
    "\n",
    "e2v_tr    = extract_e2v(train_stems, \"train\")    if USE_E2V    else {}\n",
    "sailer_tr = extract_sailer(train_stems, \"train\") if USE_SAILER else {}\n",
    "utmos_tr  = extract_utmos(train_stems, \"train\")  if USE_UTMOS_FEAT else {}\n",
    "\n",
    "def qmos_feature(sid, e2v_map, sailer_map, utmos_map):\n",
    "    \"\"\"Nối đặc trưng QMOS cho 1 wav. None nếu thiếu phần bắt buộc.\"\"\"\n",
    "    parts = []\n",
    "    if USE_E2V:\n",
    "        v = e2v_map.get(sid)\n",
    "        if v is None:\n",
    "            return None\n",
    "        parts.append(v[:-5])                      # emb e2v\n",
    "        if USE_CLASSPROB:\n",
    "            parts.append(v[-5:])                  # probs5\n",
    "    if USE_SAILER:\n",
    "        v = sailer_map.get(sid)\n",
    "        if v is None:\n",
    "            return None\n",
    "        parts.append(v[:-12])                     # emb sailer\n",
    "        if USE_CLASSPROB:\n",
    "            parts.append(v[-12:])                 # probs9 + vad3\n",
    "    if USE_UTMOS_FEAT:\n",
    "        u = utmos_map.get(sid)\n",
    "        if u is None:\n",
    "            return None\n",
    "        parts.append(np.array([u], dtype=np.float32))\n",
    "    return np.concatenate(parts).astype(np.float32)\n",
    "\n",
    "lab = qmos_df.set_index(\"wavID\")[\"qmos\"]\n",
    "X, y = [], []\n",
    "for s in train_stems:\n",
    "    f = qmos_feature(s, e2v_tr, sailer_tr, utmos_tr)\n",
    "    if f is None or s not in lab.index:\n",
    "        continue\n",
    "    X.append(f)\n",
    "    y.append(float(lab.loc[s]))\n",
    "\n",
    "X = np.stack(X).astype(np.float32)\n",
    "y = np.array(y, dtype=np.float32)\n",
    "FEAT_DIM = X.shape[1]\n",
    "print(f\"Train: X={X.shape} y={y.shape}\")\n",
    "\n",
    "feat_mean = X.mean(0, keepdims=True)\n",
    "feat_std  = X.std(0, keepdims=True) + 1e-6\n",
    "Xn = (X - feat_mean) / feat_std\n",
    "y_mu, y_sd = float(y.mean()), float(y.std() + 1e-6)\n",
    "yn = (y - y_mu) / y_sd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82cc65f8",
   "metadata": {},
   "source": [
    "## 5. Train head QMOS + so với UTMOS trên CÙNG val nội bộ\n",
    "- Head = MLP nhỏ (`Linear→ReLU→Dropout ×2 → 1`). Loss = MSE (+ tùy chọn pairwise ranking).\n",
    "- In **SRCC head** và **SRCC UTMOS** trên cùng tập val → biết head có thật sự vượt 0.414 không."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "324ab564",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "from scipy.stats import spearmanr\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "torch.manual_seed(SEED); np.random.seed(SEED)\n",
    "idx_all = np.arange(X.shape[0])\n",
    "tr_idx, va_idx = train_test_split(idx_all, test_size=VAL_FRAC, random_state=SEED)\n",
    "\n",
    "def to_t(a):\n",
    "    return torch.tensor(a, dtype=torch.float32, device=device)\n",
    "\n",
    "Xn_t = to_t(Xn); yn_t = to_t(yn).unsqueeze(1)\n",
    "\n",
    "class QMOSHead(nn.Module):\n",
    "    def __init__(self, d_in, h, p):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(d_in, h), nn.ReLU(), nn.Dropout(p),\n",
    "            nn.Linear(h, h), nn.ReLU(), nn.Dropout(p),\n",
    "            nn.Linear(h, 1),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "model = QMOSHead(FEAT_DIM, HIDDEN, DROPOUT).to(device)\n",
    "opt = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=1e-5)\n",
    "mse = nn.MSELoss()\n",
    "\n",
    "def pairwise_rank_loss(pred, target):\n",
    "    \"\"\"Khuyến khích pred xếp hạng giống target (margin ranking trên các cặp trong batch).\"\"\"\n",
    "    n = pred.shape[0]\n",
    "    if n < 2:\n",
    "        return torch.zeros((), device=device)\n",
    "    pi, pj = pred.unsqueeze(0), pred.unsqueeze(1)\n",
    "    ti, tj = target.unsqueeze(0), target.unsqueeze(1)\n",
    "    sign = torch.sign(ti - tj)                       # +1 nếu i nên cao hơn j\n",
    "    diff = pi - pj\n",
    "    # hinge: phạt khi thứ tự sai\n",
    "    return torch.relu(-sign * diff).mean()\n",
    "\n",
    "@torch.no_grad()\n",
    "def eval_val():\n",
    "    model.eval()\n",
    "    p = model(Xn_t[va_idx]).cpu().numpy().ravel()\n",
    "    srcc_head = spearmanr(p, y[va_idx]).correlation\n",
    "    out = {\"head\": float(srcc_head)}\n",
    "    if USE_UTMOS_FEAT:\n",
    "        u = X[va_idx, -1]                            # cột UTMOS (đặc trưng cuối, chưa chuẩn hóa)\n",
    "        out[\"utmos\"] = float(spearmanr(u, y[va_idx]).correlation)\n",
    "    return out\n",
    "\n",
    "best, best_state, bad = -1e9, None, 0\n",
    "tr_t = torch.tensor(tr_idx, device=device)\n",
    "for ep in range(1, EPOCHS + 1):\n",
    "    model.train()\n",
    "    perm = tr_t[torch.randperm(len(tr_t), device=device)]\n",
    "    run = 0.0\n",
    "    for i in range(0, len(perm), BATCH):\n",
    "        b = perm[i:i + BATCH]\n",
    "        opt.zero_grad()\n",
    "        pred = model(Xn_t[b])\n",
    "        loss = mse(pred, yn_t[b])\n",
    "        if RANK_LAMBDA > 0:\n",
    "            loss = loss + RANK_LAMBDA * pairwise_rank_loss(pred.ravel(), yn_t[b].ravel())\n",
    "        loss.backward(); opt.step()\n",
    "        run += loss.item() * len(b)\n",
    "    m = eval_val()\n",
    "    if m[\"head\"] > best:\n",
    "        best = m[\"head\"]\n",
    "        best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}\n",
    "        bad = 0\n",
    "    else:\n",
    "        bad += 1\n",
    "    if ep % 5 == 0 or ep == 1:\n",
    "        extra = f\" | UTMOS={m['utmos']:.4f}\" if \"utmos\" in m else \"\"\n",
    "        print(f\"epoch {ep:3d} | loss {run/len(perm):.4f} | head SRCC={m['head']:.4f}{extra} | best {best:.4f}\")\n",
    "    if bad >= PATIENCE:\n",
    "        print(f\"Early stop ở epoch {ep}.\")\n",
    "        break\n",
    "\n",
    "model.load_state_dict(best_state)\n",
    "final = eval_val()\n",
    "print(\"\\n✅ VAL (nội bộ):\")\n",
    "print(f\"   QMOS head SRCC = {final['head']:.4f}\")\n",
    "if \"utmos\" in final:\n",
    "    print(f\"   UTMOS  baseline = {final['utmos']:.4f}  (mốc leaderboard 0.414)\")\n",
    "    print(\"   →\", \"✅ HEAD VƯỢT UTMOS\" if final[\"head\"] > final[\"utmos\"] else \"⚠️ chưa vượt — thử tăng EPOCHS / RANK_LAMBDA / bật thêm đặc trưng\")\n",
    "\n",
    "torch.save({\"state\": best_state, \"feat_mean\": feat_mean, \"feat_std\": feat_std,\n",
    "            \"y_mu\": y_mu, \"y_sd\": y_sd, \"FEAT_DIM\": FEAT_DIM,\n",
    "            \"USE_E2V\": USE_E2V, \"USE_SAILER\": USE_SAILER,\n",
    "            \"USE_CLASSPROB\": USE_CLASSPROB, \"USE_UTMOS_FEAT\": USE_UTMOS_FEAT,\n",
    "            \"val_srcc\": best}, os.path.join(OUT_DIR, \"qmos_head.pt\"))\n",
    "print(\"Đã lưu\", os.path.join(OUT_DIR, \"qmos_head.pt\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d33a7aca",
   "metadata": {},
   "source": [
    "## 6. Dự đoán QMOS cho DEV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69efbd00",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def list_dev():\n",
    "    with open(DEV_SCP) as f:\n",
    "        return [ln.strip() for ln in f if ln.strip()]\n",
    "\n",
    "dev_names = list_dev()\n",
    "if LIMIT_DEV:\n",
    "    dev_names = dev_names[:LIMIT_DEV]\n",
    "dev_stems = [stem(n) for n in dev_names]\n",
    "print(\"DEV:\", len(dev_names), \"mẫu\")\n",
    "\n",
    "e2v_dev    = extract_e2v(dev_stems, \"dev\")    if USE_E2V    else {}\n",
    "sailer_dev = extract_sailer(dev_stems, \"dev\") if USE_SAILER else {}\n",
    "utmos_dev  = extract_utmos(dev_names, \"dev\")  if USE_UTMOS_FEAT else {}\n",
    "\n",
    "@torch.no_grad()\n",
    "def predict_qmos(sid):\n",
    "    f = qmos_feature(sid, e2v_dev, sailer_dev, utmos_dev)\n",
    "    if f is None:\n",
    "        return None\n",
    "    fn = (f[None, :] - feat_mean) / feat_std\n",
    "    model.eval()\n",
    "    return float(model(to_t(fn)).item()) * y_sd + y_mu     # đảo z-score\n",
    "\n",
    "qmos_pred = {}\n",
    "n_real = n_def = 0\n",
    "for n in dev_names:\n",
    "    sid = stem(n)\n",
    "    p = predict_qmos(sid)\n",
    "    if p is None:\n",
    "        p = utmos_dev.get(sid, 3.0)                          # rơi về UTMOS nếu thiếu feature\n",
    "        n_def += 1\n",
    "    else:\n",
    "        n_real += 1\n",
    "    qmos_pred[n] = p\n",
    "print(f\"QMOS dự đoán: head thật {n_real}, dự phòng UTMOS {n_def}\")\n",
    "\n",
    "# Lưu riêng (để ghép tay nếu cần)\n",
    "import csv\n",
    "qmos_csv = os.path.join(OUT_DIR, \"qmos_dev.csv\")\n",
    "with open(qmos_csv, \"w\", newline=\"\") as f:\n",
    "    w = csv.writer(f); w.writerow([\"wav\", \"QMOS\"])\n",
    "    for n in dev_names:\n",
    "        w.writerow([n, f\"{qmos_pred[n]:.6g}\"])\n",
    "print(\"Đã ghi\", qmos_csv)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3e47def",
   "metadata": {},
   "source": [
    "## 7. Ghép QMOS mới vào answer.txt của exp04 → bản nộp mới\n",
    "Giữ NGUYÊN 5 cột cảm xúc đang thắng (EMOS/CAT/VAL/ARO/DOM), chỉ thay cột QMOS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3b94589",
   "metadata": {},
   "outputs": [],
   "source": [
    "def merge_into_exp04(exp04_path, out_path):\n",
    "    if not os.path.exists(exp04_path):\n",
    "        print(f\"⚠️ Không thấy {exp04_path} → BỎ QUA ghép. Hãy dùng qmos_dev.csv để thay cột QMOS thủ công,\")\n",
    "        print(\"   hoặc trỏ EXP04_ANSWER đúng đường dẫn answer.txt của exp04 rồi chạy lại cell này.\")\n",
    "        return False\n",
    "    with open(exp04_path) as f:\n",
    "        rows = list(csv.reader(f))\n",
    "    header = rows[0]\n",
    "    qi = header.index(\"QMOS\")\n",
    "    wi = header.index(\"wav\")\n",
    "    n_swapped = n_miss = 0\n",
    "    with open(out_path, \"w\", newline=\"\") as f:\n",
    "        w = csv.writer(f); w.writerow(header)\n",
    "        for r in rows[1:]:\n",
    "            name = r[wi]\n",
    "            if name in qmos_pred:\n",
    "                r[qi] = f\"{qmos_pred[name]:.6g}\"; n_swapped += 1\n",
    "            else:\n",
    "                n_miss += 1\n",
    "            w.writerow(r)\n",
    "    print(f\"Ghép xong → {out_path} | thay {n_swapped} cột QMOS, thiếu {n_miss} (giữ QMOS cũ)\")\n",
    "    return True\n",
    "\n",
    "merged = os.path.join(OUT_DIR, \"answer.txt\")\n",
    "ok = merge_into_exp04(EXP04_ANSWER, merged)\n",
    "\n",
    "if ok:\n",
    "    # validate + zip\n",
    "    with open(merged) as f:\n",
    "        rows = list(csv.reader(f))\n",
    "    assert rows[0][0] == \"wav\" and \"QMOS\" in rows[0]\n",
    "    for i, r in enumerate(rows[1:], 2):\n",
    "        assert len(r) == len(rows[0]), f\"Dòng {i} sai số cột\"\n",
    "    print(f\"OK: {len(rows)-1} dòng, header = {rows[0]}\")\n",
    "    os.system(f\"cd {OUT_DIR} && zip -j submission_track2_exp06_qmos.zip answer.txt \"\n",
    "              f\"&& unzip -l submission_track2_exp06_qmos.zip\")\n",
    "    print(\"Sẵn sàng nộp:\", os.path.join(OUT_DIR, \"submission_track2_exp06_qmos.zip\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a517b97",
   "metadata": {},
   "source": [
    "## Ghi chú\n",
    "- **Lần đầu** đặt `LIMIT_TRAIN=300`, `LIMIT_DEV=20` để bắt lỗi; OK rồi đặt `None`.\n",
    "- **So sánh công bằng**: mục 5 in cả `head SRCC` và `UTMOS SRCC` trên CÙNG val nội bộ → chỉ nộp khi head > UTMOS.\n",
    "- Nếu head **chưa vượt** 0.414: thử (a) tăng `EPOCHS`; (b) bật `RANK_LAMBDA=0.2` (tối ưu thứ hạng);\n",
    "  (c) đảm bảo `USE_UTMOS_FEAT=True` (neo residual); (d) thử bỏ bớt đặc trưng nhiễu (tắt `USE_CLASSPROB`).\n",
    "- **Ablation QMOS cho paper**: bật/tắt `USE_E2V/USE_SAILER/USE_UTMOS_FEAT/USE_CLASSPROB` → ghi `docs/04_experiments_log.md` (exp06).\n",
    "- Cache dùng CHUNG `fusion_cache/` với exp04 → nhớ **Save Version** giữ lại (gồm `utmos_*.npz` mới).\n",
    "- Ghi config → kết quả → nhận xét vào `docs/04_experiments_log.md` (mục exp06)."
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}