scripts/measure_latency.py

#!/usr/bin/env python3
"""Measure batch-1 inference latency for baseline and LatentASR settings."""

from __future__ import annotations

import argparse
import json
import sys
import statistics
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import numpy as np
import torch
from datasets import Audio, load_dataset

from eval import (
    build_base_model_bundle,
    build_latent_bundle,
    choose_device,
    choose_dtype,
    clean_prediction,
    resolve_language_hint,
    release_model,
)
from utils import set_seed


DATASETS = {
    "fleurs": ("google/fleurs", "en_us", "test"),
    "voxpopuli": ("facebook/voxpopuli", "en", "test"),
}


def collect_samples(
    dataset_name: str,
    config: str,
    split: str,
    n: int,
    streaming: bool,
    sampling_rate: int,
) -> List[Dict[str, Any]]:
    ds = load_dataset(
        dataset_name,
        config,
        split=split,
        trust_remote_code=True,
        streaming=streaming,
    )
    if not streaming:
        ds = ds.cast_column("audio", Audio(sampling_rate=sampling_rate))
    samples: List[Dict[str, Any]] = []
    for sample in ds:
        audio = sample.get("audio")
        if isinstance(audio, dict) and "array" in audio and "sampling_rate" in audio:
            samples.append(audio)
        if len(samples) >= n:
            break
    return samples


def run_one(
    model: Any,
    processor: Any,
    audio: Dict[str, Any],
    use_baseline: bool,
    threshold: float,
    language_hint: Optional[str],
    max_new_tokens: int,
) -> Tuple[str, Dict[str, Any]]:
    prompt_text = (
        f"Transcribe the {language_hint} audio into text."
        if language_hint
        else "Transcribe the audio into text."
    )
    target_dtype = model.thinker.dtype if hasattr(model.thinker, "dtype") else torch.float32
    feat_out = processor.feature_extractor(
        np.asarray(audio["array"], dtype=np.float64),
        sampling_rate=audio["sampling_rate"],
        return_attention_mask=True,
    )
    feats = torch.tensor(
        feat_out.input_features[0],
        dtype=target_dtype,
        device=model.base_model.device,
    ).unsqueeze(0)
    n_frames = feats.size(-1)
    raw_mask = getattr(feat_out, "attention_mask", None)
    if raw_mask is not None:
        mask = torch.tensor(raw_mask[0], dtype=torch.long)
        if mask.size(-1) < n_frames:
            mask = torch.cat([mask, torch.zeros(n_frames - mask.size(-1), dtype=torch.long)])
        else:
            mask = mask[:n_frames]
        feature_attention_mask = mask.to(model.base_model.device).unsqueeze(0)
    else:
        feature_attention_mask = torch.ones((1, n_frames), dtype=torch.long, device=model.base_model.device)

    out = model.generate(
        feats,
        feature_attention_mask=feature_attention_mask,
        max_new_tokens=max_new_tokens,
        use_baseline=use_baseline,
        return_thoughts=False,
        return_stats=True,
        do_sample=False,
        eos_token_id=[151645, 151643],
        num_beams=1,
        language_hint=language_hint,
        prompt_text=prompt_text,
        dynamic_halt_threshold=threshold,
    )
    if isinstance(out, tuple):
        gen_ids = out[0]
        stats = out[1] if isinstance(out[1], dict) else {}
    else:
        gen_ids = out
        stats = {}
    text = clean_prediction(processor.tokenizer.decode(gen_ids[0], skip_special_tokens=True))
    return text, stats


