110 lines
4.2 KiB
Python
110 lines
4.2 KiB
Python
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 worker_run(data_dir: Path, wav: Path, window: dict, combination: tuple[str, ...]) -> dict:
|
|
name, model, device, segmentation = combination
|
|
with wave.open(str(wav), "rb") as handle:
|
|
sample_rate = handle.getframerate()
|
|
start_ms = int(window["start_ms"])
|
|
duration_ms = int(window["end_ms"]) - start_ms
|
|
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(window.get("language_policy", "zh_priority")),
|
|
"--chunk-seconds", str(max(15, min(30, (duration_ms + 999) // 1000))),
|
|
"--overlap-seconds", "0",
|
|
"--start-frame", str(start_ms * sample_rate // 1000),
|
|
"--max-chunks", "1", "--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
|
|
return {
|
|
"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),
|
|
"rtf": round(elapsed / max(0.001, duration_ms / 1000), 3),
|
|
"exit_code": process.returncode,
|
|
"events": events,
|
|
"stderr_tail": process.stderr[-2000:],
|
|
}
|
|
|
|
|
|
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")
|
|
for window in selected:
|
|
for combination in COMBINATIONS:
|
|
result = worker_run(args.data_dir, wav, window, combination)
|
|
result.update({"media": media.name, "window": window})
|
|
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())
|