scripts/train_whisper_emotion_head.py

#!/usr/bin/env python3
"""Whisper-Medium Encoder + Linear Emotion Head 학습.

Whisper encoder를 freeze하고 linear classifier head만 학습하여
benchmark_ser_models.py의 WhisperMediumAdapter에 사용할 체크포인트를 생성한다.

Usage:
    # AI Hub 테스트 서브셋 외의 데이터로 학습 (test leakage 방지)
    python scripts/train_whisper_emotion_head.py \\
        --train-dir data/evaluation/korean/train_audio \\
        --val-dir data/evaluation/korean/val_audio \\
        --output data/models/whisper_emotion_head.pt

    # prepare_aihub_test_subset.py의 출력으로 quick test (test leakage 주의)
    python scripts/train_whisper_emotion_head.py \\
        --train-dir data/evaluation/korean/test_audio \\
        --epochs 3 --output data/models/whisper_emotion_head.pt
"""

from __future__ import annotations

import argparse
import glob
import logging
import os
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

EVAL_LABELS = ["neutral", "joy", "sadness", "anger", "surprise", "fear"]
LABEL_TO_IDX = {label: i for i, label in enumerate(EVAL_LABELS)}


class EmotionAudioDataset(Dataset):
    """Load WAV files organized by emotion class directory."""

    def __init__(self, root_dir: str, processor, max_samples_per_class: int | None = None):
        self.samples = []
        self.processor = processor

        for label in EVAL_LABELS:
            class_dir = Path(root_dir) / label
            if not class_dir.exists():
                logger.warning("Class directory not found: %s", class_dir)
                continue

            wavs = sorted(glob.glob(str(class_dir / "*.wav")))
            if max_samples_per_class and len(wavs) > max_samples_per_class:
                wavs = wavs[:max_samples_per_class]

            for wav_path in wavs:
                self.samples.append({
                    "path": wav_path,
                    "label": LABEL_TO_IDX[label],
                })

        logger.info("Dataset: %d samples from %s", len(self.samples), root_dir)

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        sample = self.samples[idx]
        import librosa
        audio, sr = librosa.load(sample["path"], sr=16000)
        inputs = self.processor(audio, sampling_rate=16000, return_tensors="pt")
        features = inputs.input_features.squeeze(0)  # (n_mels, T)
        return features, sample["label"]


def collate_fn(batch):
    features, labels = zip(*batch)
    # Pad features to same length
    max_len = max(f.shape[1] for f in features)
    padded = []
    for f in features:
        if f.shape[1] < max_len:
            pad = torch.zeros(f.shape[0], max_len - f.shape[1])
            f = torch.cat([f, pad], dim=1)
        padded.append(f)
    return torch.stack(padded), torch.tensor(labels, dtype=torch.long)


def train(args):
    from transformers import WhisperModel, WhisperFeatureExtractor

    device = torch.device(args.device)

    # Load Whisper encoder (frozen)
    logger.info("Loading Whisper-Medium encoder...")
    processor = WhisperFeatureExtractor.from_pretrained("openai/whisper-medium")
    whisper = WhisperModel.from_pretrained("openai/whisper-medium").to(device)
    whisper.eval()
    for param in whisper.parameters():
        param.requires_grad = False

    hidden_dim = whisper.config.d_model  # 1024
    head = nn.Linear(hidden_dim, len(EVAL_LABELS)).to(device)

    # Dataset
    train_dataset = EmotionAudioDataset(args.train_dir, processor, args.max_samples_per_class)
    if len(train_dataset) == 0:
        logger.error("No training samples found")
        return

    train_loader = DataLoader(
        train_dataset, batch_size=args.batch_size, shuffle=True,
        collate_fn=collate_fn, num_workers=0,
    )

    val_loader = None
    if args.val_dir and Path(args.val_dir).exists():
        val_dataset = EmotionAudioDataset(args.val_dir, processor)
        if len(val_dataset) > 0:
            val_loader = DataLoader(
                val_dataset, batch_size=args.batch_size, shuffle=False,
                collate_fn=collate_fn, num_workers=0,
            )

    # Optimizer
    optimizer = torch.optim.Adam(head.parameters(), lr=args.lr)
    criterion = nn.CrossEntropyLoss()

    # Training loop
    best_val_acc = 0.0
    for epoch in range(args.epochs):
        head.train()
        total_loss = 0
        correct = 0
        total = 0

        for batch_idx, (features, labels) in enumerate(train_loader):
            features = features.to(device)
            labels = labels.to(device)

            with torch.no_grad():
                encoder_out = whisper.encoder(features)
                hidden = encoder_out.last_hidden_state  # (B, T, D)
                pooled = hidden.mean(dim=1)  # (B, D)

            logits = head(pooled)
            loss = criterion(logits, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item() * labels.size(0)
            preds = logits.argmax(dim=1)
            correct += (preds == labels).sum().item()
            total += labels.size(0)

        train_acc = correct / max(total, 1)
        avg_loss = total_loss / max(total, 1)
        logger.info("Epoch %d/%d: loss=%.4f, train_acc=%.3f",
                     epoch + 1, args.epochs, avg_loss, train_acc)

        # Validation
        if val_loader:
            head.eval()
            val_correct = 0
            val_total = 0
            with torch.no_grad():
                for features, labels in val_loader:
                    features = features.to(device)
                    labels = labels.to(device)
                    encoder_out = whisper.encoder(features)
                    pooled = encoder_out.last_hidden_state.mean(dim=1)
                    logits = head(pooled)
                    preds = logits.argmax(dim=1)
                    val_correct += (preds == labels).sum().item()
                    val_total += labels.size(0)
            val_acc = val_correct / max(val_total, 1)
            logger.info("  val_acc=%.3f", val_acc)

            if val_acc > best_val_acc:
                best_val_acc = val_acc
                save_checkpoint(head, args.output, epoch, val_acc)
        else:
            # No val set — save latest
            save_checkpoint(head, args.output, epoch, train_acc)

    logger.info("Training complete. Best checkpoint: %s", args.output)


def save_checkpoint(head: nn.Linear, path: str, epoch: int, accuracy: float):
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    torch.save(head.state_dict(), path)
    logger.info("Saved checkpoint: %s (epoch=%d, acc=%.3f)", path, epoch + 1, accuracy)


def main():
    parser = argparse.ArgumentParser(description="Train Whisper emotion classifier head")
    parser.add_argument("--train-dir", required=True,
                        help="학습 오디오 디렉토리 ({emotion}/*.wav 구조)")
    parser.add_argument("--val-dir", default=None,
                        help="검증 오디오 디렉토리 (없으면 train accuracy로 판단)")
    parser.add_argument("--output", default="data/models/whisper_emotion_head.pt",
                        help="출력 체크포인트 경로")
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
    parser.add_argument("--max-samples-per-class", type=int, default=None,
                        help="클래스당 최대 학습 샘플 수")
    args = parser.parse_args()

    train(args)


if __name__ == "__main__":
    main()