scripts/prepare_lora_dataset.py

#!/usr/bin/env python3
"""Unified data preprocessing pipeline for LoRA emotion2vec 7-class fine-tuning.

Extracts, preprocesses, and merges samples from three sources:
  1. AI Hub 263  — anchor (acted Korean, 7-class)
  2. AI Hub 71631 — booster (outdoor spontaneous Korean, mapped to 7-class)
  3. RAVDESS     — English (acted, 7-class)

Outputs a unified manifest (train/val JSONs) ready for LoRA training.

Usage:
    python scripts/prepare_lora_dataset.py \
        --anchor-dir   "data/AI Hub 263" \
        --booster-label-zip "data/AI Hub 71631/01-1.정식개방데이터/Training/02.라벨링데이터/TL_02.실외.zip" \
        --booster-audio-zip "data/AI Hub 71631/01-1.정식개방데이터/Training/01.원천데이터/TS_02.실외.zip" \
        --ravdess-dir  data/ravdess \
        --output-dir   data/lora_7class
"""

from __future__ import annotations

import argparse
import csv
import io
import json
import logging
import random
import zipfile
from collections import Counter, defaultdict
from pathlib import Path

import torch
import torchaudio

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(levelname)-8s  %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
LABEL2IDX: dict[str, int] = {
    "happiness": 0,
    "anger": 1,
    "disgust": 2,
    "fear": 3,
    "neutral": 4,
    "sadness": 5,
    "surprise": 6,
}

VALID_LABELS = set(LABEL2IDX.keys())

TARGET_SR = 16_000
RMS_THRESHOLD = 0.001  # 0.005→0.001: disgust 등 저에너지 발화 보존 (진짜 무음만 제거)

# ---------------------------------------------------------------------------
# Task 1 — Label Mappers
# ---------------------------------------------------------------------------

_MAP_263: dict[str, str] = {
    "angry": "anger",
    "happiness": "happiness",
    "neutral": "neutral",
    "sadness": "sadness",
    "surprise": "surprise",
    "fear": "fear",
    "disgust": "disgust",
}


def map_263_label(raw: str) -> str | None:
    """Map AI Hub 263 annotator label to 7-class. Case-insensitive."""
    if not raw:
        return None
    return _MAP_263.get(raw.strip().lower())


_MAP_71631: dict[str, str] = {
    "기쁨": "happiness",
    "화남": "anger",
    "놀라움": "surprise",
    "슬픔": "sadness",
    "두려움": "fear",
    "없음": "neutral",
    "중립": "neutral",
}


def map_71631_label(raw: str) -> str | None:
    """Map AI Hub 71631 Korean emotion label to 7-class."""
    if not raw:
        return None
    return _MAP_71631.get(raw.strip())


def map_ravdess_label(raw: str) -> str | None:
    """Map RAVDESS label to 7-class. Passthrough except joy→happiness."""
    if not raw:
        return None
    lbl = raw.strip().lower()
    if lbl == "joy":
        return "happiness"
    if lbl in VALID_LABELS:
        return lbl
    return None


def majority_vote_263(emotions: list[str | None]) -> str | None:
    """Return majority label (3/5+) or None for no majority/tie."""
    valid = [e for e in emotions if e is not None]
    if not valid:
        return None
    counts = Counter(valid)
    top_label, top_count = counts.most_common(1)[0]
    if top_count < 3:
        return None
    # Check for tie at top count
    tied = [lbl for lbl, c in counts.items() if c == top_count]
    if len(tied) > 1:
        return None
    return top_label


# ---------------------------------------------------------------------------
# Task 1 — Audio Preprocessing
# ---------------------------------------------------------------------------


import re

def _clean_text(text: str) -> str:
    """Clean text for STT-friendly format.

    Removes non-verbal tags like (웃음), (한숨), keeps ?, !, ...
    """
    # Remove non-verbal tags: (웃음), (한숨), (침묵), [noise], etc.
    text = re.sub(r"[(\[（][^)\]）]*[)\]）]", "", text)
    # Remove trailing/leading whitespace, collapse multiple spaces
    text = re.sub(r"\s+", " ", text).strip()
    return text


