记录三模型话轮检测实验、44点评估及延迟对比

2026-10-14 10:30:22 +08:00
parent 0204f1288a
commit 05f67739ce
75 changed files with 63467 additions and 4 deletions
@@ -0,0 +1,129 @@
+"""Run MaAI ch_kyoto synchronously on a local stereo recording, on CPU."""
+import argparse
+import hashlib
+import importlib.metadata
+import json
+import os
+from pathlib import Path
+import time
+
+HERE = Path(__file__).resolve().parent
+ROOT = HERE.parents[1]
+os.environ.setdefault("HF_HOME", str(HERE / ".cache/huggingface"))
+os.environ.setdefault("TORCH_HOME", str(HERE / ".cache/torch"))
+os.environ.setdefault("MPLCONFIGDIR", str(HERE / ".cache/matplotlib"))
+os.environ.setdefault("PYGAME_HIDE_SUPPORT_PROMPT", "1")
+os.environ.setdefault("HF_HUB_OFFLINE", "1")
+
+import numpy as np
+import soundfile as sf
+import soxr
+import torch
+from maai import Maai, MaaiInput
+
+WEIGHTS = HERE / "models/vap_mc_state_dict_ch_kyoto_10hz_20000msec.pt"
+CPC = HERE / "models/60k_epoch4-d0f474de.pt"
+EXPECTED = {
+    WEIGHTS.name: "65b8a51aea6a257b5b9139107f0f0cf22f94523cd4b6a2e7a356b6c43a4596a8",
+    CPC.name: "d0f474de968c7243c20f43dee1626681618f2f150d3e638e82b79e22563b17f1",
+}
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--audio", type=Path, default=ROOT / "data/韩泰机器人抢话/音频/2c83253a-7977-4e3c-a387-0aa1ea06fbe1.mp3")
+    parser.add_argument("--seconds", type=float, default=22)
+    parser.add_argument("--threads", type=int, default=4)
+    parser.add_argument("--output", type=Path, default=HERE / "results/smoke.json")
+    args = parser.parse_args()
+    if args.seconds <= 0 or args.threads < 1:
+        parser.error("seconds and threads must be positive")
+    hashes = {}
+    for path in (WEIGHTS, CPC):
+        if not path.exists():
+            raise FileNotFoundError(f"Download the official weights first: {path}")
+        hashes[path.name] = hashlib.sha256(path.read_bytes()).hexdigest()
+        if hashes[path.name] != EXPECTED[path.name]:
+            raise ValueError(f"Weight checksum mismatch: {path}")
+    torch.set_num_threads(args.threads)
+    torch.set_num_interop_threads(1)
+    audio, original_sr = sf.read(args.audio, dtype="float32", always_2d=True)
+    if audio.shape[1] != 2:
+        raise ValueError("Expected stereo audio: channel 0 user, channel 1 robot")
+    audio = audio[:round(args.seconds * original_sr)]
+    if original_sr != 16000:
+        audio = soxr.resample(audio, original_sr, 16000, quality="HQ")
+    audio = np.ascontiguousarray(audio, dtype=np.float32)
+
+    load_started = time.perf_counter()
+    # Official CPC checkpoint contains an argparse.Namespace. Allow only
+    # that class for this load, while retaining weights_only safe loading.
+    with torch.serialization.safe_globals([argparse.Namespace]):
+        model = Maai(
+            mode="vap", lang="ch_kyoto",
+            audio_ch1=MaaiInput.Chunk(), audio_ch2=MaaiInput.Chunk(),
+            frame_rate=10, context_len_sec=20, device="cpu",
+            model_type="normal", mc=True, use_kv_cache=True,
+            cpc_model=str(CPC), local_model=str(WEIGHTS),
+            cache_dir=str(HERE / ".cache/huggingface"),
+        )
+    load_seconds = time.perf_counter() - load_started
+    frame_samples = model.audio_samples_per_frame
+    frame_count = len(audio) // frame_samples
+    if frame_count < 1:
+        raise ValueError("Audio must contain at least one 100 ms frame")
+    frames, latencies = [], []
+    inference_started = time.perf_counter()
+    for i in range(frame_count):
+        start = i*frame_samples
+        before = time.perf_counter()
+        model.process(np.ascontiguousarray(audio[start:start+frame_samples, 0]),
+                      np.ascontiguousarray(audio[start:start+frame_samples, 1]))
+        latency_ms = (time.perf_counter() - before) * 1000
+        # A complete CPC frame should always produce exactly one result.
+        # Avoid blocking indefinitely if that assumption fails.
+        result = model.result_dict_queue.get_nowait()
+        assert model.result_dict_queue.empty(), "More than one result per frame"
+        row = dict(audio_end_sec=(i+1)*frame_samples/16000,
+                   inference_ms=latency_ms)
+        for key in ("p_now", "p_future", "vad"):
+            values = np.asarray(result[key], dtype=np.float64).reshape(-1)
+            if len(values) != 2 or not np.isfinite(values).all():
+                raise ValueError(f"Invalid {key}: {values}")
+            if not ((values >= 0) & (values <= 1)).all():
+                raise ValueError(f"Out-of-range {key}: {values}")
+            row[key] = values.tolist()
+        frames.append(row)
+        latencies.append(latency_ms)
+    inference_seconds = time.perf_counter() - inference_started
+    audio_seconds = frame_count*frame_samples/16000
+    warm = latencies[min(5, len(latencies)-1):]
+    report = dict(
+        status="success", audio=str(args.audio.resolve()),
+        original_sample_rate=original_sr, model_sample_rate=16000,
+        audio_seconds=audio_seconds, frame_count=frame_count,
+        model="ch_kyoto noise-robust VAP / CPC", weight_variant="vap_mc",
+        maai_version=importlib.metadata.version("maai"), torch_version=torch.__version__,
+        portaudio_dependency="Homebrew PortAudio 19.7.0 / PyAudio 0.2.14",
+        device="cpu", threads=args.threads, frame_rate=10, context_len_sec=20,
+        use_kv_cache=True, model_load_seconds=load_seconds,
+        inference_seconds=inference_seconds, realtime_factor=inference_seconds/audio_seconds,
+        warm_mean_frame_ms=float(np.mean(warm)), warm_p95_frame_ms=float(np.percentile(warm, 95)),
+        weight_sha256=hashes,
+        output_definitions={
+            "channels": "Index 0=user, index 1=robot; separate stereo channels, no mixing.",
+            "p_now": "Relative activity projection for the two speakers over roughly the next 0–0.6 seconds.",
+            "p_future": "Relative activity projection for the two speakers over roughly the next 0.6–2 seconds.",
+            "vad": "Current voice-activity probability for each speaker.",
+            "audio_end_sec": "Media time measured from file start; MaAI wall-clock result t is not used.",
+            "scope": "Basic offline CPU smoke test; VAP projections are not can_reply or EOT probabilities. Timings exclude model loading and resampling and describe this one recording/run only.",
+        }, frames=frames)
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps(report, ensure_ascii=False, indent=2)+"\n")
+    print(json.dumps({k:v for k,v in report.items() if k not in ("frames", "output_definitions", "weight_sha256")}, ensure_ascii=False, indent=2))
+    print("Last frame:", json.dumps(frames[-1], ensure_ascii=False))
+    print(f"Saved {args.output.resolve()}")
+
+
+if __name__ == "__main__":
+    main()