first tests

scripts/session_test.py

@@ -0,0 +1,92 @@
+"""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)