def _compute_rms(waveform: torch.Tensor) -> float:
    """Compute RMS of a waveform tensor."""
    return float(torch.sqrt(torch.mean(waveform.float() ** 2)))


def preprocess_audio(input_path: Path, output_path: Path) -> bool:
    """Resample to 16kHz mono, trim silence, reject if RMS < threshold.

    Returns True if file was saved, False if rejected.
    """
    input_path = Path(input_path)
    output_path = Path(output_path)

    waveform, sr = torchaudio.load(str(input_path))

    # Mono
    if waveform.shape[0] > 1:
        waveform = waveform.mean(dim=0, keepdim=True)

    # Resample
    if sr != TARGET_SR:
        waveform = torchaudio.functional.resample(waveform, sr, TARGET_SR)

    # Trim silence (leading/trailing)
    waveform_trimmed = torchaudio.functional.vad(waveform, TARGET_SR)
    if waveform_trimmed.numel() > 0:
        waveform = waveform_trimmed

    # RMS check
    rms = _compute_rms(waveform)
    if rms < RMS_THRESHOLD:
        return False

    output_path.parent.mkdir(parents=True, exist_ok=True)
    torchaudio.save(str(output_path), waveform, TARGET_SR)
    return True


def preprocess_audio_from_tensor(
    waveform: torch.Tensor, sr: int, output_path: Path
) -> bool:
    """Preprocess an in-memory waveform tensor and save to output_path.

    Used for 71631 where we slice in memory from the full conversation wav.
    """
    output_path = Path(output_path)

    # Mono
    if waveform.dim() > 1 and waveform.shape[0] > 1:
        waveform = waveform.mean(dim=0, keepdim=True)
    elif waveform.dim() == 1:
        waveform = waveform.unsqueeze(0)

    # Resample
    if sr != TARGET_SR:
        waveform = torchaudio.functional.resample(waveform, sr, TARGET_SR)

    # Trim silence (leading/trailing) — same as preprocess_audio
    waveform_trimmed = torchaudio.functional.vad(waveform, TARGET_SR)
    if waveform_trimmed.numel() > 0:
        waveform = waveform_trimmed

    # RMS check
    rms = _compute_rms(waveform)
    if rms < RMS_THRESHOLD:
        return False

    output_path.parent.mkdir(parents=True, exist_ok=True)
    torchaudio.save(str(output_path), waveform, TARGET_SR)
    return True


# ---------------------------------------------------------------------------
# Task 2 — AI Hub 263 Extraction
# ---------------------------------------------------------------------------


def parse_263_row(row: list[str]) -> dict | None:
    """Parse a CSV row from AI Hub 263.

    Columns: wav_id, 발화문, 상황, 1번감정, 1번감정세기, 2번감정, 2번감정세기,
             3번감정, 3번감정세기, 4번감정, 4번감정세기, 5번감정, 5번감정세기, 나이, 성별

    Returns dict with wav_id, label, agreement, max_intensity, or None if no majority.
    """
    if len(row) < 15:
        return None

    wav_id = row[0].strip()
    # Extract 5 annotator emotions and intensities
    annotator_emotions: list[str | None] = []
    intensities: list[int] = []
    for i in range(5):
        emo_col = 3 + i * 2      # 3, 5, 7, 9, 11
        int_col = 3 + i * 2 + 1  # 4, 6, 8, 10, 12
        raw_emo = row[emo_col].strip() if emo_col < len(row) else ""
        raw_int = row[int_col].strip() if int_col < len(row) else "0"
        mapped = map_263_label(raw_emo)
        annotator_emotions.append(mapped)
        try:
            intensities.append(int(raw_int))
        except ValueError:
            intensities.append(0)

    label = majority_vote_263(annotator_emotions)
    if label is None:
        return None

    agreement = sum(1 for e in annotator_emotions if e == label)
    # Max intensity among annotators who voted for the majority label
    max_intensity = max(
        (intensities[i] for i, e in enumerate(annotator_emotions) if e == label),
        default=0,
    )

    # 발화문 (text) — column 1
    text = row[1].strip() if len(row) > 1 else ""
    text = _clean_text(text)

    return {
        "wav_id": wav_id,
        "text": text,
        "label": label,
        "agreement": agreement,
        "max_intensity": max_intensity,
    }


