Upload finetune_wavlm.py with huggingface_hub

finetune_wavlm.py

@@ -0,0 +1,748 @@
+"""
+Fine-tune WavLM-Large Backbone for Phoneme Scoring
+===================================================
+Same as finetune_backbone.py but uses WavLM-Large instead of wav2vec2-large.
+WavLM has denoising pre-training, making it more robust for non-standard speech (e.g. children).
+"""
+
+import json
+import re
+import warnings
+from pathlib import Path
+
+import numpy as np
+import openpyxl
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+import torchaudio.functional as F_audio
+from g2p_en import G2p
+from huggingface_hub import hf_hub_download
+from sklearn.metrics import precision_recall_fscore_support, roc_auc_score
+from sklearn.model_selection import train_test_split
+from scipy import stats
+from torch.utils.data import DataLoader, Dataset
+from transformers import Wav2Vec2FeatureExtractor, Wav2Vec2Model, Wav2Vec2ForCTC, WavLMModel
+
+warnings.filterwarnings("ignore")
+
+SAMPLE_RATE = 16000
+
+ARPABET_TO_IPA = {
+    "aa": ["ɑː", "ɑ", "ɒ", "a"], "ae": ["æ"], "ah": ["ʌ", "ə", "ɐ"],
+    "ao": ["ɔː", "ɔ", "ɒ"], "aw": ["aʊ"], "ax": ["ə", "ɐ", "ʌ"], "ay": ["aɪ"],
+    "b": ["b"], "ch": ["tʃ"], "d": ["d"], "dh": ["ð"],
+    "eh": ["ɛ", "e"], "er": ["ɜː", "ɝ", "ɚ", "ɜ"], "ey": ["eɪ"],
+    "f": ["f"], "g": ["ɡ", "g"], "hh": ["h"],
+    "ih": ["ɪ", "ᵻ"], "iy": ["iː", "i"],
+    "ir": ["ɪɹ"], "jh": ["dʒ"], "k": ["k"], "l": ["l"],
+    "m": ["m"], "n": ["n"], "ng": ["ŋ"],
+    "ow": ["oʊ", "o", "əʊ"], "oy": ["ɔɪ"],
+    "p": ["p"], "r": ["ɹ", "r"], "s": ["s"], "sh": ["ʃ"],
+    "t": ["t"], "th": ["θ"], "uh": ["ʊ"], "uw": ["uː", "u"],
+    "ur": ["ʊɹ"], "v": ["v"], "w": ["w"], "y": ["j"],
+    "z": ["z"], "zh": ["ʒ"],
+    "ar": ["ɑːɹ"], "oo": ["ʊ", "uː"], "dr": ["dɹ"], "tr": ["tɹ"],
+    "ts": ["ts"], "dz": ["dz"],
+}
+
+ALL_PHONES = sorted(ARPABET_TO_IPA.keys())
+PHONE_TO_ID = {ph: i for i, ph in enumerate(ALL_PHONES)}
+N_PHONE_TYPES = len(ALL_PHONES)
+
+
+# ============================================================
+# Model: Backbone + Scoring Head (end-to-end)
+# ============================================================
+class BackbonePhoneScorer(nn.Module):
+    """
+    End-to-end model: WavLM-Large backbone → per-phoneme pooling → MLP scorer
+    """
+    def __init__(self, n_phone_types=N_PHONE_TYPES, phone_emb_dim=32,
+                 hidden_dim=1024, mlp_dim=512, unfreeze_top_n=6):
+        super().__init__()
+
+        # Load WavLM-Large backbone (disable time masking for fine-tuning)
+        self.backbone = WavLMModel.from_pretrained(
+            "microsoft/wavlm-large",
+            output_hidden_states=False,
+            mask_time_prob=0.0,
+        )
+        # Freeze all layers first
+        for param in self.backbone.parameters():
+            param.requires_grad = False
+
+        # Unfreeze top N transformer layers
+        n_layers = len(self.backbone.encoder.layers)
+        for i in range(n_layers - unfreeze_top_n, n_layers):
+            for param in self.backbone.encoder.layers[i].parameters():
+                param.requires_grad = True
+
+        # Also unfreeze layer norm
+        if hasattr(self.backbone.encoder, 'layer_norm'):
+            for param in self.backbone.encoder.layer_norm.parameters():
+                param.requires_grad = True
+
+        self.fe_backbone = Wav2Vec2FeatureExtractor.from_pretrained("microsoft/wavlm-large")
+
+        # Phone embedding + MLP scorer (same as PhoneScorerV2)
+        self.phone_emb = nn.Embedding(n_phone_types, phone_emb_dim)
+        input_dim = hidden_dim + phone_emb_dim + 2  # +2 for GOP and n_frames
+
+        self.shared = nn.Sequential(
+            nn.Linear(input_dim, mlp_dim),
+            nn.BatchNorm1d(mlp_dim),
+            nn.GELU(),
+            nn.Dropout(0.3),
+            nn.Linear(mlp_dim, mlp_dim),
+            nn.BatchNorm1d(mlp_dim),
+            nn.GELU(),
+            nn.Dropout(0.3),
+            nn.Linear(mlp_dim, 256),
+            nn.BatchNorm1d(256),
+            nn.GELU(),
+            nn.Dropout(0.2),
+        )
+        self.score_head = nn.Sequential(nn.Linear(256, 64), nn.GELU(), nn.Linear(64, 1))
+        self.pherr_head = nn.Sequential(nn.Linear(256, 64), nn.GELU(), nn.Linear(64, 1))
+
+    def forward(self, h, phone_id, gop, n_frames):
+        """Forward pass with pre-extracted hidden states (for batched training)."""
+        emb = self.phone_emb(phone_id)
+        x = torch.cat([h, emb, gop.unsqueeze(-1), n_frames.unsqueeze(-1)], dim=-1)
+        shared = self.shared(x)
+        score = self.score_head(shared).squeeze(-1)
+        pherr = self.pherr_head(shared).squeeze(-1)
+        return score, pherr
+
+    def extract_hidden(self, waveform):
+        """Extract hidden states from raw waveform. waveform: (1, T) tensor."""
+        # mask_time_prob=0.0 in config disables masking
+        out = self.backbone(waveform)
+        return out.last_hidden_state  # (1, T_frames, 1024)
+
+    def backbone_params(self):
+        """Return only the trainable backbone parameters."""
+        for name, param in self.backbone.named_parameters():
+            if param.requires_grad:
+                yield param
+
+    def head_params(self):
+        """Return head parameters (phone_emb + MLP)."""
+        yield from self.phone_emb.parameters()
+        yield from self.shared.parameters()
+        yield from self.score_head.parameters()
+        yield from self.pherr_head.parameters()
+
+
+# ============================================================
+# Alignment engine (frozen CTC model)
+# ============================================================
+class AlignmentEngine:
+    def __init__(self, device):
+        self.device = device
+        self.g2p = G2p()
+
+        print("Loading phoneme CTC model for alignment...")
+        self.ctc_model = Wav2Vec2ForCTC.from_pretrained(
+            "facebook/wav2vec2-xlsr-53-espeak-cv-ft"
+        ).to(device)
+        self.ctc_model.eval()
+        for p in self.ctc_model.parameters():
+            p.requires_grad = False
+
+        self.fe_ctc = Wav2Vec2FeatureExtractor.from_pretrained(
+            "facebook/wav2vec2-xlsr-53-espeak-cv-ft"
+        )
+
+        vocab_path = hf_hub_download("facebook/wav2vec2-xlsr-53-espeak-cv-ft", "vocab.json")
+        with open(vocab_path) as f:
+            self.vocab = json.load(f)
+        self.blank_idx = self.vocab.get("<pad>", 0)
+
+    def text_to_arpabet(self, text):
+        phones = self.g2p(text)
+        result = []
+        for ph in phones:
+            if ph == " ":
+                continue
+            clean = re.sub(r"\d", "", ph).lower()
+            if clean:
+                result.append(clean)
+        return result
+
+    def arpabet_to_idx(self, ph):
+        for ipa in ARPABET_TO_IPA.get(ph, []):
+            if ipa in self.vocab:
+                return self.vocab[ipa]
+        return self.vocab.get(ph, -1)
+
+    def viterbi_align(self, emissions, phone_indices):
+        T, C = emissions.shape
+        S = len(phone_indices)
+        if S == 0 or T < S:
+            return []
+        extended = [self.blank_idx]
+        for p in phone_indices:
+            extended.append(p)
+            extended.append(self.blank_idx)
+        S_ext = len(extended)
+        dp = np.full((T, S_ext), float("-inf"), dtype=np.float64)
+        bp = np.zeros((T, S_ext), dtype=np.int32)
+        dp[0][0] = emissions[0, extended[0]].item()
+        if S_ext > 1:
+            dp[0][1] = emissions[0, extended[1]].item()
+        for t in range(1, T):
+            for s in range(S_ext):
+                emit = emissions[t, extended[s]].item()
+                best, best_s = dp[t-1][s], s
+                if s > 0 and dp[t-1][s-1] > best:
+                    best, best_s = dp[t-1][s-1], s-1
+                if s > 1 and extended[s] != self.blank_idx and extended[s] != extended[s-2]:
+                    if dp[t-1][s-2] > best:
+                        best, best_s = dp[t-1][s-2], s-2
+                dp[t][s] = best + emit
+                bp[t][s] = best_s
+        s = S_ext-1 if (S_ext >= 2 and dp[T-1][S_ext-1] >= dp[T-1][S_ext-2]) else max(S_ext-2, 0)
+        path = []
+        for t in range(T-1, -1, -1):
+            path.append((t, extended[s]))
+            s = bp[t][s]
+        path.reverse()
+        return path
+
+    @torch.no_grad()
+    def align(self, audio_path, text):
+        """Return alignment info: list of (arpabet, ctc_frames, phone_id, gop)."""
+        waveform, sr = torchaudio.load(audio_path)
+        if waveform.shape[0] > 1:
+            waveform = waveform.mean(0, keepdim=True)
+        if sr != SAMPLE_RATE:
+            waveform = F_audio.resample(waveform, sr, SAMPLE_RATE)
+
+        wav_np = waveform.squeeze(0).numpy()
+
+        # CTC emissions
+        ctc_inputs = self.fe_ctc(wav_np, sampling_rate=SAMPLE_RATE, return_tensors="pt", padding=True)
+        ctc_logits = self.ctc_model(ctc_inputs.input_values.to(self.device)).logits
+        emissions = torch.log_softmax(ctc_logits, dim=-1).squeeze(0).cpu()
+
+        # G2P + alignment
+        arpabet = self.text_to_arpabet(text)
+        indices, valid = [], []
+        for ph in arpabet:
+            idx = self.arpabet_to_idx(ph)
+            if idx >= 0:
+                indices.append(idx)
+                valid.append(ph)
+        if not indices:
+            return [], emissions
+
+        path = self.viterbi_align(emissions, indices)
+        if not path:
+            return [], emissions
+
+        # Group frames by phone
+        segments = []
+        cur_tok, cur_frames = None, []
+        for f, tok in path:
+            if tok == self.blank_idx:
+                if cur_tok is not None:
+                    segments.append((cur_tok, cur_frames))
+                    cur_tok, cur_frames = None, []
+                continue
+            if tok != cur_tok:
+                if cur_tok is not None:
+                    segments.append((cur_tok, cur_frames))
+                cur_tok, cur_frames = tok, [f]
+            else:
+                cur_frames.append(f)
+        if cur_tok is not None:
+            segments.append((cur_tok, cur_frames))
+
+        results = []
+        for i, ph in enumerate(valid):
+            if i >= len(segments):
+                break
+            _, frames = segments[i]
+
+            # GOP
+            expected_idx = indices[i]
+            seg_emissions = emissions[frames]
+            target_lp = seg_emissions[:, expected_idx].mean().item()
+            mask = torch.ones(emissions.shape[1], dtype=torch.bool)
+            mask[self.blank_idx] = False
+            mask[expected_idx] = False
+            best_other = seg_emissions[:, mask].max(dim=-1).values.mean().item()
+            gop = target_lp - best_other
+
+            results.append({
+                "arpabet": ph,
+                "phone_id": PHONE_TO_ID.get(ph, 0),
+                "frames": frames,
+                "gop": gop,
+                "n_frames": len(frames),
+            })
+
+        return results, emissions
+
+
+# ============================================================
+# Dataset: stores pre-aligned info, extracts backbone features on-the-fly
+# ============================================================
+class PhonemeAlignedDataset(Dataset):
+    """
+    Each item is one phoneme with its alignment info.
+    During __getitem__, we return pre-computed data.
+    Backbone features are extracted in batch during training loop.
+    """
+    def __init__(self, samples, augment=False):
+        """
+        samples: list of dicts with keys:
+            audio_path, phone_id, gop, n_frames, ctc_frames,
+            y_score, y_pherr, T_ctc (CTC time steps for this audio)
+        """
+        self.samples = samples
+        self.augment = augment
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        s = self.samples[idx]
+        return {
+            "audio_path": s["audio_path"],
+            "phone_id": s["phone_id"],
+            "gop": s["gop"],
+            "n_frames": s["n_frames"],
+            "ctc_frames": s["ctc_frames"],
+            "T_ctc": s["T_ctc"],
+            "y_score": s["y_score"],
+            "y_pherr": s["y_pherr"],
+        }
+
+
+# ============================================================
+# Pre-align all data and cache alignment info
+# ============================================================
+def load_gt(xlsx_path):
+    wb = openpyxl.load_workbook(xlsx_path, read_only=True)
+    ws = wb.active
+    records = []
+    for row in ws.iter_rows(min_row=2, values_only=True):
+        fname, content, raw = row
+        if not fname or not content or not raw:
+            continue
+        s = raw
+        for _ in range(5):
+            s = s.replace("\\\\", "\\")
+        s = s.replace('\\"', '"')
+        phones = re.findall(
+            r'"char":"([^"]+)","ph2alpha":"([^"]*)".*?"pherr":(\d+).*?"score":([\d.]+)', s
+        )
+        gt = [{"phone": p.lower(), "pherr": int(e), "score": float(sc)} for p, _, e, sc in phones]
+        records.append({"file_name": fname, "content": content, "gt_phones": gt})
+    wb.close()
+    return records
+
+
+def pre_align_all(records, audio_dir, aligner, cache_path=None):
+    """Pre-compute alignment for all samples. Returns list of per-phoneme dicts."""
+    if cache_path and Path(cache_path).exists():
+        print(f"Loading cached alignments from {cache_path}")
+        return torch.load(cache_path, weights_only=False)
+
+    all_samples = []
+    n_skip = 0
+
+    for i, rec in enumerate(records):
+        audio_path = audio_dir / rec["file_name"]
+        if not audio_path.exists():
+            n_skip += 1
+            continue
+        try:
+            aligned, emissions = aligner.align(str(audio_path), rec["content"])
+            gt = rec["gt_phones"]
+            n = min(len(aligned), len(gt))
+            T_ctc = emissions.shape[0]
+
+            for j in range(n):
+                all_samples.append({
+                    "audio_path": str(audio_path),
+                    "phone_id": aligned[j]["phone_id"],
+                    "gop": aligned[j]["gop"],
+                    "n_frames": aligned[j]["n_frames"],
+                    "ctc_frames": aligned[j]["frames"],
+                    "T_ctc": T_ctc,
+                    "y_score": gt[j]["score"],
+                    "y_pherr": float(gt[j]["pherr"]),
+                })
+        except Exception as e:
+            n_skip += 1
+        if (i+1) % 200 == 0:
+            print(f"  [{i+1}/{len(records)}] phonemes={len(all_samples)}, skip={n_skip}")
+
+    print(f"  Total: {len(all_samples)} phonemes from {len(records)-n_skip} files, skip={n_skip}")
+
+    if cache_path:
+        torch.save(all_samples, cache_path)
+        print(f"  Cached to {cache_path}")
+
+    return all_samples
+
+
+# ============================================================
+# Training with backbone fine-tuning
+# ============================================================
+def extract_phoneme_features(model, audio_paths, frame_lists, T_ctcs, device, fe):
+    """
+    Extract backbone hidden states for a batch of phonemes.
+    Groups phonemes by audio file to avoid redundant forward passes.
+    Returns: (B, 1024) tensor of per-phoneme features.
+    """
+    # Group by audio path
+    path_to_indices = {}
+    for i, path in enumerate(audio_paths):
+        if path not in path_to_indices:
+            path_to_indices[path] = []
+        path_to_indices[path].append(i)
+
+    features = torch.zeros(len(audio_paths), 1024, device=device)
+
+    for path, indices in path_to_indices.items():
+        # Load audio once
+        waveform, sr = torchaudio.load(path)
+        if waveform.shape[0] > 1:
+            waveform = waveform.mean(0, keepdim=True)
+        if sr != SAMPLE_RATE:
+            waveform = F_audio.resample(waveform, sr, SAMPLE_RATE)
+
+        wav_np = waveform.squeeze(0).numpy()
+        inputs = fe(wav_np, sampling_rate=SAMPLE_RATE, return_tensors="pt", padding=True)
+        hidden = model.extract_hidden(inputs.input_values.to(device))  # (1, T_h, 1024)
+        hidden = hidden.squeeze(0)  # (T_h, 1024)
+        T_h = hidden.shape[0]
+
+        for idx in indices:
+            T_ctc = T_ctcs[idx]
+            scale = T_h / max(T_ctc, 1)
+            frames = frame_lists[idx]
+            h_start = max(0, int(min(frames) * scale))
+            h_end = min(T_h, int((max(frames) + 1) * scale))
+            if h_end <= h_start:
+                h_end = h_start + 1
+            features[idx] = hidden[h_start:h_end].mean(dim=0)
+
+    return features
+
+
+def collate_fn(batch):
+    """Custom collate that handles variable-length frame lists."""
+    return {
+        "audio_paths": [b["audio_path"] for b in batch],
+        "phone_id": torch.tensor([b["phone_id"] for b in batch], dtype=torch.long),
+        "gop": torch.tensor([b["gop"] for b in batch], dtype=torch.float32),
+        "n_frames": torch.tensor([b["n_frames"] for b in batch], dtype=torch.float32),
+        "ctc_frames": [b["ctc_frames"] for b in batch],
+        "T_ctc": [b["T_ctc"] for b in batch],
+        "y_score": torch.tensor([b["y_score"] for b in batch], dtype=torch.float32),
+        "y_pherr": torch.tensor([b["y_pherr"] for b in batch], dtype=torch.float32),
+    }
+
+
+def train(all_samples, device, epochs=30, batch_size=64, lr_backbone=1e-5, lr_head=5e-4,
+          unfreeze_top_n=6, grad_accum=4):
+    """
+    Fine-tune backbone + train head jointly.
+    """
+    # Split by audio file to avoid data leakage
+    audio_files = list(set(s["audio_path"] for s in all_samples))
+    files_train, files_test = train_test_split(audio_files, test_size=0.15, random_state=42)
+    files_train, files_val = train_test_split(files_train, test_size=0.12, random_state=42)
+
+    train_set = set(files_train)
+    val_set = set(files_val)
+    test_set = set(files_test)
+
+    train_samples = [s for s in all_samples if s["audio_path"] in train_set]
+    val_samples = [s for s in all_samples if s["audio_path"] in val_set]
+    test_samples = [s for s in all_samples if s["audio_path"] in test_set]
+
+    print(f"\nSplit by audio file:")
+    print(f"  Train: {len(train_samples)} phonemes from {len(files_train)} files")
+    print(f"  Val:   {len(val_samples)} phonemes from {len(files_val)} files")
+    print(f"  Test:  {len(test_samples)} phonemes from {len(files_test)} files")
+
+    train_ds = PhonemeAlignedDataset(train_samples, augment=True)
+    val_ds = PhonemeAlignedDataset(val_samples)
+    test_ds = PhonemeAlignedDataset(test_samples)
+
+    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,
+                              collate_fn=collate_fn, num_workers=0)
+    val_loader = DataLoader(val_ds, batch_size=batch_size, collate_fn=collate_fn, num_workers=0)
+
+    # Model
+    model = BackbonePhoneScorer(unfreeze_top_n=unfreeze_top_n).to(device)
+
+    # Count parameters
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"\n  Total params: {total_params/1e6:.1f}M")
+    print(f"  Trainable params: {trainable_params/1e6:.1f}M")
+
+    # Differential learning rate
+    optimizer = torch.optim.AdamW([
+        {"params": model.backbone_params(), "lr": lr_backbone},
+        {"params": model.head_params(), "lr": lr_head},
+    ], weight_decay=1e-3)
+
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
+
+    # Pos weight for imbalanced pherr
+    pherr_vals = torch.tensor([s["y_pherr"] for s in train_samples])
+    pos_rate = pherr_vals.mean().item()
+    pos_weight = torch.tensor([(1 - pos_rate) / max(pos_rate, 0.01)]).to(device)
+    print(f"  Pherr pos rate: {pos_rate:.3f}, pos_weight: {pos_weight.item():.2f}")
+
+    best_val_loss = float("inf")
+    best_state = None
+    patience, patience_counter = 8, 0
+
+    fe = model.fe_backbone
+
+    n_steps_per_epoch = len(train_loader)
+    print(f"  Steps per epoch: {n_steps_per_epoch}")
+
+    for epoch in range(epochs):
+        model.train()
+        losses = []
+
+        optimizer.zero_grad()
+        for step, batch in enumerate(train_loader):
+            # Extract backbone features (with gradient for unfrozen layers)
+            h = extract_phoneme_features(
+                model, batch["audio_paths"], batch["ctc_frames"],
+                batch["T_ctc"], device, fe
+            )
+
+            # Add noise augmentation
+            h = h + torch.randn_like(h) * 0.01
+
+            pid = batch["phone_id"].to(device)
+            gop = batch["gop"].to(device)
+            nf = batch["n_frames"].to(device)
+            ys = batch["y_score"].to(device)
+            yp = batch["y_pherr"].to(device)
+
+            pred_s, pred_p = model(h, pid, gop, nf)
+            loss_s = F.mse_loss(pred_s, ys)
+            loss_p = F.binary_cross_entropy_with_logits(pred_p, yp, pos_weight=pos_weight)
+            loss = (loss_s / 100.0 + loss_p) / grad_accum
+
+            loss.backward()
+
+            if (step + 1) % grad_accum == 0 or (step + 1) == len(train_loader):
+                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+                optimizer.step()
+                optimizer.zero_grad()
+
+            losses.append(loss.item() * grad_accum)
+
+            if (step + 1) % 50 == 0:
+                avg_loss = np.mean(losses[-50:])
+                print(f"    Epoch {epoch+1} Step {step+1}/{n_steps_per_epoch}  loss={avg_loss:.4f}")
+
+        scheduler.step()
+
+        # Validate
+        model.eval()
+        val_losses = []
+        with torch.no_grad():
+            for batch in val_loader:
+                h = extract_phoneme_features(
+                    model, batch["audio_paths"], batch["ctc_frames"],
+                    batch["T_ctc"], device, fe
+                )
+                pid = batch["phone_id"].to(device)
+                gop = batch["gop"].to(device)
+                nf = batch["n_frames"].to(device)
+                ys = batch["y_score"].to(device)
+                yp = batch["y_pherr"].to(device)
+
+                pred_s, pred_p = model(h, pid, gop, nf)
+                loss = F.mse_loss(pred_s, ys) / 100.0 + \
+                       F.binary_cross_entropy_with_logits(pred_p, yp, pos_weight=pos_weight)
+                val_losses.append(loss.item())
+
+        vl = np.mean(val_losses)
+        tl = np.mean(losses)
+
+        if vl < best_val_loss:
+            best_val_loss = vl
+            best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
+            patience_counter = 0
+            marker = " *"
+        else:
+            patience_counter += 1
+            marker = ""
+
+        print(f"  Epoch {epoch+1:3d}/{epochs}  train={tl:.4f}  val={vl:.4f}  "
+              f"best={best_val_loss:.4f}  patience={patience_counter}/{patience}{marker}")
+
+        if patience_counter >= patience:
+            print(f"  Early stopping at epoch {epoch+1}")
+            break
+
+    # Load best model
+    model.load_state_dict(best_state)
+    model = model.to(device)
+
+    # ============================================================
+    # Test evaluation
+    # ============================================================
+    print(f"\n{'='*60}")
+    print("TEST EVALUATION")
+    print(f"{'='*60}")
+
+    model.eval()
+    all_pred_s, all_pred_p, all_gt_s, all_gt_p = [], [], [], []
+
+    test_loader = DataLoader(test_ds, batch_size=batch_size, collate_fn=collate_fn, num_workers=0)
+    with torch.no_grad():
+        for batch in test_loader:
+            h = extract_phoneme_features(
+                model, batch["audio_paths"], batch["ctc_frames"],
+                batch["T_ctc"], device, fe
+            )
+            pid = batch["phone_id"].to(device)
+            gop = batch["gop"].to(device)
+            nf = batch["n_frames"].to(device)
+
+            pred_s, pred_p = model(h, pid, gop, nf)
+            all_pred_s.append(pred_s.cpu())
+            all_pred_p.append(torch.sigmoid(pred_p).cpu())
+            all_gt_s.append(batch["y_score"])
+            all_gt_p.append(batch["y_pherr"])
+
+    pred_s = torch.cat(all_pred_s).numpy()
+    pred_p = torch.cat(all_pred_p).numpy()
+    gt_s = torch.cat(all_gt_s).numpy()
+    gt_p = torch.cat(all_gt_p).numpy()
+
+    pred_s_clip = np.clip(pred_s, 0, 100)
+
+    # Score metrics
+    mae = np.mean(np.abs(gt_s - pred_s_clip))
+    rmse = np.sqrt(np.mean((gt_s - pred_s_clip) ** 2))
+    corr, _ = stats.pearsonr(gt_s, pred_s_clip)
+    sp, _ = stats.spearmanr(gt_s, pred_s_clip)
+
+    print(f"\nPHONE SCORE PREDICTION")
+    print(f"  MAE:      {mae:.2f}")
+    print(f"  RMSE:     {rmse:.2f}")
+    print(f"  Pearson:  {corr:.3f}")
+    print(f"  Spearman: {sp:.3f}")
+    abs_err = np.abs(gt_s - pred_s_clip)
+    for t in [5, 10, 15, 20, 30]:
+        print(f"  |err|<{t:2d}: {np.mean(abs_err<t)*100:.1f}%")
+
+    # Pherr metrics
+    auc = roc_auc_score(gt_p, pred_p) if len(np.unique(gt_p)) > 1 else 0
+    best_f1, best_thresh = 0, 0.5
+    for th in np.arange(0.1, 0.9, 0.05):
+        pb = (pred_p >= th).astype(int)
+        _, _, f1, _ = precision_recall_fscore_support(gt_p, pb, average="binary", zero_division=0)
+        if f1 > best_f1:
+            best_f1, best_thresh = f1, th
+
+    pb = (pred_p >= best_thresh).astype(int)
+    prec, rec, f1, _ = precision_recall_fscore_support(gt_p, pb, average="binary")
+    tp = ((pb==1) & (gt_p==1)).sum()
+    fp = ((pb==1) & (gt_p==0)).sum()
+    fn = ((pb==0) & (gt_p==1)).sum()
+    tn = ((pb==0) & (gt_p==0)).sum()
+
+    print(f"\nPHONE ERROR DETECTION")
+    print(f"  AUC-ROC:   {auc:.3f}")
+    print(f"  Threshold: {best_thresh:.2f}")
+    print(f"  Precision: {prec:.3f}")
+    print(f"  Recall:    {rec:.3f}")
+    print(f"  F1:        {f1:.3f}")
+    print(f"  TP={tp:4d}  FP={fp:4d}")
+    print(f"  FN={fn:4d}  TN={tn:4d}")
+
+    # Comparison
+    print(f"\n{'='*60}")
+    print("COMPARISON WITH PREVIOUS METHODS")
+    print(f"{'='*60}")
+    print(f"  {'Method':<35s}  {'AUC':>6s}  {'F1':>6s}  {'Prec':>6s}  {'Rec':>6s}  {'Pearson':>8s}  {'MAE':>6s}")
+    methods = [
+        ("GOP threshold (v1.0, 1K data)",   0.738, 0.476, 0.379, 0.638, 0.372, 27.44),
+        ("E2E MLP frozen (v2, 1K data)",    0.814, 0.565, 0.500, 0.650, 0.528, 22.57),
+        ("Phoneme comparison (1K data)",     0.691, 0.492, 0.379, 0.703, None,  None),
+        (f"WavLM finetune (11K data)",       auc,   best_f1, prec, rec, corr,   mae),
+    ]
+    for name, a, f, p, r, c, m in methods:
+        c_str = f"{c:.3f}" if c is not None else "  N/A"
+        m_str = f"{m:.2f}" if m is not None else "  N/A"
+        print(f"  {name:<35s}  {a:>6.3f}  {f:>6.3f}  {p:>6.3f}  {r:>6.3f}  {c_str:>8s}  {m_str:>6s}")
+
+    # Save model
+    save_dir = Path("/mnt/weka/home/jianshu.she/Voice-correction")
+    save_path = save_dir / "wavlm_finetuned.pt"
+    torch.save({
+        "model_state": model.state_dict(),
+        "unfreeze_top_n": unfreeze_top_n,
+        "metrics": {
+            "auc": auc, "f1": best_f1, "precision": float(prec),
+            "recall": float(rec), "pearson": corr, "mae": mae,
+        },
+    }, save_path)
+    print(f"\nModel saved to {save_path}")
+
+    return model
+
+
+def main():
+    device = torch.device("cuda:0")
+
+    # Load both datasets
+    print("Loading datasets...")
+    records_sent = load_gt("/mnt/weka/home/jianshu.she/Voice-correction/eval_log.xlsx")
+    records_word = load_gt("/mnt/weka/home/jianshu.she/Voice-correction/word_eval_log.xlsx")
+    print(f"  Sentence data: {len(records_sent)} records")
+    print(f"  Word data: {len(records_word)} records")
+
+    # Tag audio directories
+    for r in records_sent:
+        r["audio_dir"] = "audio_files"
+    for r in records_word:
+        r["audio_dir"] = "words_audio_files"
+
+    # Pre-align all data
+    aligner = AlignmentEngine(device)
+
+    base_dir = Path("/mnt/weka/home/jianshu.she/Voice-correction")
+
+    cache_sent = base_dir / "align_cache_sentences.pt"
+    cache_word = base_dir / "align_cache_words.pt"
+
+    print("\nAligning sentence data...")
+    samples_sent = pre_align_all(records_sent, base_dir / "audio_files", aligner, cache_sent)
+
+    print("\nAligning word data...")
+    samples_word = pre_align_all(records_word, base_dir / "words_audio_files", aligner, cache_word)
+
+    all_samples = samples_sent + samples_word
+    print(f"\nTotal: {len(all_samples)} phonemes")
+
+    pherr_rate = np.mean([s["y_pherr"] for s in all_samples])
+    print(f"Overall error rate: {pherr_rate:.3f}")
+
+    # Train
+    train(all_samples, device, epochs=30, batch_size=64, unfreeze_top_n=6)
+
+
+if __name__ == "__main__":
+    main()