eval_ab.py

"""A/B evaluator for Vocence miner checkpoints using subnet-like scoring.

Compares two model sources (A and B) on a JSONL prompt set and reports:
- mean score
- pass rate at threshold (default 0.9)
- head-to-head wins (A>B, B>A, ties)

The scorer follows Vocence's element weighting and trait rules used in subnet78.
"""
from __future__ import annotations

import argparse
import asyncio
import json
import math
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import soundfile as sf
from audiojudge import AudioJudge
from qwen_tts import Qwen3TTSModel


VOICE_TRAIT_ENUMS: dict[str, list[str]] = {
    "gender": ["male", "female", "neutral"],
    "pitch": ["low", "mid", "high"],
    "speed": ["slow", "normal", "fast"],
    "age_group": ["child", "young_adult", "adult", "senior"],
    "emotion": ["neutral", "happy", "sad", "angry", "calm", "excited", "serious", "fearful"],
    "tone": ["warm", "cold", "friendly", "formal", "casual", "authoritative"],
    "accent": ["us", "uk", "au", "in", "neutral", "other"],
}
ORDINAL_TRAITS = {"pitch", "speed", "age_group"}
WEIGHTS = {
    "script": 0.30,
    "naturalness": 0.15,
    "gender": 0.10,
    "speed": 0.10,
    "emotion": 0.10,
    "age_group": 0.10,
    "pitch": 0.05,
    "accent": 0.05,
    "tone": 0.05,
}
PASS_THRESHOLD = 0.9

DESCRIPTION_SYSTEM = """You are an expert at analyzing speech for text-to-speech evaluation.
Analyze the audio and return JSON with these keys and exact enum values:
- transcription (string)
- gender: male|female|neutral
- pitch: low|mid|high
- speed: slow|normal|fast
- age_group: child|young_adult|adult|senior
- emotion: neutral|happy|sad|angry|calm|excited|serious|fearful
- tone: warm|cold|friendly|formal|casual|authoritative
- accent: us|uk|au|in|neutral|other
Return ONLY JSON."""


@dataclass
class Sample:
    text: str
    instruction: str
    source_audio: str | None = None


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Evaluate two Vocence miners A/B")
    p.add_argument("--a-source", required=True, help="Model A source (HF repo id or local path)")
    p.add_argument("--b-source", required=True, help="Model B source (HF repo id or local path)")
    p.add_argument("--dataset", required=True, help="JSONL file with {text, instruction[, source_audio]}")
    p.add_argument("--openai-key", default="", help="OpenAI API key (or use OPENAI_API_KEY env)")
    p.add_argument("--model", default="gpt-4o-audio-preview", help="Judge model")
    p.add_argument("--limit", type=int, default=0, help="Max rows from dataset (0 = all)")
    p.add_argument("--device", default="cuda:0", help="Torch device map for Qwen models")
    p.add_argument("--precision", default="bfloat16", choices=("bfloat16", "float16", "float32"))
    p.add_argument("--pass-threshold", type=float, default=PASS_THRESHOLD)
    p.add_argument("--no-naturalness", action="store_true", help="Disable naturalness element")
    return p.parse_args()


