{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9f2c52f2",
   "metadata": {},
   "source": [
    "# VMC2026 Track 2 — exp15 PREDICT-ONLY (nạp checkpoint → chấm DEV, KHÔNG train) — Kaggle\n",
    "\n",
    "**Mục đích:** bạn ĐÃ có checkpoint exp15 (`ft_mamba_emotion_full*.pt`, lưu cả backbone WavLM + Mamba enc + heads).\n",
    "File này **chỉ inference**: dựng lại đúng kiến trúc → nạp trọng số + thống kê chuẩn hóa TỪ ckpt →\n",
    "dự đoán 5 cột cảm xúc trên tập DEV → ghép QMOS (exp07/UTMOSv2) → `answer.txt` → zip nộp.\n",
    "**KHÔNG** train, **KHÔNG** cần train.csv (chỉ cần wav DEV + metadata.csv để lấy cảm xúc target cho EMOS).\n",
    "\n",
    "## Vì sao nhanh\n",
    "- Không có vòng train → chỉ 1 lượt forward qua DEV (~2730 mẫu). Việc lâu nhất là trích audeering DEV\n",
    "  (~vài phút; có cache thì gần như tức thì).\n",
    "\n",
    "## Chuẩn bị input trên Kaggle (Add Input)\n",
    "1. Dataset Track 2 (wav + `metadata.csv` + `sets/dev.scp`).\n",
    "2. **Checkpoint** exp15: dataset chứa `ft_mamba_emotion_full*.pt` (vd `cache_exp8`). Auto-dò; hoặc trỏ `CKPT_PATH`.\n",
    "3. (tùy chọn) cache audeering `aud_dev.npz` để khỏi trích lại.\n",
    "4. (tùy chọn) `answer.txt` exp07 để mượn cột QMOS 0.548.\n",
    "\n",
    "**Cách chạy:** GPU **T4** + Internet **On** → Add Input → Run All."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adbc7c65",
   "metadata": {},
   "source": [
    "## 0. Cấu hình — SỬA Ở ĐÂY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7eb066d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, glob\n",
    "\n",
    "# ── TỰ DÒ DATA_ROOT (quét /kaggle/input tìm thư mục có sets + wav/ + metadata.csv) ──\n",
    "def find_data_root(search_root=\"/kaggle/input\"):\n",
    "    cands = []\n",
    "    for dev_scp in glob.glob(os.path.join(search_root, \"**\", \"sets\", \"dev.scp\"), recursive=True):\n",
    "        root = os.path.dirname(os.path.dirname(dev_scp))\n",
    "        score = os.path.isdir(os.path.join(root, \"wav\")) + os.path.exists(os.path.join(root, \"metadata.csv\"))\n",
    "        cands.append((score, root))\n",
    "    cands.sort(reverse=True)\n",
    "    return cands\n",
    "\n",
    "_cands = find_data_root(\"/kaggle/input\")\n",
    "if _cands:\n",
    "    print(\"🔎 Ứng viên DATA_ROOT:\")\n",
    "    for sc, r in _cands:\n",
    "        print(f\"   [{sc}/2] {r}\")\n",
    "    DATA_ROOT = _cands[0][1]\n",
    "    print(f\"👉 Tự chọn DATA_ROOT = {DATA_ROOT}\")\n",
    "else:\n",
    "    DATA_ROOT = \"/kaggle/input/datasets/minhtoan2\"   # dự phòng — sửa tay\n",
    "    print(f\"❌ Không thấy sets/dev.scp → dùng dự phòng {DATA_ROOT} (đã Add Input chưa?)\")\n",
    "\n",
    "WAV_DIR      = f\"{DATA_ROOT}/wav\"\n",
    "METADATA_CSV = f\"{DATA_ROOT}/metadata.csv\"     # wavID|emotion|transcript (KHÔNG header) — lấy cảm xúc target cho EMOS\n",
    "DEV_SCP      = f\"{DATA_ROOT}/sets/dev.scp\"\n",
    "\n",
    "OUT_DIR   = \"/kaggle/working\"\n",
    "CACHE_DIR = \"/kaggle/working/ft_cache\"\n",
    "os.makedirs(CACHE_DIR, exist_ok=True)\n",
    "\n",
    "# ── CHECKPOINT exp15 (đủ backbone + Mamba + heads) ───────────────────────────\n",
    "CKPT_PATH = \"\"    # << \"\" = auto-dò ft_mamba_emotion_full*.pt; hoặc \"/kaggle/input/<slug>/ft_mamba_emotion_full (2).pt\"\n",
    "\n",
    "def find_ckpt(explicit):\n",
    "    \"\"\"Tìm checkpoint exp15. Khớp cả tên bị thêm hậu tố trùng, vd 'ft_mamba_emotion_full (2).pt'.\"\"\"\n",
    "    if explicit and os.path.exists(explicit):\n",
    "        return explicit\n",
    "    for base in [\"/kaggle/input\", \"/kaggle/working\"]:\n",
    "        hits = sorted(glob.glob(os.path.join(base, \"**\", \"ft_mamba_emotion_full*.pt\"), recursive=True))\n",
    "        if hits:\n",
    "            return hits[0]\n",
    "    return \"\"\n",
    "\n",
    "CKPT_PATH = find_ckpt(CKPT_PATH)\n",
    "assert CKPT_PATH, \"❌ Không thấy checkpoint ft_mamba_emotion_full*.pt. Đã Add Input dataset chứa ckpt chưa?\"\n",
    "print(\"✅ Dùng checkpoint:\", CKPT_PATH)\n",
    "\n",
    "# (Tùy chọn) tái dùng cache audeering DEV — quét đệ quy (file có thể nằm trong archive/)\n",
    "CACHE_INPUT = \"/kaggle/input/cache-exp8\"   # << SỬA slug (hoặc \"\")\n",
    "if CACHE_INPUT and os.path.isdir(CACHE_INPUT):\n",
    "    import shutil\n",
    "    _n = 0\n",
    "    for _fp in glob.glob(os.path.join(CACHE_INPUT, \"**\", \"aud_*.npz\"), recursive=True):\n",
    "        shutil.copy(_fp, os.path.join(CACHE_DIR, os.path.basename(_fp))); _n += 1\n",
    "    print(f\"📦 Copy {_n} file aud_*.npz từ {CACHE_INPUT}\")\n",
    "\n",
    "# Mượn cột QMOS exp07 (0.548). Trỏ answer.txt exp07 nếu có; không thì UTMOSv2.\n",
    "EXP07_ANSWER = \"/kaggle/input/exp07-answer/answer.txt\"   # << (tùy chọn)\n",
    "\n",
    "# ── Siêu tham số PHẢI KHỚP lúc train exp15 (ckpt không lưu các số này của Mamba) ──\n",
    "MAMBA_DMODEL  = 256\n",
    "MAMBA_LAYERS  = 2\n",
    "MAMBA_DSTATE  = 16\n",
    "BIDIRECTIONAL = True\n",
    "TRUNK_HIDDEN  = 512\n",
    "HEAD_HIDDEN   = 128\n",
    "DROPOUT       = 0.3       # không ảnh hưởng eval (model.eval() tắt dropout) — chỉ để dựng đúng shape\n",
    "\n",
    "DEVICE       = \"cuda\"\n",
    "SR           = 16000\n",
    "MAX_SECONDS  = 6          # khớp lúc train (exp15 = 6)\n",
    "USE_AMP      = True\n",
    "LIMIT_DEV    = None       # << để None chấm ĐỦ 2730; đặt 20 để smoke-test nhanh\n",
    "\n",
    "EMOTIONS5 = [\"angry\", \"happy\", \"neutral\", \"sad\", \"surprised\"]\n",
    "_EMO_ALIAS = {\n",
    "    \"angry\": \"angry\", \"anger\": \"angry\",\n",
    "    \"happy\": \"happy\", \"happiness\": \"happy\", \"joy\": \"happy\",\n",
    "    \"neutral\": \"neutral\", \"calm\": \"neutral\",\n",
    "    \"sad\": \"sad\", \"sadness\": \"sad\",\n",
    "    \"surprise\": \"surprised\", \"surprised\": \"surprised\", \"surprising\": \"surprised\",\n",
    "}\n",
    "\n",
    "def norm_emotion(label):\n",
    "    key = str(label).strip().lower()\n",
    "    return _EMO_ALIAS.get(key, key if key in EMOTIONS5 else None)\n",
    "\n",
    "def stem(p):\n",
    "    return os.path.splitext(os.path.basename(str(p)))[0]\n",
    "\n",
    "print(\"DATA_ROOT:\", DATA_ROOT)\n",
    "for p in [WAV_DIR, METADATA_CSV, DEV_SCP, CKPT_PATH]:\n",
    "    print((\"  ✅ \" if os.path.exists(p) else \"  ❌ THIẾU \") + p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "febe8bdc",
   "metadata": {},
   "source": [
    "## 1. Cài đặt + tải code SAILER (để dựng đúng kiến trúc WavLM rồi nạp ckpt đè lên)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7732e245",
   "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(\"loralib\", \"speechbrain\", \"speechmos\", \"librosa\", \"soundfile\",\n",
    "            \"scipy\", \"scikit-learn\", \"pandas\", \"tqdm\")\n",
    "\n",
    "# Mamba kernel CUDA (tùy chọn — không có thì dùng Mamba thuần PyTorch, inference vẫn ổn vì chỉ 1 lượt forward)\n",
    "INSTALL_MAMBA_SSM = True\n",
    "if INSTALL_MAMBA_SSM:\n",
    "    try:\n",
    "        subprocess.run([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"ninja\"], check=True)\n",
    "        subprocess.run([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"--no-build-isolation\", \"causal-conv1d>=1.2.0\"], check=True)\n",
    "        subprocess.run([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"--no-build-isolation\", \"mamba-ssm\"], check=True)\n",
    "        print(\"✅ Cài mamba-ssm xong (dùng kernel CUDA nếu import được).\")\n",
    "    except Exception as e:\n",
    "        print(\"⚠️ Cài mamba-ssm thất bại:\", repr(e), \"→ Mamba thuần PyTorch (inference vẫn chạy).\")\n",
    "\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": "fba12581",
   "metadata": {},
   "source": [
    "## 2. Nạp checkpoint → dựng WavLM → load trọng số backbone đã fine-tune"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61199736",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\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 (chậm)\")\n",
    "\n",
    "ckpt = torch.load(CKPT_PATH, map_location=\"cpu\", weights_only=False)  # ckpt có numpy → cần False\n",
    "assert \"wavlm\" in ckpt, \"❌ Checkpoint KHÔNG có 'wavlm' (backbone) → không inference được. Cần ft_mamba_emotion_full*.pt đủ.\"\n",
    "print(\"✅ Nạp ckpt | keys:\", list(ckpt.keys()))\n",
    "\n",
    "# Lấy cấu hình KIẾN TRÚC từ ckpt (để dựng đúng shape head)\n",
    "USE_MAMBA  = bool(ckpt.get(\"USE_MAMBA\", True))\n",
    "Z_DIM      = int(ckpt.get(\"Z_DIM\", 256))\n",
    "AUD_DIM    = int(ckpt.get(\"AUD_DIM\", 0))\n",
    "USE_AUDEERING = AUD_DIM > 0\n",
    "UNFREEZE_TOP_LAYERS = int(ckpt.get(\"UNFREEZE_TOP_LAYERS\", 6))\n",
    "print(f\"Từ ckpt: USE_MAMBA={USE_MAMBA} · Z_DIM={Z_DIM} · AUD_DIM={AUD_DIM} (audeering={'ON' if USE_AUDEERING else 'OFF'})\")\n",
    "\n",
    "def find_hf_backbone(module):\n",
    "    cands = []\n",
    "    for name, m in module.named_modules():\n",
    "        enc = getattr(m, \"encoder\", None)\n",
    "        if getattr(m, \"feature_extractor\", None) is not None and enc is not None \\\n",
    "                and getattr(enc, \"layers\", None) is not None:\n",
    "            cands.append((name, m))\n",
    "    if not cands:\n",
    "        return None, None\n",
    "    cands.sort(key=lambda nm: sum(p.numel() for p in nm[1].parameters()), reverse=True)\n",
    "    return cands[0]\n",
    "\n",
    "wavlm = None\n",
    "try:\n",
    "    from src.model.emotion.wavlm_emotion import WavLMWrapper   # noqa: E402\n",
    "    _wrapper = WavLMWrapper.from_pretrained(\"tiantiaf/wavlm-large-categorical-emotion\")\n",
    "    name, wavlm = find_hf_backbone(_wrapper)\n",
    "    if wavlm is not None:\n",
    "        print(f\"✅ Dựng backbone WavLM từ SAILER wrapper tại '.{name}'\")\n",
    "except Exception as e:\n",
    "    print(\"⚠️ Lỗi nạp SAILER wrapper:\", repr(e), \"→ fallback WavLM trắng.\")\n",
    "\n",
    "if wavlm is None:\n",
    "    from transformers import WavLMModel\n",
    "    wavlm = WavLMModel.from_pretrained(\"microsoft/wavlm-large\")\n",
    "    print(\"ℹ️ Fallback: microsoft/wavlm-large.\")\n",
    "\n",
    "wavlm = wavlm.to(device)\n",
    "WAVLM_DIM = int(wavlm.config.hidden_size)\n",
    "wavlm.config.layerdrop = 0.0\n",
    "\n",
    "miss, unexp = wavlm.load_state_dict(ckpt[\"wavlm\"], strict=False)\n",
    "print(f\"🔁 load wavlm từ ckpt: thiếu {len(miss)} / dư {len(unexp)} key (kỳ vọng ~0)\")\n",
    "if len(miss) > 20 or len(unexp) > 20:\n",
    "    print(\"   ⚠️ Lệch key nhiều → kiểm tra backbone có khớp ckpt không.\")\n",
    "wavlm.eval()\n",
    "\n",
    "def frame_mask(T, attn_mask):\n",
    "    if attn_mask is None:\n",
    "        return torch.ones((1, T), dtype=torch.bool, device=device)\n",
    "    try:\n",
    "        return wavlm._get_feature_vector_attention_mask(T, attn_mask).bool()\n",
    "    except Exception:\n",
    "        return torch.ones((attn_mask.shape[0], T), dtype=torch.bool, device=attn_mask.device)\n",
    "\n",
    "def masked_mean(hidden, attn_mask):\n",
    "    if attn_mask is None:\n",
    "        return hidden.mean(dim=1)\n",
    "    fm = frame_mask(hidden.shape[1], attn_mask).unsqueeze(-1).to(hidden.dtype)\n",
    "    return (hidden * fm).sum(1) / fm.sum(1).clamp(min=1e-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "421e0b6a",
   "metadata": {},
   "source": [
    "## 3. audeering MSP-dim (FROZEN) — chỉ dựng nếu ckpt có dùng (AUD_DIM>0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d37d3d53",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import librosa\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "aud_backbone = aud_head = aud_proc = None\n",
    "if USE_AUDEERING:\n",
    "    from transformers import Wav2Vec2Model, Wav2Vec2Config, Wav2Vec2Processor\n",
    "    from huggingface_hub import hf_hub_download\n",
    "    AUD_NAME = \"audeering/wav2vec2-large-robust-12-ft-emotion-msp-dim\"\n",
    "    aud_proc = Wav2Vec2Processor.from_pretrained(AUD_NAME)\n",
    "    aud_cfg = Wav2Vec2Config.from_pretrained(AUD_NAME)\n",
    "    aud_backbone = Wav2Vec2Model(aud_cfg)\n",
    "    try:\n",
    "        _sd = __import__(\"safetensors.torch\", fromlist=[\"load_file\"]).load_file(\n",
    "            hf_hub_download(AUD_NAME, \"model.safetensors\"))\n",
    "    except Exception:\n",
    "        _sd = torch.load(hf_hub_download(AUD_NAME, \"pytorch_model.bin\"), map_location=\"cpu\")\n",
    "    bb_sd = {k[len(\"wav2vec2.\"):]: v for k, v in _sd.items() if k.startswith(\"wav2vec2.\")}\n",
    "    aud_backbone.load_state_dict(bb_sd, strict=False)\n",
    "    _hid = _sd[\"classifier.dense.weight\"].shape[0]\n",
    "    aud_head = nn.Sequential(nn.Linear(_hid, _hid), nn.Tanh(), nn.Linear(_hid, _sd[\"classifier.out_proj.weight\"].shape[0]))\n",
    "    aud_head[0].weight.data.copy_(_sd[\"classifier.dense.weight\"]); aud_head[0].bias.data.copy_(_sd[\"classifier.dense.bias\"])\n",
    "    aud_head[2].weight.data.copy_(_sd[\"classifier.out_proj.weight\"]); aud_head[2].bias.data.copy_(_sd[\"classifier.out_proj.bias\"])\n",
    "    aud_backbone = aud_backbone.to(device).eval()\n",
    "    aud_head = aud_head.to(device).eval()\n",
    "    assert _hid + 3 == AUD_DIM, f\"⚠️ AUD_DIM dựng ({_hid+3}) ≠ ckpt ({AUD_DIM}) → audeering không khớp!\"\n",
    "    print(f\"✅ audeering frozen ({AUD_DIM}-D)\")\n",
    "\n",
    "def load_wav(name_or_stem):\n",
    "    p = name_or_stem if os.path.isabs(str(name_or_stem)) else os.path.join(\n",
    "        WAV_DIR, name_or_stem if str(name_or_stem).endswith(\".wav\") else str(name_or_stem) + \".wav\")\n",
    "    if not os.path.exists(p):\n",
    "        return None\n",
    "    wave, _ = librosa.load(p, sr=SR, mono=True)\n",
    "    return wave[: MAX_SECONDS * SR].astype(np.float32)\n",
    "\n",
    "@torch.no_grad()\n",
    "def extract_audeering(stems, tag):\n",
    "    if not USE_AUDEERING:\n",
    "        return {}\n",
    "    cache_path = os.path.join(CACHE_DIR, f\"aud_{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\"[aud/{tag}] nạp cache: {len(store)}\")\n",
    "    todo = [s for s in stems if s not in store]\n",
    "    for i, s in enumerate(tqdm(todo, desc=f\"audeering {tag}\")):\n",
    "        wave = load_wav(s)\n",
    "        if wave is None:\n",
    "            continue\n",
    "        x = aud_proc(wave, sampling_rate=SR).input_values[0]\n",
    "        x = torch.from_numpy(np.asarray(x, dtype=np.float32)).unsqueeze(0).to(device)\n",
    "        h = aud_backbone(x)[0].mean(dim=1)\n",
    "        out = aud_head(h)[0].cpu().numpy()\n",
    "        vad = np.array([1 + 4 * out[2], 1 + 4 * out[0], 1 + 4 * out[1]], dtype=np.float32)  # [VAL,ARO,DOM]\n",
    "        store[s] = np.concatenate([h[0].cpu().numpy(), vad]).astype(np.float32)\n",
    "        if (i + 1) % 500 == 0:\n",
    "            np.savez(cache_path, **store)\n",
    "    if todo:\n",
    "        np.savez(cache_path, **store)\n",
    "    return store"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a04ef30",
   "metadata": {},
   "source": [
    "## 4. Cảm xúc target theo wavID (cho one-hot điều kiện của head EMOS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c092318",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def load_target_emotions():\n",
    "    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",
    "                tgt[stem(parts[0])] = norm_emotion(parts[1])\n",
    "    return tgt\n",
    "\n",
    "target_map = load_target_emotions()\n",
    "print(\"Target cảm xúc:\", len(target_map), \"wav\")\n",
    "\n",
    "def onehot_target(tgt):\n",
    "    v = np.zeros(len(EMOTIONS5), dtype=np.float32)\n",
    "    if tgt in EMOTIONS5:\n",
    "        v[EMOTIONS5.index(tgt)] = 1.0\n",
    "    return v"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0d7021a",
   "metadata": {},
   "source": [
    "## 5. Khối Mamba (giống exp15) + MambaEncoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8c31f88",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "try:\n",
    "    from mamba_ssm import Mamba as _OfficialMamba\n",
    "    _HAS_MAMBA_SSM = True\n",
    "    print(\"✅ Dùng mamba-ssm (CUDA kernel)\")\n",
    "except Exception:\n",
    "    _HAS_MAMBA_SSM = False\n",
    "    print(\"ℹ️ Không có mamba-ssm → Mamba thuần PyTorch\")\n",
    "\n",
    "class MambaBlockTorch(nn.Module):\n",
    "    def __init__(self, d_model, d_state=16, d_conv=4, expand=2):\n",
    "        super().__init__()\n",
    "        self.d_inner = expand * d_model\n",
    "        self.dt_rank = math.ceil(d_model / 16)\n",
    "        self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=False)\n",
    "        self.conv1d = nn.Conv1d(self.d_inner, self.d_inner, kernel_size=d_conv,\n",
    "                                groups=self.d_inner, padding=d_conv - 1, bias=True)\n",
    "        self.x_proj = nn.Linear(self.d_inner, self.dt_rank + d_state * 2, bias=False)\n",
    "        self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True)\n",
    "        A = torch.arange(1, d_state + 1, dtype=torch.float32).repeat(self.d_inner, 1)\n",
    "        self.A_log = nn.Parameter(torch.log(A))\n",
    "        self.D = nn.Parameter(torch.ones(self.d_inner))\n",
    "        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)\n",
    "        self.d_state = d_state\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, L, _ = x.shape\n",
    "        xin, z = self.in_proj(x).chunk(2, dim=-1)\n",
    "        xin = xin.transpose(1, 2)\n",
    "        xin = self.conv1d(xin)[..., :L].transpose(1, 2)\n",
    "        xin = F.silu(xin)\n",
    "        y = self._ssm(xin) * F.silu(z)\n",
    "        return self.out_proj(y)\n",
    "\n",
    "    def _ssm(self, x):\n",
    "        A = -torch.exp(self.A_log)\n",
    "        delta, Bm, Cm = torch.split(self.x_proj(x), [self.dt_rank, self.d_state, self.d_state], dim=-1)\n",
    "        delta = F.softplus(self.dt_proj(delta))\n",
    "        dA = torch.exp(delta.unsqueeze(-1) * A)\n",
    "        dB_x = delta.unsqueeze(-1) * Bm.unsqueeze(2) * x.unsqueeze(-1)\n",
    "        h = torch.zeros(x.shape[0], self.d_inner, self.d_state, device=x.device, dtype=x.dtype)\n",
    "        ys = []\n",
    "        for t in range(x.shape[1]):\n",
    "            h = dA[:, t] * h + dB_x[:, t]\n",
    "            ys.append((h * Cm[:, t].unsqueeze(1)).sum(-1))\n",
    "        return torch.stack(ys, dim=1) + x * self.D\n",
    "\n",
    "class MambaLayer(nn.Module):\n",
    "    def __init__(self, d_model, d_state):\n",
    "        super().__init__()\n",
    "        self.norm = nn.LayerNorm(d_model)\n",
    "        self.mix = _OfficialMamba(d_model=d_model, d_state=d_state, d_conv=4, expand=2) \\\n",
    "            if _HAS_MAMBA_SSM else MambaBlockTorch(d_model, d_state=d_state)\n",
    "    def forward(self, x):\n",
    "        return x + self.mix(self.norm(x))\n",
    "\n",
    "class MambaEncoder(nn.Module):\n",
    "    def __init__(self, d_in, d_model, n_layers, d_state, z_dim, bidir):\n",
    "        super().__init__()\n",
    "        self.bidir = bidir\n",
    "        self.proj = nn.Linear(d_in, d_model)\n",
    "        self.fwd = nn.ModuleList([MambaLayer(d_model, d_state) for _ in range(n_layers)])\n",
    "        if bidir:\n",
    "            self.bwd = nn.ModuleList([MambaLayer(d_model, d_state) for _ in range(n_layers)])\n",
    "        self.attn = nn.Linear(d_model, 1)\n",
    "        self.out = nn.Linear(d_model, z_dim)\n",
    "\n",
    "    @staticmethod\n",
    "    def _run(layers, h):\n",
    "        for L in layers:\n",
    "            h = L(h)\n",
    "        return h\n",
    "\n",
    "    def forward(self, x, mask):\n",
    "        with torch.cuda.amp.autocast(enabled=False):\n",
    "            x = x.float()\n",
    "            h = self.proj(x)\n",
    "            out = self._run(self.fwd, h)\n",
    "            if self.bidir:\n",
    "                out = out + torch.flip(self._run(self.bwd, torch.flip(h, dims=[1])), dims=[1])\n",
    "            a = self.attn(out).squeeze(-1).masked_fill(~mask, float(\"-inf\"))\n",
    "            w = torch.softmax(a, dim=1).unsqueeze(-1)\n",
    "            return self.out((out * w).sum(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8369a6b",
   "metadata": {},
   "source": [
    "## 6. Dựng enc + heads → nạp trọng số từ ckpt + lấy chuẩn hóa từ ckpt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c5e8556",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "N_EMO = len(EMOTIONS5)\n",
    "WAVLM_BRANCH = Z_DIM if USE_MAMBA else WAVLM_DIM\n",
    "TRUNK_IN = WAVLM_BRANCH + (AUD_DIM if USE_AUDEERING else 0)\n",
    "\n",
    "enc = MambaEncoder(WAVLM_DIM, MAMBA_DMODEL, MAMBA_LAYERS, MAMBA_DSTATE, Z_DIM, BIDIRECTIONAL).to(device) \\\n",
    "    if USE_MAMBA else None\n",
    "\n",
    "class EmoHeads(nn.Module):\n",
    "    def __init__(self, d_in, trunk_h, head_h, p, n_emo):\n",
    "        super().__init__()\n",
    "        self.trunk = nn.Sequential(nn.Linear(d_in, trunk_h), nn.ReLU(), nn.Dropout(p),\n",
    "                                   nn.Linear(trunk_h, trunk_h), nn.ReLU(), nn.Dropout(p))\n",
    "        self.emos = nn.Sequential(nn.Linear(trunk_h + n_emo, head_h), nn.ReLU(), nn.Dropout(p), nn.Linear(head_h, 1))\n",
    "        self.cat = nn.Sequential(nn.Linear(trunk_h, head_h), nn.ReLU(), nn.Dropout(p), nn.Linear(head_h, n_emo))\n",
    "        self.vad = nn.Sequential(nn.Linear(trunk_h, head_h), nn.ReLU(), nn.Dropout(p), nn.Linear(head_h, 3))\n",
    "    def forward(self, feat, tgt):\n",
    "        h = self.trunk(feat)\n",
    "        return self.emos(torch.cat([h, tgt], 1)), self.cat(h), self.vad(h)\n",
    "\n",
    "heads = EmoHeads(TRUNK_IN, TRUNK_HIDDEN, HEAD_HIDDEN, DROPOUT, N_EMO).to(device)\n",
    "hm, hu = heads.load_state_dict(ckpt[\"heads\"], strict=False)\n",
    "print(f\"🔁 load heads từ ckpt: thiếu {len(hm)} / dư {len(hu)} key (kỳ vọng 0)\")\n",
    "if USE_MAMBA:\n",
    "    assert ckpt.get(\"enc\") is not None, \"❌ ckpt USE_MAMBA=True nhưng KHÔNG có 'enc' → không inference đúng được.\"\n",
    "    em, eu = enc.load_state_dict(ckpt[\"enc\"], strict=False)\n",
    "    print(f\"🔁 load Mamba enc từ ckpt: thiếu {len(em)} / dư {len(eu)} key (kỳ vọng 0)\")\n",
    "heads.eval()\n",
    "if USE_MAMBA:\n",
    "    enc.eval()\n",
    "\n",
    "# Chuẩn hóa LẤY TỪ ckpt (head dự đoán ở thang z-score này → phải giải chuẩn đúng thang)\n",
    "emos_mu = float(ckpt[\"emos_mu\"]); emos_sd = float(ckpt[\"emos_sd\"])\n",
    "vad_mu = np.asarray(ckpt[\"vad_mu\"], dtype=np.float32); vad_sd = np.asarray(ckpt[\"vad_sd\"], dtype=np.float32)\n",
    "print(f\"Chuẩn hóa từ ckpt: emos μ={emos_mu:.3f} σ={emos_sd:.3f} | vad μ={np.round(vad_mu,2)}\")\n",
    "\n",
    "def wavlm_branch(input_values, attn_mask):\n",
    "    out = wavlm(input_values, attention_mask=attn_mask).last_hidden_state\n",
    "    if USE_MAMBA:\n",
    "        return enc(out, frame_mask(out.shape[1], attn_mask))\n",
    "    return masked_mean(out, attn_mask)\n",
    "\n",
    "print(f\"Trunk input = {TRUNK_IN} (wavlm-branch {WAVLM_BRANCH} [{'Mamba' if USE_MAMBA else 'mean-pool'}] + aud {AUD_DIM if USE_AUDEERING else 0})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdcf05c2",
   "metadata": {},
   "source": [
    "## 7. Dự đoán DEV → answer.txt (5 cột cảm xúc; QMOS mượn exp07/UTMOSv2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d225f54",
   "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",
    "aud_dev = extract_audeering(dev_stems, \"dev\")\n",
    "\n",
    "def load_exp07_qmos():\n",
    "    if EXP07_ANSWER and os.path.exists(EXP07_ANSWER):\n",
    "        import csv\n",
    "        d = {}\n",
    "        with open(EXP07_ANSWER) as f:\n",
    "            for row in csv.DictReader(f):\n",
    "                d[row[\"wav\"]] = float(row[\"QMOS\"]); d[stem(row[\"wav\"])] = float(row[\"QMOS\"])\n",
    "        print(f\"✅ Mượn QMOS exp07 ({EXP07_ANSWER}): {len(d)//2} wav\")\n",
    "        return d\n",
    "    return None\n",
    "\n",
    "qmos_map = load_exp07_qmos()\n",
    "if qmos_map is None:\n",
    "    print(\"ℹ️ Không có answer.txt exp07 → chấm QMOS bằng UTMOSv2 (T05, vô địch VMC2024).\")\n",
    "    pip_install(\"git+https://github.com/sarulab-speech/UTMOSv2.git\")\n",
    "    import utmosv2\n",
    "    v2 = utmosv2.create_model(pretrained=True)\n",
    "    qmos_map = {}\n",
    "    for n in tqdm(dev_names, desc=\"UTMOSv2\"):\n",
    "        wav = os.path.join(WAV_DIR, n if str(n).endswith(\".wav\") else str(n) + \".wav\")\n",
    "        if not os.path.exists(wav):\n",
    "            continue\n",
    "        out = v2.predict(input_path=wav)\n",
    "        qmos_map[n] = float(out[\"predicted_mos\"]) if isinstance(out, dict) else float(out)\n",
    "    del v2; torch.cuda.empty_cache() if device == \"cuda\" else None\n",
    "\n",
    "@torch.no_grad()\n",
    "def predict_emotion(sid):\n",
    "    wave = load_wav(sid)\n",
    "    if wave is None or (USE_AUDEERING and sid not in aud_dev):\n",
    "        return None\n",
    "    iv = torch.from_numpy(wave).unsqueeze(0).to(device)\n",
    "    am = torch.ones((1, len(wave)), dtype=torch.long, device=device)\n",
    "    tgt = torch.from_numpy(onehot_target(target_map.get(sid))).unsqueeze(0).to(device)\n",
    "    with torch.cuda.amp.autocast(enabled=USE_AMP and device == \"cuda\"):\n",
    "        fw = wavlm_branch(iv, am)\n",
    "        feat = torch.cat([fw, torch.from_numpy(aud_dev[sid]).unsqueeze(0).to(device)], dim=1) if USE_AUDEERING else fw\n",
    "        emos_p, cat_l, vad_p = heads(feat, tgt)\n",
    "    emos = float(emos_p.item()) * emos_sd + emos_mu\n",
    "    cat5 = F.softmax(cat_l, 1)[0].float().cpu().numpy()\n",
    "    vad3 = vad_p[0].float().cpu().numpy() * vad_sd + vad_mu\n",
    "    return emos, cat5, vad3\n",
    "\n",
    "def fmt_cat(p5):\n",
    "    return \"|\".join(f\"{e}:{p5[i]:.6g}\" for i, e in enumerate(EMOTIONS5))\n",
    "\n",
    "def build_answer(out_path):\n",
    "    n_real = n_def = 0\n",
    "    with open(out_path, \"w\") as f:\n",
    "        f.write(\"wav,QMOS,EMOS,CAT,VAL,ARO,DOM\\n\")\n",
    "        for name in tqdm(dev_names, desc=\"answer\"):\n",
    "            sid = stem(name)\n",
    "            pr = predict_emotion(sid)\n",
    "            if pr is None:\n",
    "                emos, cat5, vad3 = 3.0, np.full(5, 0.2, np.float32), np.array([3.0, 3.0, 3.0]); n_def += 1\n",
    "            else:\n",
    "                emos, cat5, vad3 = pr; n_real += 1\n",
    "            qmos = qmos_map.get(name, qmos_map.get(sid, 3.0))\n",
    "            f.write(f\"{name},{qmos:.6g},{emos:.6g},{fmt_cat(cat5)},{vad3[0]:.6g},{vad3[1]:.6g},{vad3[2]:.6g}\\n\")\n",
    "    print(f\"Ghi {len(dev_names)} dòng → {out_path} | cảm xúc thật {n_real}, mặc định {n_def}\")\n",
    "\n",
    "answer_path = os.path.join(OUT_DIR, \"answer.txt\")\n",
    "build_answer(answer_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42503595",
   "metadata": {},
   "source": [
    "## 8. Validate + đóng zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42dec31f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def validate(path):\n",
    "    import csv\n",
    "    with open(path) as f:\n",
    "        rows = list(csv.reader(f))\n",
    "    assert rows[0][0] == \"wav\" and \"QMOS\" in rows[0] and \"EMOS\" in rows[0], \"Header sai\"\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",
    "\n",
    "validate(answer_path)\n",
    "os.system(f\"cd {OUT_DIR} && zip -j submission_track2_exp15_predict.zip answer.txt \"\n",
    "          f\"&& unzip -l submission_track2_exp15_predict.zip\")\n",
    "print(\"Sẵn sàng nộp:\", os.path.join(OUT_DIR, \"submission_track2_exp15_predict.zip\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbef2a21",
   "metadata": {},
   "source": [
    "## Ghi chú\n",
    "- File này **chỉ inference** — không train, không cần train.csv. Dùng khi đã có `ft_mamba_emotion_full*.pt`.\n",
    "- ⚠️ **Siêu tham số Mamba/heads (MAMBA_DMODEL/LAYERS/DSTATE, TRUNK_HIDDEN, HEAD_HIDDEN) PHẢI khớp lúc train**\n",
    "  (ckpt không lưu các số này) — nếu lúc train exp15 bạn đổi, hãy sửa cho khớp ở cell 0, nếu không load_state_dict\n",
    "  sẽ lệch key / sai shape.\n",
    "- `USE_MAMBA`, `Z_DIM`, `AUD_DIM`, `UNFREEZE_TOP_LAYERS` thì **đọc tự động từ ckpt**.\n",
    "- QMOS: tốt nhất Add Input `answer.txt` exp07 (0.548); không có thì tự chấm UTMOSv2.\n",
    "- Smoke-test: đặt `LIMIT_DEV=20` chạy thử cho nhanh, OK rồi đặt lại `None` để chấm đủ 2730."
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}