296 lines
17 KiB
Python
296 lines
17 KiB
Python
|
"""MaAI ch_kyoto standard/noise-robust Mimi VAP: causal stereo CPU replay."""
|
||
|
|
import argparse
|
||
|
|
from collections import Counter
|
||
|
|
import contextlib
|
||
|
|
from datetime import datetime, timezone
|
||
|
|
import importlib.metadata as metadata
|
||
|
|
import io
|
||
|
|
import json
|
||
|
|
import platform
|
||
|
|
from pathlib import Path
|
||
|
|
import sys
|
||
|
|
import time
|
||
|
|
|
||
|
|
from smoke_run import HERE, ROOT, Maai, MaaiInput, np, sf, soxr, torch
|
||
|
|
from evaluate_can_reply import metrics, read, write
|
||
|
|
from run_user_only import LABELS, HISTORY, seed, sha
|
||
|
|
|
||
|
|
OUT = HERE / "results/mimi-standard"
|
||
|
|
MODELS = HERE / "models/mimi"
|
||
|
|
VAP = MODELS / "vap_mimi_state_dict_ch_kyoto_12.5hz_20000msec.pt"
|
||
|
|
MC = False
|
||
|
|
ONNX = MODELS / "continuous_mimi_fp32.onnx"
|
||
|
|
META = MODELS / "continuous_mimi_fp32.json"
|
||
|
|
MANIFEST = HERE / "mimi-model-manifest.json"
|
||
|
|
FRAME_MS, FRAME_SAMPLES = 80, 1280
|
||
|
|
|
||
|
|
|
||
|
|
def validate_weights():
|
||
|
|
manifest = json.loads(MANIFEST.read_text())
|
||
|
|
assert {x["filename"] for x in manifest["files"]} == {
|
||
|
|
"vap_mimi_state_dict_ch_kyoto_12.5hz_20000msec.pt",
|
||
|
|
"vap_mc_mimi_state_dict_ch_kyoto_12.5hz_20000msec.pt", ONNX.name, META.name}
|
||
|
|
for item in manifest["files"]:
|
||
|
|
path = MODELS / item["filename"]
|
||
|
|
assert sha(path) == item["sha256"] and path.stat().st_size == item["size"], path.name
|
||
|
|
return manifest
|
||
|
|
|
||
|
|
|
||
|
|
def model():
|
||
|
|
seed()
|
||
|
|
messages = io.StringIO()
|
||
|
|
with contextlib.redirect_stdout(messages):
|
||
|
|
maai = Maai(mode="vap", lang="ch_kyoto", mc=MC,
|
||
|
|
audio_ch1=MaaiInput.Chunk(), audio_ch2=MaaiInput.Chunk(),
|
||
|
|
model_type="normal-ver2", frame_rate=12.5, context_len_sec=20,
|
||
|
|
device="cpu", use_kv_cache=True, local_model=str(VAP),
|
||
|
|
use_mimi_onnx=True, mimi_onnx_precision="fp32",
|
||
|
|
mimi_local_onnx_fp32_path=str(ONNX), mimi_local_onnx_fp32_meta_path=str(META),
|
||
|
|
mimi_onnx_cpu_intra_threads=4, mimi_onnx_cpu_inter_threads=1,
|
||
|
|
cache_dir=str(HERE / ".cache/huggingface"))
|
||
|
|
# Official local_model path uses the common 'vap' architecture marker for
|
||
|
|
# either checkpoint. The filename/hash and loading audit identify MC weights.
|
||
|
|
assert maai.variant == "vap" and not maai.is_mono
|
||
|
|
assert maai.encoder_type == "mimi" and not maai.single_tower
|
||
|
|
assert maai.audio_samples_per_frame == FRAME_SAMPLES
|
||
|
|
maai.CALC_PROCESS_TIME_INTERVAL = 10**9
|
||
|
|
return maai, messages.getvalue()
|
||
|
|
|
||
|
|
|
||
|
|
def loading_audit(maai, messages):
|
||
|
|
raw = torch.load(VAP, map_location="cpu", weights_only=True)
|
||
|
|
sd = {**raw.get("encoder_state_dict", {}), **raw["state_dict"]} if "state_dict" in raw else raw
|
||
|
|
missing, mismatch, bypassed = [], [], []
|
||
|
|
for name, parameter in maai.vap.named_parameters():
|
||
|
|
key = name
|
||
|
|
if name.startswith(("encoder1.", "encoder2.")):
|
||
|
|
encoder = getattr(maai.vap, name.split('.', 1)[0])
|
||
|
|
# At native 12.5Hz, _align_frame_rate returns embeddings directly;
|
||
|
|
# these frozen conversion convolutions are not used or pretrained.
|
||
|
|
if '.frame_rate_conv.' in name and not parameter.requires_grad:
|
||
|
|
assert encoder.frame_hz == encoder.frame_hz_mimi == 12.5
|
||
|
|
bypassed.append(name)
|
||
|
|
continue
|
||
|
|
key = "encoder." + name.split(".", 1)[1]
|
||
|
|
if key not in sd:
|
||
|
|
missing.append(name)
|
||
|
|
elif not torch.equal(parameter.detach().cpu(), sd[key]):
|
||
|
|
mismatch.append(name)
|
||
|
|
assert not missing and not mismatch, dict(missing=missing, mismatched=mismatch)
|
||
|
|
return dict(variant="vap_mc" if MC else "vap", loader_architecture_marker=maai.variant,
|
||
|
|
encoder=maai.encoder_type, checkpoint_mc=MC,
|
||
|
|
single_tower=maai.single_tower, missing_trainable_parameters=missing,
|
||
|
|
mismatched_parameters=mismatch, unused_frozen_frame_rate_parameters=bypassed, loader_messages=messages,
|
||
|
|
encoder_provider=maai.vap.encoder1._onnx_sess.get_providers())
|
||
|
|
|
||
|
|
|
||
|
|
def prepare(pcm, sr, decision):
|
||
|
|
duration = min(20000, decision // FRAME_MS * FRAME_MS)
|
||
|
|
if duration < FRAME_MS * 2:
|
||
|
|
raise ValueError("Insufficient history for Mimi startup")
|
||
|
|
origin = decision - duration
|
||
|
|
start, end = origin * sr // 1000, decision * sr // 1000
|
||
|
|
if not 0 <= start < end <= len(pcm):
|
||
|
|
raise ValueError("Point outside recording")
|
||
|
|
# Preserve BOTH recorded channels. Crop before resampling to avoid future input.
|
||
|
|
segment = np.ascontiguousarray(soxr.resample(pcm[start:end], sr, 16000, quality="HQ"), dtype=np.float32)
|
||
|
|
assert segment.shape == (duration // FRAME_MS * FRAME_SAMPLES, 2)
|
||
|
|
return segment, origin
|
||
|
|
|
||
|
|
|
||
|
|
def consume(maai, frame):
|
||
|
|
maai.process(np.ascontiguousarray(frame[:, 0]), np.ascontiguousarray(frame[:, 1]))
|
||
|
|
if maai.result_dict_queue.empty():
|
||
|
|
return None # Official ONNX Mimi skips its first encoder output.
|
||
|
|
result = maai.result_dict_queue.get_nowait()
|
||
|
|
assert maai.result_dict_queue.empty()
|
||
|
|
values = {k: np.asarray(result[k], dtype=float).reshape(-1) for k in ("p_now", "p_future", "vad")}
|
||
|
|
for value in values.values():
|
||
|
|
assert value.shape == (2,) and np.isfinite(value).all()
|
||
|
|
assert ((value >= 0) & (value <= 1)).all()
|
||
|
|
return values
|
||
|
|
|
||
|
|
|
||
|
|
def score(values):
|
||
|
|
return 1.0 - float(values["p_now"][0])
|
||
|
|
|
||
|
|
|
||
|
|
def replay(segment):
|
||
|
|
began = time.perf_counter()
|
||
|
|
maai, _ = model()
|
||
|
|
load_ms = (time.perf_counter() - began) * 1000
|
||
|
|
times, values, first_output_ms = [], None, None
|
||
|
|
for start in range(0, len(segment), FRAME_SAMPLES):
|
||
|
|
t = time.perf_counter()
|
||
|
|
values = consume(maai, segment[start:start+FRAME_SAMPLES])
|
||
|
|
times.append((time.perf_counter()-t)*1000)
|
||
|
|
if values is not None and first_output_ms is None:
|
||
|
|
first_output_ms = (start+FRAME_SAMPLES)*1000/16000
|
||
|
|
if values is None:
|
||
|
|
raise ValueError("No score at the exact final input frame")
|
||
|
|
return values, dict(model_load_ms=load_ms, replay_inference_ms=sum(times),
|
||
|
|
mean_update_ms=float(np.mean(times)), p95_update_ms=float(np.percentile(times, 95)),
|
||
|
|
first_output_ms=first_output_ms, frame_count=len(times))
|
||
|
|
|
||
|
|
|
||
|
|
def witness(pcm, sr, decision):
|
||
|
|
maai, messages = model()
|
||
|
|
audit = loading_audit(maai, messages)
|
||
|
|
del maai
|
||
|
|
segment, _ = prepare(pcm, sr, decision)
|
||
|
|
changed = pcm.copy()
|
||
|
|
changed[decision*sr//1000:] = np.random.default_rng(42).normal(0, .1, changed[decision*sr//1000:].shape)
|
||
|
|
other, _ = prepare(changed, sr, decision)
|
||
|
|
assert np.array_equal(segment, other), "Future audio changed a causal input"
|
||
|
|
a, timing = replay(segment)
|
||
|
|
b, _ = replay(other)
|
||
|
|
assert all(np.array_equal(a[k], b[k]) for k in a), "Seeded inference differs"
|
||
|
|
assert np.count_nonzero(segment[:, 0]) and np.count_nonzero(segment[:, 1]), "Smoke must exercise both speakers"
|
||
|
|
print(f"WITNESS Mimi {'mc' if MC else 'standard'} CPU seed=42 score={score(a):.7f}; exact repeat, future invariance, both speakers and loading audit passed", flush=True)
|
||
|
|
return dict(repeat_identical=True, future_invariance=True, both_channels_nonzero=True,
|
||
|
|
score=score(a), timing=timing, loading_audit=audit)
|
||
|
|
|
||
|
|
|
||
|
|
def continuous(pcm, sr):
|
||
|
|
maai, _ = model()
|
||
|
|
resampler = soxr.ResampleStream(sr, 16000, 2, dtype="float32", quality="HQ")
|
||
|
|
pending = np.empty((0, 2), dtype=np.float32)
|
||
|
|
points = []
|
||
|
|
block_samples = sr * FRAME_MS // 1000
|
||
|
|
assert block_samples * 1000 == sr * FRAME_MS
|
||
|
|
for start in range(0, len(pcm), block_samples):
|
||
|
|
end = min(len(pcm), start+block_samples)
|
||
|
|
pending = np.concatenate([pending, resampler.resample_chunk(pcm[start:end], last=end == len(pcm))])
|
||
|
|
latest = None
|
||
|
|
while len(pending) >= FRAME_SAMPLES:
|
||
|
|
frame, pending = pending[:FRAME_SAMPLES], pending[FRAME_SAMPLES:]
|
||
|
|
result = consume(maai, frame)
|
||
|
|
if result is not None:
|
||
|
|
latest = score(result)
|
||
|
|
if latest is not None:
|
||
|
|
points.append([round(end*1000/sr, 3), round(latest, 7)])
|
||
|
|
return points
|
||
|
|
|
||
|
|
|
||
|
|
def provenance(manifest):
|
||
|
|
import maai.encoder, maai.model
|
||
|
|
return dict(machine=platform.platform(), hardware="Apple M4 Mac mini / 16 GB", python=sys.version,
|
||
|
|
environment="experiments/maai-ch-kyoto/.venv (uv.lock)", device="cpu",
|
||
|
|
torch_threads=4, torch_interop_threads=1, onnx_intra_threads=4, onnx_inter_threads=1,
|
||
|
|
dependencies={p: metadata.version(p) for p in ("maai", "torch", "transformers", "onnxruntime", "numpy", "soundfile", "soxr")},
|
||
|
|
model_type="normal-ver2", variant="vap_mc" if MC else "vap", mc=MC,
|
||
|
|
checkpoint_filename=VAP.name, encoder="mimi", encoder_precision="fp32",
|
||
|
|
channels={"0": "recorded user", "1": "recorded robot"}, sample_rate=16000,
|
||
|
|
frame_rate=12.5, frame_ms=FRAME_MS, context_ms=20000, seed=42, use_kv_cache=True,
|
||
|
|
weights=manifest, label_sha256=sha(LABELS), history_sha256=sha(HISTORY),
|
||
|
|
uv_lock_sha256=sha(HERE/"uv.lock"), script_sha256=sha(HERE/"run_mimi.py"),
|
||
|
|
maai_encoder_source_sha256=sha(Path(maai.encoder.__file__)),
|
||
|
|
maai_model_source_sha256=sha(Path(maai.model.__file__)))
|
||
|
|
|
||
|
|
|
||
|
|
def references(prov):
|
||
|
|
directory = HERE/"results/user-only"
|
||
|
|
summary = json.loads((directory/"summary.json").read_text())
|
||
|
|
source = summary["provenance"]
|
||
|
|
for key in ("label_sha256", "history_sha256", "uv_lock_sha256"):
|
||
|
|
assert source[key] == prov[key], f"Stale CPC reference: {key}; rerun run_user_only.py all"
|
||
|
|
return {(r["call_id"], int(r["decision_ms"])): r for r in read(directory/"predictions.csv")}, source, sha(directory/"predictions.csv")
|
||
|
|
|
||
|
|
|
||
|
|
def main(mode, variant):
|
||
|
|
global MC, VAP, OUT
|
||
|
|
MC = variant == 'mc'
|
||
|
|
VAP = MODELS / f"vap_{'mc_' if MC else ''}mimi_state_dict_ch_kyoto_12.5hz_20000msec.pt"
|
||
|
|
OUT = HERE / ('results/mimi-mc' if MC else 'results/mimi-standard')
|
||
|
|
manifest = validate_weights()
|
||
|
|
prov = provenance(manifest)
|
||
|
|
labels = read(LABELS)
|
||
|
|
annotated = [r for r in labels if r["can_reply_now"] in ("0", "1")]
|
||
|
|
assert annotated and len({(r['call_id'],r['decision_ms']) for r in labels}) == len(labels)
|
||
|
|
audio_cache, audio_hashes = {}, {}
|
||
|
|
for call in sorted({r["call_id"] for r in labels}):
|
||
|
|
path = ROOT/"data/韩泰机器人抢话/音频"/f"{call}.mp3"
|
||
|
|
pcm, sr = sf.read(path, dtype="float32", always_2d=True)
|
||
|
|
assert pcm.shape[1] == 2
|
||
|
|
audio_cache[call] = pcm, sr
|
||
|
|
audio_hashes[call] = sha(path)
|
||
|
|
began = time.perf_counter()
|
||
|
|
first = annotated[0]
|
||
|
|
checks = witness(*audio_cache[first["call_id"]], int(first["decision_ms"]))
|
||
|
|
if mode == "smoke":
|
||
|
|
print(json.dumps(dict(witness=checks, provenance=prov), ensure_ascii=False, indent=2))
|
||
|
|
return
|
||
|
|
baseline, source_prov, baseline_hash = references(prov)
|
||
|
|
assert source_prov["audio_sha256"] == audio_hashes, "CPC reference uses different source recordings"
|
||
|
|
OUT.mkdir(parents=True, exist_ok=True)
|
||
|
|
rows = []
|
||
|
|
for index, label in enumerate(annotated):
|
||
|
|
call, decision = label['call_id'], int(label['decision_ms'])
|
||
|
|
segment, origin = prepare(*audio_cache[call], decision)
|
||
|
|
values, timing = replay(segment)
|
||
|
|
value = score(values)
|
||
|
|
other = baseline[(call, decision)]
|
||
|
|
assert other['can_reply_now'] == label['can_reply_now']
|
||
|
|
row = dict(label)
|
||
|
|
row.update(can_reply_score=value, prediction_05=int(value > .5), correct_05=int(int(value > .5) == int(label['can_reply_now'])),
|
||
|
|
p_now_user=float(values['p_now'][0]), p_now_robot=float(values['p_now'][1]),
|
||
|
|
p_future_user=float(values['p_future'][0]), p_future_robot=float(values['p_future'][1]),
|
||
|
|
vad_user=float(values['vad'][0]), vad_robot=float(values['vad'][1]),
|
||
|
|
cpc_mc_stereo_score=float(other['stereo_can_reply_score']), cpc_mc_user_only_score=float(other['can_reply_score']),
|
||
|
|
audio_start_ms=origin, audio_end_ms=decision, context_ms=decision-origin, **timing)
|
||
|
|
rows.append(row)
|
||
|
|
print(f"{index+1}/{len(annotated)} {call[:8]} {decision}: label={label['can_reply_now']}, Mimi={value:.4f}, CPC-mc={float(other['stereo_can_reply_score']):.4f}", flush=True)
|
||
|
|
thresholds = [dict(threshold=t, **metrics(rows, 'can_reply_score', t)) for t in (.11, .5, .7, .8, .9, .95, .98)]
|
||
|
|
summary = dict(model=f"VAP / MaAI Mimi {variant} (both speakers)", created_at=datetime.now(timezone.utc).isoformat(),
|
||
|
|
total_points=len(rows), main_rule="1 - p_now[user] > 0.5", main_metrics=metrics(rows,'can_reply_score'),
|
||
|
|
cpc_mc_stereo_reference=metrics(rows,'cpc_mc_stereo_score'), cpc_mc_user_only_reference=metrics(rows,'cpc_mc_user_only_score'),
|
||
|
|
thresholds=thresholds, per_call={c:metrics([r for r in rows if r['call_id']==c],'can_reply_score') for c in audio_cache},
|
||
|
|
changed_vs_cpc_mc_stereo=dict(fixed=sum((r['cpc_mc_stereo_score']>.5)!=bool(int(r['can_reply_now'])) and bool(r['correct_05']) for r in rows),
|
||
|
|
regressed=sum((r['cpc_mc_stereo_score']>.5)==bool(int(r['can_reply_now'])) and not r['correct_05'] for r in rows)),
|
||
|
|
annotation_sources=dict(Counter(r['annotator'] or 'unspecified' for r in rows)),
|
||
|
|
provenance={**prov, 'audio_sha256':audio_hashes, 'cpc_reference_predictions_sha256':baseline_hash}, witness=checks,
|
||
|
|
method=f"{variant} (mc={MC}) ch_kyoto VAP with official FP32 ONNX Mimi encoder. Recorded user+robot stereo; crop preceding <=20s before 16kHz resampling; 80ms native frames, final frame ends at exact decision. Reset Mimi and VAP state per point, retain official skipped startup encoder output. No future input, annotation text, external end signal, silence gate or forced timeout. Score=1-p_now[user], not calibrated EOT probability.",
|
||
|
|
limitations=["Comparison changes encoder/checkpoints and native framing; no held-out benchmark or claim that differences are caused by the encoder alone.",
|
||
|
|
"Both models use <=20s histories but native 80ms/100ms framing may differ at recording start by less than one frame.",
|
||
|
|
"Continuous full-call states differ from independent point replay; checkpoint scores need not lie on the continuous curve.",
|
||
|
|
"Replay computation includes the entire history; it is not one-step online update latency."])
|
||
|
|
if mode == "all":
|
||
|
|
calls = []
|
||
|
|
for call,(pcm,sr) in audio_cache.items():
|
||
|
|
points = continuous(pcm,sr)
|
||
|
|
cs = sorted([r for r in rows if r['call_id']==call],key=lambda r:int(r['decision_ms']))
|
||
|
|
calls.append(dict(call_id=call,audio_path=f'韩泰机器人抢话/音频/{call}.mp3',audio_sha256=audio_hashes[call],
|
||
|
|
duration_ms=round(len(pcm)*1000/sr,3),points=points,
|
||
|
|
checkpoints=[[int(r['decision_ms']),round(r['can_reply_score'],7)] for r in cs]))
|
||
|
|
print(f"{call[:8]} timeline: {len(points)} scores, {len(cs)} checkpoints",flush=True)
|
||
|
|
summary['timeline_points'] = sum(len(c['points']) for c in calls)
|
||
|
|
exported = dict(schema_version=1,model='maai_mimi_mc' if MC else 'maai_mimi_standard',
|
||
|
|
name='VAP / MaAI Mimi（降噪双声道）' if MC else 'VAP / MaAI Mimi（标准双声道）',
|
||
|
|
score='1 - p_now[user]',sample_interval_ms=FRAME_MS,method=summary['method'],
|
||
|
|
provenance=summary['provenance'],calls=calls)
|
||
|
|
assert prov['label_sha256']==sha(LABELS) and prov['history_sha256']==sha(HISTORY),'Annotations changed during inference'
|
||
|
|
write(OUT/'predictions.csv',rows)
|
||
|
|
errors = [r for r in rows if not r['correct_05']]
|
||
|
|
if errors:write(OUT/'errors.csv',errors)
|
||
|
|
else:(OUT/'errors.csv').unlink(missing_ok=True)
|
||
|
|
write(OUT/'thresholds.csv',thresholds)
|
||
|
|
summary['elapsed_seconds'] = time.perf_counter()-began
|
||
|
|
if mode=='all':(OUT/'score-timeline.json').write_text(json.dumps(exported,ensure_ascii=False,separators=(',',':'))+'\n')
|
||
|
|
(OUT/'summary.json').write_text(json.dumps(summary,ensure_ascii=False,indent=2)+'\n')
|
||
|
|
print(json.dumps({k:summary[k] for k in ('main_metrics','cpc_mc_stereo_reference','changed_vs_cpc_mc_stereo','elapsed_seconds')},ensure_ascii=False,indent=2))
|
||
|
|
if mode=='all':
|
||
|
|
from evaluate_streaming import refresh_streaming
|
||
|
|
streaming, _ = refresh_streaming('maai_mimi_mc' if MC else 'maai_mimi_standard')
|
||
|
|
print('PRIMARY streaming metrics:',json.dumps(streaming['main_metrics']),flush=True)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__=='__main__':
|
||
|
|
parser=argparse.ArgumentParser(description=__doc__)
|
||
|
|
parser.add_argument('mode',choices=('smoke','evaluate','all'))
|
||
|
|
parser.add_argument('--variant',choices=('standard','mc','both'),default='standard')
|
||
|
|
args=parser.parse_args()
|
||
|
|
torch.set_num_threads(4)
|
||
|
|
torch.set_num_interop_threads(1)
|
||
|
|
for selected in (('standard','mc') if args.variant=='both' else (args.variant,)):
|
||
|
|
main(args.mode,selected)
|