def extract_anchor_263(
    anchor_dir: Path,
    output_dir: Path,
    cap_per_class: int = 2100,
) -> list[dict]:
    """Extract and preprocess AI Hub 263 dataset.

    Parses 3 CSVs (cp949), majority-votes annotator labels,
    priority-sorts (agreement desc, intensity desc), caps per class,
    extracts wavs from ZIPs, preprocesses audio.
    """
    anchor_dir = Path(anchor_dir)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    csv_files = sorted(anchor_dir.glob("*.csv"))
    zip_files = sorted(anchor_dir.glob("*.zip"))

    log.info("263: Found %d CSVs, %d ZIPs", len(csv_files), len(zip_files))

    # Step 1: Parse all CSVs
    all_parsed: list[dict] = []
    for csv_path in csv_files:
        with open(csv_path, encoding="cp949", newline="") as f:
            reader = csv.reader(f)
            header = next(reader)  # skip header
            for row in reader:
                result = parse_263_row(row)
                if result is not None:
                    result["csv_source"] = csv_path.stem
                    all_parsed.append(result)

    log.info("263: Parsed %d rows with majority vote", len(all_parsed))

    # Step 2: Group by label, priority sort, cap
    by_label: dict[str, list[dict]] = defaultdict(list)
    for item in all_parsed:
        by_label[item["label"]].append(item)

    selected: list[dict] = []
    for label, items in by_label.items():
        # Sort by agreement desc, then intensity desc
        items.sort(key=lambda x: (x["agreement"], x["max_intensity"]), reverse=True)
        capped = items[:cap_per_class]
        selected.extend(capped)
        log.info("263: %s — %d available, %d selected", label, len(items), len(capped))

    # Step 3: Build wav_id → zip lookup
    wav_to_zip: dict[str, tuple[zipfile.ZipFile, str]] = {}
    zip_handles = [zipfile.ZipFile(zp) for zp in zip_files]
    for zf in zip_handles:
        for name in zf.namelist():
            if name.endswith(".wav"):
                basename = Path(name).stem
                wav_to_zip[basename] = (zf, name)

    # Step 4: Extract and preprocess
    samples: list[dict] = []
    skipped = 0
    for item in selected:
        wav_id = item["wav_id"]
        if wav_id not in wav_to_zip:
            skipped += 1
            continue

        zf, zip_entry = wav_to_zip[wav_id]
        out_path = output_dir / item["label"] / f"{wav_id}.wav"

        try:
            with zf.open(zip_entry) as src:
                audio_bytes = src.read()

            # Write to temp, then preprocess
            tmp_path = output_dir / f"_tmp_{wav_id}.wav"
            tmp_path.write_bytes(audio_bytes)

            ok = preprocess_audio(tmp_path, out_path)
            tmp_path.unlink(missing_ok=True)

            if ok:
                samples.append({
                    "path": str(out_path),
                    "label": item["label"],
                    "label_idx": LABEL2IDX[item["label"]],
                    "source": "263",
                    "speaker_id": f"263_{wav_id[:8]}",
                    "text": item.get("text", ""),
                    "agreement": item["agreement"],
                    "intensity": item["max_intensity"],
                })
            else:
                skipped += 1
        except Exception as e:
            log.warning("263: Failed to process %s: %s", wav_id, e)
            skipped += 1

    # Close zip handles
    for zf in zip_handles:
        zf.close()

    log.info("263: Extracted %d samples, skipped %d", len(samples), skipped)
    return samples


