first tests

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")