searchi/tools/speech-eval/run.py

"""Isolated actual whisper.cpp CPU diagnostic; never downloads or builds tools."""
import argparse
import hashlib
import json
import math
import os
from pathlib import Path
import signal
import subprocess
import tempfile
import time
import wave

UPSTREAM = "d1be6fde11ac6e0407606b4e42fe72d34add8037"
MODELS = {
    "base": (147951465, "60ed5bc3dd14eea856493d334349b405782ddcaf0028d4b5df4088345fba2efe"),
    "small": (487601967, "1be3a9b2063867b937e64e2ec7483364a79917e157fa98c5d94b5c1fffea987b"),
}
ASSETS = ("clip04", "music01", "music02", "sound01", "sound02")
FIXTURE_SHA = "b4e465a0eb9582b826eebd2a20797e1680e4907e336d02d0ef2194e2ba0af393"
MAX_OUTPUT = 4 * 1024 * 1024
# Capture before any expensive work; end-of-run reads can describe another file.
HARNESS_SHA256_AT_START = hashlib.sha256(Path(__file__).read_bytes()).hexdigest()


def sha256(path):
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def stop(child):
    if os.name == "nt":
        subprocess.run(["taskkill", "/PID", str(child.pid), "/T", "/F"],
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                       creationflags=subprocess.CREATE_NO_WINDOW, check=False)
    else:
        os.killpg(child.pid, signal.SIGKILL)
    child.wait()


def run_owned(command, scratch, name, outputs, timeout):
    """Bound time/output; terminate only our process tree on failure/interruption."""
    log = scratch / f"{name}.log"
    kwargs = {"creationflags": subprocess.CREATE_NO_WINDOW} if os.name == "nt" else {"start_new_session": True}
    started = time.perf_counter()
    with log.open("wb") as stream:
        child = subprocess.Popen(command, stdin=subprocess.DEVNULL,
                                 stdout=stream, stderr=subprocess.STDOUT, **kwargs)
        try:
            while child.poll() is None:
                if time.perf_counter() - started > timeout:
                    raise TimeoutError(f"{name} exceeded {timeout}s")
                if any(path.exists() and path.stat().st_size > limit for path, limit in [(log, MAX_OUTPUT), *outputs]):
                    raise ValueError(f"{name} exceeded output budget")
                time.sleep(0.05)
            if child.returncode:
                raise RuntimeError(f"{name} failed with exit {child.returncode}; owned diagnostic log: {log}")
            if any(path.exists() and path.stat().st_size > limit for path, limit in [(log, MAX_OUTPUT), *outputs]):
                raise ValueError(f"{name} exceeded output budget")
        except BaseException as error:
            if child.poll() is None:
                stop(child)
            stream.flush()
            # Preserve bounded actual diagnostic bytes outside ephemeral scratch.
            with tempfile.NamedTemporaryFile(prefix=f"speech-{name}-failure-", suffix=".log", dir=scratch.parent, delete=False) as retained:
                with log.open("rb") as source:
                    retained.write(source.read(MAX_OUTPUT))
                print(f"Owned failure diagnostic retained: {retained.name}")
                diagnostic = {"stage": name, "elapsed_ms": (time.perf_counter() - started) * 1000,
                              "exit_code": child.returncode, "owned_pid": child.pid, "error_type": type(error).__name__,
                              "bounded_log": retained.name, "complete": False}
            Path(retained.name).with_suffix(".json").write_text(json.dumps(diagnostic, indent=2) + "\n", encoding="utf-8")
            raise
        finally:
            if child.poll() is None:
                stop(child)
    return (time.perf_counter() - started) * 1000


def segments(document, offset_ms, duration_ms):
    raw = document.get("transcription")
    if not isinstance(raw, list) or len(raw) > 1000:
        raise ValueError("Missing or oversized transcription array")
    parsed = []
    previous = 0
    for row in raw:
        start, end = row["offsets"]["from"], row["offsets"]["to"]
        text = row["text"]
        if type(start) is not int or type(end) is not int or not 0 <= start <= end <= duration_ms or start < previous:
            raise ValueError("Invalid local ASR timestamps")
        if not isinstance(text, str) or len(text.encode("utf-8")) > 4096:
            raise ValueError("Invalid recognized text")
        previous = start
        if text.strip():
            parsed.append({"start_ms": offset_ms + start, "end_ms": offset_ms + end, "text": text})
    return parsed


def recognized_window(raw, offset, length, decode_ms, asr_ms):
    report = {"offset_ms": offset, "duration_ms": length, "decode_ms": decode_ms,
              "asr_process_ms": asr_ms, "actual_whisper_json": raw}
    try:
        report["segments"] = segments(raw, offset, length)
    except (ValueError, KeyError, TypeError) as error:
        report["segments"] = []
        report["validation_error"] = str(error)
    return report


