93 lines
3.6 KiB
Python
93 lines
3.6 KiB
Python
"""Check StreamingSession: streamed ASR must match offline ASR; two-turn chat from the cache; latency after release."""
|
|
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import librosa
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import torch
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
import server # noqa: E402
|
|
from session import StreamingSession # noqa: E402
|
|
|
|
engine = server.Engine()
|
|
tok = engine.tokenizer
|
|
wav, rate = sf.read(Path(__file__).resolve().parent.parent / "smoke_out/tts.wav", dtype="float32")
|
|
CHUNK = 2048 # what the browser worklet sends
|
|
|
|
|
|
def system_block(text):
|
|
return engine.processor.apply_chat_template([{"role": "system", "content": text}], tokenize=False)
|
|
|
|
|
|
def stream_in(s, audio, rate):
|
|
for i in range(0, len(audio), CHUNK):
|
|
s.feed_audio(audio[i:i + CHUNK], rate)
|
|
|
|
|
|
def decode_text(out, suffix_len):
|
|
seq = out.sequences[0, suffix_len:].tolist()
|
|
return tok.decode([t for t in seq if t != engine.model.config.audio_out_token_id], skip_special_tokens=True)
|
|
|
|
|
|
# offline reference
|
|
_, inputs = engine.build_inputs([{"role": "system", "content": "Perform ASR."}, {"role": "user", "content": [{"type": "audio"}]}],
|
|
[librosa.resample(wav, orig_sr=rate, target_sr=16000)], None)
|
|
with torch.inference_mode():
|
|
ref = engine.model.generate(**inputs, max_new_tokens=64, generation_mode="text", return_dict_in_generate=True)
|
|
print("offline ASR :", tok.decode(ref.sequences[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True))
|
|
|
|
s = StreamingSession(engine)
|
|
s.begin_turn(("asr", ""), system_block("Perform ASR."), keep_history=False)
|
|
stream_in(s, wav, rate)
|
|
t = time.perf_counter()
|
|
s.feed_audio(np.zeros(0, np.float32), rate, final=True)
|
|
out = s.generate("<|im_end|>\n<|im_start|>assistant\n", max_new_tokens=64, generation_mode="text")
|
|
print(f"streamed ASR: {decode_text(out, 5)} ({(time.perf_counter() - t) * 1000:.0f} ms after release, {s.audio_tokens} audio tokens)")
|
|
|
|
|
|
class FirstFrame:
|
|
def __init__(self): self.t = None
|
|
def put(self, v): self.t = self.t or time.perf_counter()
|
|
def end(self): pass
|
|
|
|
|
|
sysb = system_block("Respond with interleaved text and audio.")
|
|
conv = [{"role": "system", "content": "Respond with interleaved text and audio."}, {"role": "user", "content": [{"type": "audio"}]}]
|
|
w16 = librosa.resample(wav, orig_sr=rate, target_sr=16000)
|
|
for label in ("warmup", "offline"):
|
|
ff = None
|
|
|
|
class FF:
|
|
t = None
|
|
def put(self, v): self.t = self.t or time.perf_counter()
|
|
def end(self): pass
|
|
ff = FF()
|
|
t = time.perf_counter()
|
|
_, inputs = engine.build_inputs(conv, [w16], "<|mixed_start|>")
|
|
with torch.inference_mode():
|
|
engine.model.generate(**inputs, max_new_tokens=512, audio_streamer=ff, return_dict_in_generate=True)
|
|
print(f"{label} chat: first_frame={(ff.t - t) * 1000:.0f}ms after release (prep + full prefill)")
|
|
|
|
s = StreamingSession(engine)
|
|
for turn in range(3):
|
|
info = s.begin_turn(("chat", ""), sysb, keep_history=True)
|
|
stream_in(s, wav, rate)
|
|
t = time.perf_counter()
|
|
s.feed_audio(np.zeros(0, np.float32), rate, final=True)
|
|
t_flush = time.perf_counter()
|
|
ff = FirstFrame()
|
|
out = s.generate("<|im_end|>\n<|im_start|>assistant\n<|mixed_start|>", max_new_tokens=512, audio_streamer=ff)
|
|
print(f"chat turn {turn}: {info} flush={(t_flush - t) * 1000:.0f}ms first_frame={(ff.t - t) * 1000:.0f}ms "
|
|
f"total={(time.perf_counter() - t):.1f}s ctx={s.length} close={s.needs_close}\n {decode_text(out, 6)!r}")
|
|
|
|
# snapshot / restore round trip
|
|
snap = s.snapshot()
|
|
n = s.length
|
|
s.begin_turn(("chat", ""), sysb, keep_history=True)
|
|
s.restore(snap)
|
|
print("restore ok:", s.length == n)
|