# ---------------------------------------------------------------------------
# Task 3 — AI Hub 71631 Outdoor Extraction
# ---------------------------------------------------------------------------


def parse_71631_utterance(conv_entry: dict) -> dict | None:
    """Parse a conversation entry from 71631 JSON.

    Filters by VerifyEmotionLevel (보통/강함 only, rejects 약함).
    Returns dict with label, intensity, start_time, end_time, speaker_no or None.
    """
    level = conv_entry.get("VerifyEmotionLevel", "")
    if level not in ("보통", "강함"):
        return None

    emotion = conv_entry.get("VerifyEmotionTarget", "")
    label = map_71631_label(emotion)
    if label is None:
        return None

    try:
        start_time = float(conv_entry["StartTime"])
        end_time = float(conv_entry["EndTime"])
    except (KeyError, ValueError):
        return None

    if end_time <= start_time:
        return None

    # Text from conversation entry
    text = _clean_text(conv_entry.get("Text", ""))

    return {
        "label": label,
        "intensity": level,
        "start_time": start_time,
        "end_time": end_time,
        "speaker_no": conv_entry.get("SpeakerNo", ""),
        "text": text,
    }


def extract_booster_71631(
    label_zip: Path,
    audio_zip: Path,
    output_dir: Path,
    cap: int = 3500,
    max_per_speaker: int = 20,
) -> list[dict]:
    """Extract and preprocess AI Hub 71631 outdoor dataset.

    Parses label ZIP JSONs, filters intensity, applies speaker cap,
    slices wav segments, resamples to 16kHz, RMS-filters neutral.
    """
    label_zip = Path(label_zip)
    audio_zip = Path(audio_zip)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Step 1: Parse all label JSONs
    all_utterances: list[dict] = []
    with zipfile.ZipFile(label_zip) as lzf:
        json_files = [n for n in lzf.namelist() if n.endswith(".json")]
        log.info("71631: Found %d label JSONs", len(json_files))

        for jf in json_files:
            try:
                with lzf.open(jf) as f:
                    data = json.load(f)
            except Exception as e:
                log.warning("71631: Failed to parse %s: %s", jf, e)
                continue

            filename = data.get("File", {}).get("FileName", "")
            conv_id = filename  # Use filename as conversation_id
            spk1_id = data.get("Speaker1", {}).get("ID", "")
            spk2_id = data.get("Speaker2", {}).get("ID", "")

            for entry in data.get("Conversation", []):
                parsed = parse_71631_utterance(entry)
                if parsed is None:
                    continue
                # Determine speaker ID
                spk_no = parsed["speaker_no"]
                if spk_no == "Speaker1":
                    spk_id = spk1_id
                elif spk_no == "Speaker2":
                    spk_id = spk2_id
                else:
                    spk_id = spk_no

                parsed["conversation_id"] = conv_id
                parsed["speaker_id"] = f"71631_{spk_id}"
                parsed["wav_filename"] = filename
                parsed["text_no"] = entry.get("TextNo", "")
                all_utterances.append(parsed)

    log.info("71631: Parsed %d utterances (보통/강함)", len(all_utterances))

    # Step 2: Speaker cap
    speaker_counts: Counter = Counter()
    speaker_capped: list[dict] = []
    # Priority: 강함 first, then 보통
    all_utterances.sort(key=lambda x: (0 if x["intensity"] == "강함" else 1))
    for utt in all_utterances:
        spk = utt["speaker_id"]
        if speaker_counts[spk] < max_per_speaker:
            speaker_capped.append(utt)
            speaker_counts[spk] += 1

    log.info("71631: After speaker cap (%d/spk): %d utterances", max_per_speaker, len(speaker_capped))

    # Step 3: Group by label, cap per class
    by_label: dict[str, list[dict]] = defaultdict(list)
    for utt in speaker_capped:
        by_label[utt["label"]].append(utt)

    selected: list[dict] = []
    for label, items in by_label.items():
        # Priority: 강함 first (already sorted)
        capped = items[:cap]
        selected.extend(capped)
        log.info("71631: %s — %d available, %d selected", label, len(items), len(capped))

    # Step 4: Group by wav filename for efficient audio loading
    by_wav: dict[str, list[dict]] = defaultdict(list)
    for utt in selected:
        by_wav[utt["wav_filename"]].append(utt)

    # Step 5: Extract audio segments
    samples: list[dict] = []
    skipped = 0

    with zipfile.ZipFile(audio_zip) as azf:
        wav_lookup: dict[str, str] = {}
        for name in azf.namelist():
            if name.endswith(".wav"):
                stem = Path(name).stem
                wav_lookup[stem] = name

        for wav_filename, utterances in by_wav.items():
            if wav_filename not in wav_lookup:
                log.warning("71631: WAV not found in zip: %s", wav_filename)
                skipped += len(utterances)
                continue

            zip_entry = wav_lookup[wav_filename]
            try:
                with azf.open(zip_entry) as src:
                    audio_bytes = src.read()

                buf = io.BytesIO(audio_bytes)
                waveform, sr = torchaudio.load(buf)
            except Exception as e:
                log.warning("71631: Failed to load %s: %s", wav_filename, e)
                skipped += len(utterances)
                continue

            # Mono
            if waveform.shape[0] > 1:
                waveform = waveform.mean(dim=0, keepdim=True)

            for utt in utterances:
                start_sample = int(utt["start_time"] * sr)
                end_sample = int(utt["end_time"] * sr)

                if end_sample > waveform.shape[1]:
                    end_sample = waveform.shape[1]
                if start_sample >= end_sample:
                    skipped += 1
                    continue

                segment = waveform[:, start_sample:end_sample]

                out_name = f"{wav_filename}_{utt['text_no']}.wav"
                out_path = output_dir / utt["label"] / out_name

                ok = preprocess_audio_from_tensor(segment, sr, out_path)
                if ok:
                    samples.append({
                        "path": str(out_path),
                        "label": utt["label"],
                        "label_idx": LABEL2IDX[utt["label"]],
                        "source": "71631",
                        "speaker_id": utt["speaker_id"],
                        "text": utt.get("text", ""),
                        "conversation_id": utt["conversation_id"],
                        "intensity": utt["intensity"],
                    })
                else:
                    skipped += 1

    log.info("71631: Extracted %d samples, skipped %d", len(samples), skipped)
    return samples