def _dtype_from_str(name: str):
    import torch

    return {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[name]


def load_samples(path: Path, limit: int) -> list[Sample]:
    rows: list[Sample] = []
    with path.open("r", encoding="utf-8") as fh:
        for line in fh:
            raw = line.strip()
            if not raw:
                continue
            item = json.loads(raw)
            rows.append(
                Sample(
                    text=str(item["text"]),
                    instruction=str(item["instruction"]),
                    source_audio=str(item["source_audio"]) if item.get("source_audio") else None,
                )
            )
            if limit > 0 and len(rows) >= limit:
                break
    return rows


def _tokenize(text: str) -> list[str]:
    import re

    return re.findall(r"\w+", (text or "").lower())


def word_error_rate(reference: str, hypothesis: str) -> float:
    ref = _tokenize(reference)
    hyp = _tokenize(hypothesis)
    if not ref:
        return 1.0 if hyp else 0.0
    n, m = len(ref), len(hyp)
    prev = list(range(m + 1))
    for i in range(1, n + 1):
        curr = [i] + [0] * m
        for j in range(1, m + 1):
            cost = 0 if ref[i - 1] == hyp[j - 1] else 1
            curr[j] = min(prev[j] + 1, curr[j - 1] + 1, prev[j - 1] + cost)
        prev = curr
    return min(1.0, prev[m] / n)


def score_element(key: str, expected: Any, actual: Any) -> float:
    if key == "script":
        return max(0.0, 1.0 - word_error_rate(str(expected or ""), str(actual or "")))
    enum = VOICE_TRAIT_ENUMS.get(key)
    if not enum:
        return 1.0 if expected == actual else 0.0
    if key in ORDINAL_TRAITS:
        try:
            i = enum.index(str(expected))
            j = enum.index(str(actual))
        except ValueError:
            return 0.0
        dist = abs(i - j)
        return 1.0 if dist == 0 else (0.5 if dist == 1 else 0.0)
    return 1.0 if expected == actual else 0.0


def parse_traits(raw: str) -> dict[str, Any]:
    text = (raw or "").strip()
    try:
        parsed = json.loads(text)
    except json.JSONDecodeError:
        parsed = {}
    out: dict[str, Any] = {"transcription": str(parsed.get("transcription") or "").strip()}
    for k, enum in VOICE_TRAIT_ENUMS.items():
        v = str(parsed.get(k) or "").strip().lower().replace(" ", "_").replace("-", "_")
        out[k] = v if v in enum else enum[0]
    return out


def extract_traits(judge: AudioJudge, audio_path: str, model_name: str) -> dict[str, Any]:
    result = judge.judge_audio_pointwise(
        audio_path=audio_path,
        system_prompt=DESCRIPTION_SYSTEM,
        user_prompt=None,
        model=model_name,
        concatenation_method="no_concatenation",
        temperature=0.0,
        max_tokens=500,
    )
    if not result.get("success"):
        return {
            "transcription": "",
            "gender": "neutral",
            "pitch": "mid",
            "speed": "normal",
            "age_group": "adult",
            "emotion": "neutral",
            "tone": "casual",
            "accent": "neutral",
        }
    return parse_traits(result.get("response") or "")


def naturalness_win(judge: AudioJudge, source_audio: str, generated_audio: str, model_name: str, task: str) -> bool:
    prompt = (
        "You are an audio naturalness judge. You will hear two clips for the same TTS task.\n"
        f"Task: {task}\n"
        "Which sounds more natural as human speech? Reply FIRST or SECOND only."
    )
    result = judge.judge_audio(
        audio1_path=source_audio,
        audio2_path=generated_audio,
        system_prompt=prompt,
        user_prompt=None,
        model=model_name,
        concatenation_method="no_concatenation",
        temperature=0.0,
        max_tokens=40,
    )
    if not result.get("success"):
        return False
    first_line = (result.get("response") or "").strip().split("\n", 1)[0].strip().upper()
    return "SECOND" in first_line


def compute_score(source_traits: dict[str, Any], miner_traits: dict[str, Any], naturalness: bool | None) -> float:
    weight_sum = 0.0
    weighted = 0.0
    for key, w in WEIGHTS.items():
        if key == "naturalness":
            if naturalness is None:
                continue
            s = 1.0 if naturalness else 0.0
        else:
            sk = "transcription" if key == "script" else key
            s = score_element(key, source_traits.get(sk), miner_traits.get(sk))
        weighted += w * s
        weight_sum += w
    return weighted / weight_sum if weight_sum else 0.0


def save_wav(path: Path, wave: Any, sr: int) -> None:
    sf.write(str(path), wave, sr)


def run_model_once(model: Qwen3TTSModel, text: str, instruction: str) -> tuple[Any, int]:
    wavs, sr = model.generate_voice_design(text=text, instruct=instruction, language="english")
    if not wavs:
        raise RuntimeError("Model returned empty wave list")
    return wavs[0], int(sr)


async def main() -> int:
    args = parse_args()
    dataset = load_samples(Path(args.dataset), args.limit)
    if not dataset:
        raise SystemExit("Dataset is empty")

    openai_key = args.openai_key or __import__("os").environ.get("OPENAI_API_KEY", "")
    if not openai_key:
        raise SystemExit("Missing OpenAI key. Provide --openai-key or OPENAI_API_KEY.")

    dtype = _dtype_from_str(args.precision)
    model_a = Qwen3TTSModel.from_pretrained(args.a_source, device_map=args.device, dtype=dtype)
    model_b = Qwen3TTSModel.from_pretrained(args.b_source, device_map=args.device, dtype=dtype)
    judge = AudioJudge(openai_api_key=openai_key, google_api_key=None)

    a_scores: list[float] = []
    b_scores: list[float] = []
    a_pass = b_pass = 0
    a_wins = b_wins = ties = 0

    with tempfile.TemporaryDirectory(prefix="vocence_eval_") as tmp:
        tdir = Path(tmp)
        for idx, row in enumerate(dataset, start=1):
            aw, asr = await asyncio.to_thread(run_model_once, model_a, row.text, row.instruction)
            bw, bsr = await asyncio.to_thread(run_model_once, model_b, row.text, row.instruction)
            a_path = tdir / f"a_{idx}.wav"
            b_path = tdir / f"b_{idx}.wav"
            await asyncio.to_thread(save_wav, a_path, aw, asr)
            await asyncio.to_thread(save_wav, b_path, bw, bsr)

            if row.source_audio:
                src_traits = await asyncio.to_thread(extract_traits, judge, row.source_audio, args.model)
            else:
                # Fallback: use prompt text + parsed instruction as pseudo-spec.
                src_traits = {
                    "transcription": row.text,
                    "gender": "neutral",
                    "pitch": "mid",
                    "speed": "normal",
                    "age_group": "adult",
                    "emotion": "neutral",
                    "tone": "casual",
                    "accent": "neutral",
                }
                for chunk in row.instruction.split("|"):
                    if ":" not in chunk:
                        continue
                    k, v = chunk.split(":", 1)
                    key = k.strip().lower()
                    val = v.strip().lower().replace(" ", "_")
                    if key in VOICE_TRAIT_ENUMS and val in VOICE_TRAIT_ENUMS[key]:
                        src_traits[key] = val

            a_traits = await asyncio.to_thread(extract_traits, judge, str(a_path), args.model)
            b_traits = await asyncio.to_thread(extract_traits, judge, str(b_path), args.model)

            nat_a = nat_b = None
            if row.source_audio and not args.no_naturalness:
                nat_a = await asyncio.to_thread(
                    naturalness_win, judge, row.source_audio, str(a_path), args.model, row.instruction
                )
                nat_b = await asyncio.to_thread(
                    naturalness_win, judge, row.source_audio, str(b_path), args.model, row.instruction
                )

            sa = compute_score(src_traits, a_traits, nat_a)
            sb = compute_score(src_traits, b_traits, nat_b)
            a_scores.append(sa)
            b_scores.append(sb)
            a_pass += int(sa >= args.pass_threshold)
            b_pass += int(sb >= args.pass_threshold)

            if sa > sb + 1e-6:
                a_wins += 1
            elif sb > sa + 1e-6:
                b_wins += 1
            else:
                ties += 1

            print(f"[{idx}/{len(dataset)}] A={sa:.4f} B={sb:.4f} passA={sa >= args.pass_threshold} passB={sb >= args.pass_threshold}")

    n = len(a_scores)
    mean_a = sum(a_scores) / n
    mean_b = sum(b_scores) / n
    std_a = math.sqrt(sum((x - mean_a) ** 2 for x in a_scores) / n)
    std_b = math.sqrt(sum((x - mean_b) ** 2 for x in b_scores) / n)

    print("\n=== Vocence A/B Report ===")
    print(f"Samples: {n}")
    print(f"A mean score: {mean_a:.4f} (std {std_a:.4f})")
    print(f"B mean score: {mean_b:.4f} (std {std_b:.4f})")
    print(f"A pass rate @ {args.pass_threshold:.2f}: {a_pass}/{n} = {a_pass / n:.1%}")
    print(f"B pass rate @ {args.pass_threshold:.2f}: {b_pass}/{n} = {b_pass / n:.1%}")
    print(f"Head-to-head: A wins {a_wins}, B wins {b_wins}, ties {ties}")

    return 0


if __name__ == "__main__":
    raise SystemExit(asyncio.run(main()))