def measure_setting(
    model: Any,
    processor: Any,
    samples: List[Dict[str, Any]],
    use_baseline: bool,
    threshold: float,
    language_hint: Optional[str],
    warmup: int,
    measure: int,
    max_new_tokens: int,
) -> Dict[str, Any]:
    n_total = min(len(samples), warmup + measure)
    if n_total <= warmup:
        raise ValueError(f"Need > warmup samples, got {len(samples)}")
    times: List[float] = []
    steps: List[int] = []
    for idx, audio in enumerate(samples[:n_total]):
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        t0 = time.perf_counter()
        _, stats = run_one(
            model=model,
            processor=processor,
            audio=audio,
            use_baseline=use_baseline,
            threshold=threshold,
            language_hint=language_hint,
            max_new_tokens=max_new_tokens,
        )
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        dt_ms = (time.perf_counter() - t0) * 1000.0
        if idx >= warmup:
            times.append(dt_ms)
            iters = stats.get("deq_iters")
            if hasattr(iters, "item"):
                steps.append(int(iters.item()))
    return {
        "mean_ms": statistics.mean(times),
        "median_ms": statistics.median(times),
        "std_ms": statistics.pstdev(times) if len(times) > 1 else 0.0,
        "n": len(times),
        "avg_steps": (sum(steps) / len(steps)) if steps else None,
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model-id", default="Qwen/Qwen3-ASR-0.6B")
    parser.add_argument("--latent-ckpt", default="latent_qwen_asr_best.pth")
    parser.add_argument("--output-json", required=True)
    parser.add_argument("--warmup", type=int, default=10)
    parser.add_argument("--samples", type=int, default=100)
    parser.add_argument("--max-new-tokens", type=int, default=128)
    parser.add_argument("--streaming", action="store_true")
    parser.add_argument("--device", default="auto", choices=["auto", "cuda", "cpu"])
    parser.add_argument("--dtype", default="auto", choices=["auto", "float32", "float16", "bfloat16"])
    args = parser.parse_args()

    set_seed(42)
    device = choose_device(args.device)
    dtype = choose_dtype(args.dtype, device)
    output: Dict[str, Any] = {
        "model_id": args.model_id,
        "latent_ckpt": args.latent_ckpt,
        "device": device,
        "dtype": str(dtype),
        "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
        "warmup": args.warmup,
        "samples": args.samples,
        "results": {},
    }

    base_bundle = build_base_model_bundle(args.model_id, device=device, dtype=dtype)
    latent_bundle = build_latent_bundle(
        args.model_id,
        checkpoint_path=args.latent_ckpt,
        n_latent_override=-1,
        device=device,
        dtype=dtype,
    )
    try:
        sampling_rate = int(getattr(base_bundle.processor.feature_extractor, "sampling_rate", 16000) or 16000)
        for tag, (dataset_name, config, split) in DATASETS.items():
            print(f"[latency] collect {tag}: {dataset_name}/{config}/{split}")
            samples = collect_samples(
                dataset_name=dataset_name,
                config=config,
                split=split,
                n=args.warmup + args.samples,
                streaming=args.streaming,
                sampling_rate=sampling_rate,
            )
            language_hint = resolve_language_hint(dataset_name, config)
            output["results"][tag] = {}
            settings = [
                ("baseline", base_bundle.model, base_bundle.processor, True, 0.0),
                ("theta_full", latent_bundle.model, latent_bundle.processor, False, -2.0),
                ("theta_zero", latent_bundle.model, latent_bundle.processor, False, 0.0),
                ("theta_skip", latent_bundle.model, latent_bundle.processor, False, 0.5),
            ]
            for name, model, processor, use_baseline, threshold in settings:
                print(f"[latency] {tag} {name}")
                output["results"][tag][name] = measure_setting(
                    model=model,
                    processor=processor,
                    samples=samples,
                    use_baseline=use_baseline,
                    threshold=threshold,
                    language_hint=language_hint,
                    warmup=args.warmup,
                    measure=args.samples,
                    max_new_tokens=args.max_new_tokens,
                )
                print(output["results"][tag][name])
    finally:
        release_model(latent_bundle)
        release_model(base_bundle)

    out_path = Path(args.output_json)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(output, indent=2), encoding="utf-8")
    print(out_path)
    return 0


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