# ---------------------------------------------------------------------------
# Task 4 — RAVDESS Extraction
# ---------------------------------------------------------------------------


def extract_ravdess(ravdess_dir: Path, output_dir: Path) -> list[dict]:
    """Extract and preprocess RAVDESS dataset from manifest.csv."""
    ravdess_dir = Path(ravdess_dir)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    manifest_path = ravdess_dir / "manifest.csv"
    if not manifest_path.exists():
        log.error("RAVDESS manifest not found: %s", manifest_path)
        return []

    samples: list[dict] = []
    skipped = 0

    with open(manifest_path) as f:
        reader = csv.DictReader(f)
        for row in reader:
            clean_path = Path(row["clean_path"])
            emotion_raw = row.get("emotion", "")
            label = map_ravdess_label(emotion_raw)
            if label is None:
                skipped += 1
                continue

            actor_id = int(row["actor_id"])
            out_name = clean_path.name
            out_path = output_dir / label / out_name

            if not clean_path.exists():
                skipped += 1
                continue

            ok = preprocess_audio(clean_path, out_path)
            if ok:
                samples.append({
                    "path": str(out_path),
                    "label": label,
                    "label_idx": LABEL2IDX[label],
                    "source": "ravdess",
                    "actor_id": actor_id,
                    "speaker_id": f"ravdess_{actor_id}",
                    "text": "",  # RAVDESS uses fixed sentences, not useful for text emotion
                })
            else:
                skipped += 1

    log.info("RAVDESS: Extracted %d samples, skipped %d", len(samples), skipped)
    return samples