def decode_command(ffmpeg, source, offset, length, wav, audio_index=None):
    return [str(ffmpeg.resolve()), "-nostdin", "-hide_banner", "-v", "error",
            "-protocol_whitelist", "file,pipe", "-threads", "1", "-seek_timestamp", "1", "-ss", str(offset / 1000), "-t", str(length / 1000),
            "-i", str(source.resolve()), "-map", f"0:{audio_index}" if audio_index is not None else "0:a:0", "-vn", "-ac", "1", "-ar", "16000",
            "-af", f"aresample=16000,atrim=end_sample={length * 16},asetpts=PTS-STARTPTS",
            "-c:a", "pcm_s16le", "-threads", "1", "-y", str(wav)]


def decoded_frames(wav, length):
    with wave.open(str(wav), "rb") as audio:
        if audio.getframerate() != 16000 or audio.getnchannels() != 1 or audio.getsampwidth() != 2:
            raise ValueError("Invalid decoded PCM")
        frames = audio.getnframes()
    # Container duration rounding allows at most one millisecond's shortfall.
    if not max(1, length * 16 - 16) <= frames <= length * 16:
        raise ValueError(f"Decoded sample coverage mismatch: expected {length * 16}, got {frames}")
    return frames


def primary_audio_range(probe):
    audio = next((stream for stream in probe.get("streams", []) if stream.get("codec_type") == "audio"), None)
    if audio is None:
        raise ValueError("No primary audio stream")
    start = float(audio.get("start_time", 0))
    if "duration" not in audio:
        raise ValueError("Primary audio stream duration unavailable; proof requires explicit stream duration, not container fallback")
    duration = float(audio["duration"])
    if not math.isfinite(start) or not math.isfinite(duration) or start < 0 or not 0 < duration <= 600:
        raise ValueError("Proof audio stream range exceeds bounds")
    return math.ceil(start * 1000), math.ceil((start + duration) * 1000), audio["index"]


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--whisper", type=Path, required=True)
    parser.add_argument("--whisper-sha256", required=True, help="Previously recorded pinned build SHA, not computed for acceptance here")
    parser.add_argument("--ffmpeg", type=Path, required=True)
    parser.add_argument("--ffprobe", type=Path, required=True)
    parser.add_argument("--model", type=Path, required=True)
    parser.add_argument("--candidate", choices=MODELS, required=True)
    parser.add_argument("--token-sentences", action="store_true", help="Frozen ordinary token-time diagnostic; retains original segments")
    parser.add_argument("--dtw-sentences", action="store_true", help="Separate frozen DTW point-span diagnostic")
    parser.add_argument("--fixture", type=Path, default=Path("fixtures/media/fixture.json"))
    parser.add_argument("--scratch-parent", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    if args.output.exists():
        raise ValueError("Never overwrite experiment evidence")
    if args.token_sentences and args.dtw_sentences:
        raise ValueError("Select one separate alignment experiment")
    alignment_mode = args.token_sentences or args.dtw_sentences
    alignment_file = "dtw_alignment.py" if args.dtw_sentences else "token_alignment.py"
    policy_file = "dtw-policy.json" if args.dtw_sentences else "token-policy.json"
    if alignment_mode and args.candidate != "base":
        raise ValueError("Frozen token experiment accepts multilingual base only")
    if alignment_mode:
        policy = json.loads(Path(__file__).with_name(policy_file).read_text(encoding="utf-8"))
        if policy["worker_sha256"] != args.whisper_sha256:
            raise ValueError("Frozen token experiment worker mismatch")
        alignment_contract_sha = sha256(Path(__file__).with_name(alignment_file))
        alignment_policy_sha = sha256(Path(__file__).with_name(policy_file))
        # Load helpers now; future file edits cannot change a later lazy import.
        if args.dtw_sentences:
            from dtw_alignment import aligned_segments
            from token_alignment import aligned_segments as ordinary_segments
        else:
            from token_alignment import aligned_segments
    size, digest = MODELS[args.candidate]
    if args.model.stat().st_size != size or sha256(args.model) != digest:
        raise ValueError("Model size/SHA verification failed")
    if sha256(args.whisper) != args.whisper_sha256:
        raise ValueError("Pinned worker build SHA verification failed")
    if sha256(args.fixture) != FIXTURE_SHA:
        raise ValueError("Frozen media-v1 fixture SHA mismatch")
    fixture = json.loads(args.fixture.read_text(encoding="utf-8"))
    if not set(ASSETS).issubset({asset["id"] for asset in fixture["assets"]}):
        raise ValueError("Missing proof assets")
    rows = []
    args.scratch_parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix="speech-proof-", dir=args.scratch_parent) as directory:
        scratch = Path(directory)
        for asset in fixture["assets"]:
            if asset["id"] not in ASSETS:
                continue
            source = args.fixture.parent / asset["path"]
            if source.stat().st_size != asset["bytes"] or sha256(source) != asset["sha256"]:
                raise ValueError(f"Frozen asset mismatch: {asset['id']}")
            # Probe stdout is bounded by the same owned log, then parsed directly.
            run_owned([str(args.ffprobe.resolve()), "-v", "error", "-protocol_whitelist", "file,pipe",
                       "-show_entries", "stream=index,codec_type,duration,start_time:format=duration,start_time", "-of", "json", str(source.resolve())],
                      scratch, "probe", [], 30)
            probe = json.loads((scratch / "probe.log").read_text(encoding="utf-8"))
            start_ms, end_ms, audio_index = primary_audio_range(probe)
            result = {"asset": asset["id"], "source_sha256": asset["sha256"], "duration_ms": end_ms - start_ms,
                      "audio_start_ms": start_ms, "audio_end_ms": end_ms, "audio_stream_index": audio_index,
                      "actual_probe": probe, "windows": []}
            for offset in range(start_ms, end_ms, 30000):
                length = min(30000, end_ms - offset)
                wav = scratch / "window.wav"
                prefix = scratch / "recognized"
                output = prefix.with_suffix(".json")
                output.unlink(missing_ok=True)
                decode_ms = run_owned(decode_command(args.ffmpeg, source, offset, length, wav, audio_index),
                    scratch, "decode", [(wav, 970000)], 60)
                try:
                    frames = decoded_frames(wav, length)
                except ValueError as error:
                    with wave.open(str(wav), "rb") as audio:
                        actual_frames = audio.getnframes()
                    result["windows"].append({"offset_ms": offset, "duration_ms": length, "decode_ms": decode_ms,
                        "decoded_frames": actual_frames, "validation_error": str(error), "error_stage": "decode_coverage",
                        "segments": [], "asr_not_run": True})
                else:
                    asr_ms = run_owned([str(args.whisper.resolve()), "--no-gpu", "--threads", "1", "--processors", "1",
                        "--language", "auto", "--model", str(args.model.resolve()), "--file", str(wav),
                        "--output-json-full" if alignment_mode else "--output-json",
                        *(["--dtw", "base", "--no-flash-attn"] if args.dtw_sentences else []),
                        "--output-file", str(prefix)], scratch, "asr", [(output, MAX_OUTPUT)], 600)
                    raw = json.loads(output.read_text(encoding="utf-8"))
                    result["windows"].append(recognized_window(raw, offset, length, decode_ms, asr_ms))
                    result["windows"][-1]["decoded_frames"] = frames
                    if alignment_mode:
                        if args.dtw_sentences:
                            try:
                                ordinary_segments(raw, offset, length)
                            except (ValueError, KeyError, TypeError) as error:
                                result["windows"][-1]["ordinary_alignment_validation_error"] = str(error)
                        try:
                            result["windows"][-1]["aligned_segments"] = aligned_segments(raw, offset, length)
                        except (ValueError, KeyError, TypeError) as error:
                            result["windows"][-1]["alignment_validation_error"] = str(error)
                args.output.parent.mkdir(parents=True, exist_ok=True)
                checkpoint = {"complete": False, "fixture_sha256": sha256(args.fixture),
                              "harness_sha256": HARNESS_SHA256_AT_START,
                              "model_sha256": digest, "worker_sha256": args.whisper_sha256,
                              "rows": [*rows, result]}
                if alignment_mode:
                    checkpoint["alignment_policy"] = policy["policy"]
                    checkpoint["alignment_policy_sha256"] = alignment_policy_sha
                    checkpoint["alignment_contract_sha256"] = alignment_contract_sha
                args.output.write_text(json.dumps(checkpoint, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
            if source.stat().st_size != asset["bytes"] or sha256(source) != asset["sha256"]:
                raise ValueError("Source changed during recognition")
            rows.append(result)
    failures = [{"asset": row["asset"], "offset_ms": window["offset_ms"], "error": window["validation_error"]}
                for row in rows for window in row["windows"] if "validation_error" in window]
    if alignment_mode:
        failures.extend({"asset": row["asset"], "offset_ms": window["offset_ms"],
                         "error": window["alignment_validation_error"]}
                        for row in rows for window in row["windows"] if "alignment_validation_error" in window)
    report = {"version": 1, "complete": True, "quality_failed": bool(failures), "validation_failures": failures,
              "scope": "Actual ASR diagnostic, not Google retrieval or production integration",
              "fixture_sha256": sha256(args.fixture), "upstream_commit": UPSTREAM,
              "worker_sha256": args.whisper_sha256, "model_sha256": digest, "candidate": args.candidate,
              "ffmpeg_sha256": sha256(args.ffmpeg), "ffprobe_sha256": sha256(args.ffprobe),
              "backend": "forced CPU", "threads": 1,
              "harness_sha256": HARNESS_SHA256_AT_START,
              "harness_file_changed_during_run": sha256(Path(__file__)) != HARNESS_SHA256_AT_START,
              "policy": "primary-audio-range-v3;aresample16000-before-atrim;samplecoverage-1ms;nonoverlapping30s;independent-process-per-window;auto-language;segment-times",
              "rows": rows}
    if alignment_mode:
        report["alignment_policy"] = policy["policy"]
        report["alignment_contract_sha256"] = alignment_contract_sha
        report["alignment_policy_sha256"] = alignment_policy_sha
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()