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)

scripts/smoke_test.py

@@ -0,0 +1,52 @@
+"""Headless check of the server's generation path: TTS -> ASR -> voice chat, with streaming."""
+
+import asyncio
+import sys
+import threading
+import time
+from pathlib import Path
+
+import librosa
+import soundfile as sf
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+import server  # noqa: E402
+
+OUT = Path(__file__).resolve().parent.parent / "smoke_out"
+OUT.mkdir(exist_ok=True)
+
+
+def run(engine, conv, clips, prefill, mode):
+    loop = asyncio.new_event_loop()
+    queue = asyncio.Queue()
+    stats, cancel = {}, threading.Event()
+    emit = server.Emitter(loop, queue, time.perf_counter())
+    emit.send = lambda item: None
+    _, inputs = engine.build_inputs(conv, clips, prefill)
+    text = server.TextStream(engine, emit, cancel, stats)
+    audio = server.AudioStream(engine, emit, cancel, stats) if mode != "asr" else None
+    t = time.perf_counter()
+    with torch.inference_mode():
+        engine.model.generate(**inputs, max_new_tokens=512, streamer=text, audio_streamer=audio,
+                              return_dict_in_generate=True, **({"generation_mode": "text"} if mode == "asr" else {}))
+    gen = time.perf_counter() - t
+    pcm = torch.cat([torch.from_numpy(c) for c in audio.collected]).numpy() if audio and audio.collected else None
+    dur = pcm.size / server.OUT_RATE if pcm is not None else 0
+    print(f"[{mode}] prompt={inputs['input_ids'].shape[1]} gen={gen:.2f}s audio={dur:.2f}s stats={stats}\n  text={text.text!r}")
+    return text.text, pcm
+
+
+engine = server.Engine()
+print(engine.info())
+_, pcm = run(engine, [{"role": "system", "content": "Perform TTS in English using a female voice."},
+                      {"role": "user", "content": "What is the tallest mountain in Europe?"}], [], "<|audio_start|>", "tts")
+sf.write(OUT / "tts.wav", pcm, server.OUT_RATE)
+q16 = librosa.resample(pcm, orig_sr=server.OUT_RATE, target_sr=server.IN_RATE)
+run(engine, [{"role": "system", "content": "Perform ASR."}, {"role": "user", "content": [{"type": "audio"}]}], [q16], None, "asr")
+reply, pcm = run(engine, [{"role": "system", "content": "Respond with interleaved text and audio."},
+                          {"role": "user", "content": [{"type": "audio"}]}], [q16], "<|mixed_start|>", "chat")
+sf.write(OUT / "chat.wav", pcm, server.OUT_RATE)
+# sanity-check the chat reply's speech by transcribing it
+r16 = librosa.resample(pcm, orig_sr=server.OUT_RATE, target_sr=server.IN_RATE)
+run(engine, [{"role": "system", "content": "Perform ASR."}, {"role": "user", "content": [{"type": "audio"}]}], [r16], None, "asr")

scripts/tools_test.py

@@ -0,0 +1,64 @@
+"""Exercise "Voice + tools" over the WebSocket: typed requests, then a spoken one (made with the server's TTS)."""
+
+import asyncio
+import json
+
+import numpy as np
+import websockets
+
+URL = "ws://127.0.0.1:8000/ws"
+PROMPTS = [
+    "What's the weather like in Paris right now?",
+    "What is 17 times 23, divided by 4?",
+    "Please make the screen green.",
+    "What time is it in Tokyo?",
+    "Tell me a fun fact about octopuses.",  # should answer without a tool
+]
+
+
+async def turn(ws, send):
+    await send()
+    pcm, events = [], []
+    while True:
+        m = await ws.recv()
+        if isinstance(m, bytes):
+            pcm.append(np.frombuffer(m, dtype=np.float32))
+            continue
+        m = json.loads(m)
+        if m["type"] in ("tool_call", "tool_result", "ui"):
+            events.append({k: v for k, v in m.items() if k not in ("t_ms",)})
+        if m["type"] in ("done", "error"):
+            return m, events, np.concatenate(pcm) if pcm else np.zeros(0, np.float32)
+
+
+async def main():
+    async with websockets.connect(URL, max_size=None) as ws:
+        await ws.recv()
+        await ws.send(json.dumps({"type": "settings", "mode": "tools", "keep_history": False}))
+        for p in PROMPTS:
+            done, events, pcm = await turn(ws, lambda: ws.send(json.dumps({"type": "text", "text": p})))
+            print(f"\n> {p}")
+            for e in events:
+                print("  ", e)
+            s = done.get("stats", {})
+            print(f"  reply: {done.get('text')!r}\n  rounds={s.get('rounds')} audio={s.get('audio_out_s')}s "
+                  f"first_text={s.get('first_text_ms')}ms gen={s.get('gen_ms')}ms")
+
+        # spoken request: synthesize it with TTS, then send it as mic audio
+        q = "What's the weather in Stockholm?"
+        _, _, speech = await turn(ws, lambda: ws.send(json.dumps({"type": "tts", "text": q})))
+
+        async def speak():
+            await ws.send(json.dumps({"type": "start", "rate": 24000}))
+            for i in range(0, len(speech), 2048):
+                await ws.send(speech[i:i + 2048].tobytes())
+            await ws.send(json.dumps({"type": "stop"}))
+
+        done, events, _ = await turn(ws, speak)
+        print(f"\n> (spoken) {q}")
+        for e in events:
+            print("  ", e)
+        print(f"  reply: {done.get('text')!r} rounds={done['stats'].get('rounds')}")
+
+
+asyncio.run(main())

scripts/ws_client_test.py

@@ -0,0 +1,42 @@
+"""Drive the server's WebSocket like the browser does: stream a wav in real time in 2048-sample chunks, 3 turns."""
+
+import asyncio
+import json
+import sys
+
+import numpy as np
+import soundfile as sf
+import websockets
+
+
+async def main(path, url="ws://127.0.0.1:8000/ws"):
+    audio, rate = sf.read(path, dtype="float32")
+    async with websockets.connect(url, max_size=None) as ws:
+        print("<", json.loads(await ws.recv())["type"])
+        for turn in range(3):  # later turns exercise the cached history
+            await ws.send(json.dumps({"type": "start", "rate": rate}))
+            for i in range(0, len(audio), 2048):
+                await ws.send(audio[i:i + 2048].tobytes())
+                await asyncio.sleep(2048 / rate)  # real-time pacing, like a mic
+            await ws.send(json.dumps({"type": "stop"}))
+            pcm_bytes = 0
+            while True:
+                msg = await ws.recv()
+                if isinstance(msg, bytes):
+                    pcm_bytes += len(msg)
+                    continue
+                m = json.loads(msg)
+                if m["type"] in ("token", "audio_chunk", "log"):
+                    continue
+                if m["type"] == "done":
+                    keys = ("first_text_ms", "first_frame_ms", "first_audio_ms", "flush_ms", "prefill_calls",
+                            "prefill_busy_ms", "context_tokens", "context_tokens_after", "rtf", "gen_ms", "codec_decode_ms", "audio_out_s")
+                    print("< done", m["text"][:80], {k: m["stats"].get(k) for k in keys})
+                elif m["type"] != "prompt":
+                    print("<", m["type"], {k: v for k, v in m.items() if k != "type"})
+                if m["type"] in ("done", "error"):
+                    break
+            print(f"  received {pcm_bytes / 4 / 24000:.2f}s of PCM")
+
+
+asyncio.run(main(sys.argv[1]))