# ---------------------------------------------------------------------------
# Task 4 — Train/Val Splits
# ---------------------------------------------------------------------------


def speaker_isolated_split(
    samples: list[dict], val_ratio: float = 0.1
) -> tuple[list[dict], list[dict]]:
    """Split by speaker_id — no leakage between train/val."""
    if not samples:
        return [], []

    # Group by speaker
    by_speaker: dict[str, list[dict]] = defaultdict(list)
    for s in samples:
        by_speaker[s["speaker_id"]].append(s)

    speakers = list(by_speaker.keys())
    random.shuffle(speakers)

    total = len(samples)
    target_val = int(total * val_ratio)

    val_samples: list[dict] = []
    val_speakers: set[str] = set()
    for spk in speakers:
        if len(val_samples) >= target_val:
            break
        val_samples.extend(by_speaker[spk])
        val_speakers.add(spk)

    train_samples = [s for s in samples if s["speaker_id"] not in val_speakers]
    return train_samples, val_samples


def conversation_isolated_split(
    samples: list[dict], val_ratio: float = 0.1
) -> tuple[list[dict], list[dict]]:
    """Split by conversation_id — no leakage between train/val."""
    if not samples:
        return [], []

    by_conv: dict[str, list[dict]] = defaultdict(list)
    for s in samples:
        by_conv[s["conversation_id"]].append(s)

    convs = list(by_conv.keys())
    random.shuffle(convs)

    total = len(samples)
    target_val = int(total * val_ratio)

    val_samples: list[dict] = []
    val_convs: set[str] = set()
    for conv in convs:
        if len(val_samples) >= target_val:
            break
        val_samples.extend(by_conv[conv])
        val_convs.add(conv)

    train_samples = [s for s in samples if s["conversation_id"] not in val_convs]
    return train_samples, val_samples


def actor_split_ravdess(
    samples: list[dict], val_actors: list[int]
) -> tuple[list[dict], list[dict]]:
    """Split RAVDESS by actor — specified actors go to val."""
    val_set = set(val_actors)
    train = [s for s in samples if s["actor_id"] not in val_set]
    val = [s for s in samples if s["actor_id"] in val_set]
    return train, val


# ---------------------------------------------------------------------------
# Task 4 — Manifest & Stats
# ---------------------------------------------------------------------------


