imagefind/scripts/audio_small_ab.py

from __future__ import annotations

import argparse
import json
import subprocess
import sys
import tempfile
import time
import wave
from pathlib import Path


COMBINATIONS = (
    ("small_vad_gpu_1beam", "small", "GPU", "vad"),
    ("small_context_gpu_1beam", "small", "GPU", "continuous"),
    ("small_context_cpu_5beam", "small", "CPU", "continuous"),
    ("medium_current_baseline", "medium", "GPU", "vad"),
)


def arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Read-only Small/Medium audio A/B diagnostics")
    parser.add_argument("--data-dir", required=True, type=Path)
    parser.add_argument("--media", required=True, type=Path, action="append")
    parser.add_argument("--windows", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--ffmpeg", default="ffmpeg")
    parser.add_argument("--force", action="store_true")
    return parser.parse_args()


def _write_window_batch(source: Path, destination: Path, windows: list[dict]) -> int:
    """Concatenate equal-sized diagnostic windows without reloading Whisper."""
    with wave.open(str(source), "rb") as handle:
        params = handle.getparams()
        sample_rate = handle.getframerate()
        frame_count = handle.getnframes()
        durations = {int(item["end_ms"]) - int(item["start_ms"]) for item in windows}
        if len(durations) != 1:
            raise ValueError("all diagnostic windows for one media item must have equal duration")
        duration_ms = durations.pop()
        frames_per_window = duration_ms * sample_rate // 1000
        with wave.open(str(destination), "wb") as output:
            output.setparams(params)
            for item in windows:
                start_frame = int(item["start_ms"]) * sample_rate // 1000
                handle.setpos(min(start_frame, frame_count))
                raw = handle.readframes(frames_per_window)
                expected = frames_per_window * params.nchannels * params.sampwidth
                if len(raw) < expected:
                    raw += b"\0" * (expected - len(raw))
                output.writeframesraw(raw)
    return duration_ms


def worker_run(data_dir: Path, wav: Path, windows: list[dict], combination: tuple[str, ...]) -> list[dict]:
    name, model, device, segmentation = combination
    with wave.open(str(wav), "rb") as handle:
        duration_ms = round(handle.getnframes() / handle.getframerate() * 1000 / len(windows))
    command = [
        sys.executable, "-m", "imagefind.audio_worker",
        "--data-dir", str(data_dir), "--wav", str(wav),
        "--device", device, "--model-variant", model,
        "--segmentation", segmentation, "--quality-profile", "accuracy",
        "--language-policy", str(windows[0].get("language_policy", "zh_priority")),
        "--chunk-seconds", str(max(15, min(30, (duration_ms + 999) // 1000))),
        "--overlap-seconds", "0",
        "--start-frame", "0",
        "--max-chunks", str(len(windows)), "--cpu-threads", "2",
    ]
    started = time.monotonic()
    process = subprocess.run(command, text=True, capture_output=True, encoding="utf-8", errors="replace")
    elapsed = time.monotonic() - started
    events = []
    for line in process.stdout.splitlines():
        try:
            events.append(json.loads(line))
        except json.JSONDecodeError:
            pass
    results = []
    for index, window in enumerate(windows, 1):
        selected = [event for event in events if event.get("chunk_index") == index]
        if index == len(windows):
            selected.extend(event for event in events if event.get("event") in {"complete", "error"})
        results.append({
            "combination": name, "model": model, "device": device,
            "beam": 1 if device == "GPU" else 5, "segmentation": segmentation,
            "elapsed_seconds": round(elapsed, 3), "window_seconds": round(duration_ms / 1000, 3),
            "batch_window_count": len(windows),
            "batch_rtf": round(elapsed / max(0.001, duration_ms * len(windows) / 1000), 3),
            "exit_code": process.returncode, "events": selected,
            "stderr_tail": process.stderr[-2000:] if process.returncode else "", "window": window,
        })
    return results


def main() -> int:
    args = arguments()
    if args.output.exists() and not args.force:
        raise SystemExit("output exists; pass --force to replace this diagnostic artifact")
    windows = json.loads(args.windows.read_text(encoding="utf-8"))
    artifact = {
        "schema_version": 1,
        "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "database_writes": False,
        "runs": [],
    }
    with tempfile.TemporaryDirectory(prefix="imagefind-audio-ab-") as work:
        for media in args.media:
            wav = Path(work) / f"{media.stem}.wav"
            subprocess.run(
                [
                    args.ffmpeg, "-nostdin", "-y", "-i", str(media), "-vn",
                    "-ac", "1", "-ar", "16000", "-c:a", "pcm_s16le", str(wav),
                ],
                check=True,
                capture_output=True,
            )
            selected = windows.get(media.name)
            if not isinstance(selected, list) or len(selected) != 8:
                raise SystemExit(f"{media.name}: windows JSON must contain exactly 8 labelled windows")
            batched_wav = Path(work) / f"{media.stem}-windows.wav"
            _write_window_batch(wav, batched_wav, selected)
            for combination in COMBINATIONS:
                for result in worker_run(args.data_dir, batched_wav, selected, combination):
                    result["media"] = media.name
                    artifact["runs"].append(result)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(artifact, ensure_ascii=False, indent=2), encoding="utf-8")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())