def save_manifest(samples: list[dict], path: Path) -> None:
    """Save manifest as JSON Lines."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w") as f:
        json.dump(samples, f, indent=2, ensure_ascii=False)
    log.info("Saved manifest: %s (%d samples)", path, len(samples))


def save_stats(train: list[dict], val: list[dict], path: Path) -> None:
    """Save dataset statistics."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    def _count_stats(samples: list[dict]) -> dict:
        by_label: Counter = Counter()
        by_source: Counter = Counter()
        for s in samples:
            by_label[s["label"]] += 1
            by_source[s["source"]] += 1
        return {
            "total": len(samples),
            "by_label": dict(sorted(by_label.items())),
            "by_source": dict(sorted(by_source.items())),
        }

    stats = {
        "train": _count_stats(train),
        "val": _count_stats(val),
        "label2idx": LABEL2IDX,
    }

    with open(path, "w") as f:
        json.dump(stats, f, indent=2, ensure_ascii=False)
    log.info("Saved stats: %s", path)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main():
    parser = argparse.ArgumentParser(
        description="Prepare unified LoRA 7-class dataset"
    )
    parser.add_argument(
        "--anchor-dir",
        type=Path,
        default=Path("data/AI Hub 263"),
        help="Path to AI Hub 263 directory with CSVs + ZIPs",
    )
    parser.add_argument(
        "--booster-label-zip",
        type=Path,
        default=Path(
            "data/AI Hub 71631/01-1.정식개방데이터/Training/"
            "02.라벨링데이터/TL_02.실외.zip"
        ),
        help="Path to 71631 label ZIP",
    )
    parser.add_argument(
        "--booster-audio-zip",
        type=Path,
        default=Path(
            "data/AI Hub 71631/01-1.정식개방데이터/Training/"
            "01.원천데이터/TS_02.실외.zip"
        ),
        help="Path to 71631 audio ZIP",
    )
    parser.add_argument(
        "--ravdess-dir",
        type=Path,
        default=Path("data/ravdess"),
        help="Path to RAVDESS directory with manifest.csv",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("data/lora_7class"),
        help="Output directory for processed dataset",
    )
    parser.add_argument("--cap-263", type=int, default=2100, help="Cap per class for 263")
    parser.add_argument("--cap-71631", type=int, default=3500, help="Cap per class for 71631")
    parser.add_argument("--max-per-speaker-71631", type=int, default=20, help="Max utterances per speaker for 71631")
    parser.add_argument("--val-ratio", type=float, default=0.1, help="Val ratio for splits")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    parser.add_argument(
        "--skip-263", action="store_true", help="Skip AI Hub 263 extraction"
    )
    parser.add_argument(
        "--skip-71631", action="store_true", help="Skip AI Hub 71631 extraction"
    )
    parser.add_argument(
        "--skip-ravdess", action="store_true", help="Skip RAVDESS extraction"
    )
    args = parser.parse_args()

    random.seed(args.seed)

    output_dir = args.output_dir
    output_dir.mkdir(parents=True, exist_ok=True)

    all_train: list[dict] = []
    all_val: list[dict] = []

    # ---- 263 Anchor ----
    if not args.skip_263:
        log.info("=" * 60)
        log.info("Extracting AI Hub 263 (anchor)")
        samples_263 = extract_anchor_263(
            args.anchor_dir,
            output_dir / "263",
            cap_per_class=args.cap_263,
        )
        train_263, val_263 = speaker_isolated_split(samples_263, args.val_ratio)
        log.info("263: train=%d, val=%d", len(train_263), len(val_263))
        all_train.extend(train_263)
        all_val.extend(val_263)

    # ---- 71631 Booster ----
    if not args.skip_71631:
        log.info("=" * 60)
        log.info("Extracting AI Hub 71631 (booster)")
        samples_71631 = extract_booster_71631(
            args.booster_label_zip,
            args.booster_audio_zip,
            output_dir / "71631",
            cap=args.cap_71631,
            max_per_speaker=args.max_per_speaker_71631,
        )
        train_71631, val_71631 = conversation_isolated_split(
            samples_71631, args.val_ratio
        )
        log.info("71631: train=%d, val=%d", len(train_71631), len(val_71631))
        all_train.extend(train_71631)
        all_val.extend(val_71631)

    # ---- RAVDESS ----
    if not args.skip_ravdess:
        log.info("=" * 60)
        log.info("Extracting RAVDESS")
        samples_ravdess = extract_ravdess(
            args.ravdess_dir,
            output_dir / "ravdess",
        )
        val_actors = [21, 22, 23, 24]
        train_ravdess, val_ravdess = actor_split_ravdess(
            samples_ravdess, val_actors
        )
        log.info("RAVDESS: train=%d, val=%d", len(train_ravdess), len(val_ravdess))
        all_train.extend(train_ravdess)
        all_val.extend(val_ravdess)

    # ---- Save ----
    log.info("=" * 60)
    log.info("Total: train=%d, val=%d", len(all_train), len(all_val))

    save_manifest(all_train, output_dir / "train_manifest.json")
    save_manifest(all_val, output_dir / "val_manifest.json")
    save_stats(all_train, all_val, output_dir / "stats.json")

    log.info("Done! Output: %s", output_dir)


if __name__ == "__main__":
    main()