Initial commit: FunASR Speech Recognition Toolkit

Add complete FunASR codebase including models, runtime, and documentation.

examples/industrial_data_pretraining/fun_asr_nano/client_test.py

@@ -0,0 +1,244 @@
+#!/usr/bin/env python3
+"""Fun-ASR-Nano Streaming ASR - Automated Test Script.
+
+Tests the WebSocket server with audio files and validates the response format.
+
+Usage:
+    python client_test.py --server ws://localhost:10095 --file test_audio.wav
+    python client_test.py --server ws://localhost:10095 --file test_audio.wav --hotwords "热词1,热词2"
+"""
+
+import asyncio
+import argparse
+import json
+import time
+import sys
+import numpy as np
+
+try:
+    import websockets
+except ImportError:
+    print("ERROR: pip install websockets")
+    sys.exit(1)
+
+try:
+    import soundfile as sf
+except ImportError:
+    print("ERROR: pip install soundfile")
+    sys.exit(1)
+
+
+SAMPLE_RATE = 16000
+
+
+def load_audio(file_path):
+    audio, sr = sf.read(file_path)
+    if sr != SAMPLE_RATE:
+        try:
+            import librosa
+            audio = librosa.resample(audio, orig_sr=sr, target_sr=SAMPLE_RATE)
+        except ImportError:
+            print(f"ERROR: Audio is {sr}Hz. Install librosa for resampling.")
+            sys.exit(1)
+    if audio.ndim > 1:
+        audio = audio[:, 0]
+    return audio.astype(np.float32)
+
+
+async def test_basic(server, audio, hotwords=None):
+    """Test basic streaming ASR flow."""
+    print("=" * 60)
+    print("TEST: Basic Streaming ASR")
+    print("=" * 60)
+
+    t_start = time.perf_counter()
+
+    async with websockets.connect(server, ping_interval=None) as ws:
+        await ws.send("START")
+        resp = json.loads(await ws.recv())
+        assert resp["event"] == "started", f"Expected 'started', got: {resp}"
+        print("  [PASS] START -> event:started")
+
+        if hotwords:
+            await ws.send(f"HOTWORDS:{hotwords}")
+            resp = json.loads(await ws.recv())
+            assert resp["event"] == "hotwords_set", f"Expected 'hotwords_set', got: {resp}"
+            print(f"  [PASS] HOTWORDS -> {len(resp['hotwords'])} words set")
+
+        int16 = (audio * 32768).clip(-32768, 32767).astype(np.int16)
+        chunk_size = 4096
+        partial_count = 0
+
+        for i in range(0, len(int16), chunk_size):
+            chunk = int16[i:i+chunk_size]
+            await ws.send(chunk.tobytes())
+            await asyncio.sleep(0.02)
+
+            try:
+                msg = await asyncio.wait_for(ws.recv(), timeout=0.01)
+                data = json.loads(msg)
+                if "sentences" in data:
+                    partial_count += 1
+                    assert "partial" in data
+                    assert "duration_ms" in data
+                    assert "is_final" in data
+                    assert isinstance(data["sentences"], list)
+            except asyncio.TimeoutError:
+                pass
+
+        await ws.send("STOP")
+
+        final_result = None
+        while True:
+            msg = await asyncio.wait_for(ws.recv(), timeout=30)
+            data = json.loads(msg)
+            if data.get("is_final"):
+                final_result = data
+            if data.get("event") == "stopped":
+                break
+
+        t_end = time.perf_counter()
+
+    assert final_result is not None, "No final result received"
+    assert final_result["is_final"] is True
+    sentences = final_result["sentences"]
+    assert len(sentences) > 0, "No sentences in final result"
+
+    print(f"  [PASS] Received final result: {len(sentences)} sentences")
+    print(f"  [PASS] Partial updates received: {partial_count}")
+
+    for s in sentences:
+        assert "text" in s, f"Missing 'text' in sentence: {s}"
+        assert "start" in s, f"Missing 'start' in sentence: {s}"
+        assert "end" in s, f"Missing 'end' in sentence: {s}"
+        assert isinstance(s["text"], str) and len(s["text"]) > 0
+        assert s["end"] > s["start"]
+
+    print("  [PASS] All sentences have valid format (text, start, end)")
+
+    has_spk = all("spk" in s for s in sentences)
+    if has_spk:
+        spk_ids = set(s["spk"] for s in sentences)
+        print(f"  [PASS] Speaker diarization: {len(spk_ids)} speakers detected")
+    else:
+        print("  [INFO] No speaker IDs in result")
+
+    elapsed = t_end - t_start
+    audio_duration = len(audio) / SAMPLE_RATE
+    rtf = elapsed / audio_duration
+    print(f"\n  Audio: {audio_duration:.1f}s | Time: {elapsed:.2f}s | RTF: {rtf:.3f}")
+
+    print("\n  --- Results ---")
+    for s in sentences:
+        spk = f" SPK{s['spk']}" if "spk" in s else ""
+        print(f"  [{s['start']/1000:.1f}-{s['end']/1000:.1f}s]{spk}: {s['text']}")
+
+    return True
+
+
+async def test_empty_audio(server):
+    """Test with very short/empty audio."""
+    print("\n" + "=" * 60)
+    print("TEST: Empty/Short Audio Handling")
+    print("=" * 60)
+
+    async with websockets.connect(server, ping_interval=None) as ws:
+        await ws.send("START")
+        await ws.recv()
+
+        short_audio = np.zeros(800, dtype=np.int16)
+        await ws.send(short_audio.tobytes())
+        await asyncio.sleep(0.1)
+
+        await ws.send("STOP")
+        final = None
+        while True:
+            msg = await asyncio.wait_for(ws.recv(), timeout=10)
+            data = json.loads(msg)
+            if data.get("is_final"):
+                final = data
+            if data.get("event") == "stopped":
+                break
+
+        assert final is not None
+        assert final["sentences"] == [] or all(s["text"].strip() for s in final["sentences"])
+        print("  [PASS] Short audio handled gracefully")
+
+    return True
+
+
+async def test_multiple_sessions(server, audio):
+    """Test multiple consecutive sessions on same connection."""
+    print("\n" + "=" * 60)
+    print("TEST: Multiple Sessions")
+    print("=" * 60)
+
+    audio_short = audio[:SAMPLE_RATE * 5]
+    int16 = (audio_short * 32768).clip(-32768, 32767).astype(np.int16)
+
+    async with websockets.connect(server, ping_interval=None) as ws:
+        for session_num in range(2):
+            await ws.send("START")
+            resp = json.loads(await ws.recv())
+            assert resp["event"] == "started"
+
+            for i in range(0, len(int16), 4096):
+                await ws.send(int16[i:i+4096].tobytes())
+                await asyncio.sleep(0.01)
+
+            await ws.send("STOP")
+            got_final = False
+            while True:
+                msg = await asyncio.wait_for(ws.recv(), timeout=15)
+                data = json.loads(msg)
+                if data.get("is_final"):
+                    got_final = True
+                if data.get("event") == "stopped":
+                    break
+            assert got_final, f"Session {session_num+1}: no final result"
+            print(f"  [PASS] Session {session_num+1} completed")
+
+    return True
+
+
+async def run_tests(args):
+    audio = load_audio(args.file)
+    print(f"Loaded: {args.file} ({len(audio)/SAMPLE_RATE:.1f}s)")
+    print(f"Server: {args.server}\n")
+
+    passed = 0
+    failed = 0
+
+    tests = [
+        ("Basic Streaming", lambda: test_basic(args.server, audio, args.hotwords)),
+        ("Empty Audio", lambda: test_empty_audio(args.server)),
+        ("Multiple Sessions", lambda: test_multiple_sessions(args.server, audio)),
+    ]
+
+    for name, test_fn in tests:
+        try:
+            await test_fn()
+            passed += 1
+        except Exception as e:
+            print(f"  [FAIL] {name}: {e}")
+            failed += 1
+
+    print("\n" + "=" * 60)
+    print(f"RESULTS: {passed} passed, {failed} failed")
+    print("=" * 60)
+    return failed == 0
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Fun-ASR-Nano Streaming Test")
+    parser.add_argument("--server", type=str, default="ws://localhost:10095")
+    parser.add_argument("--file", type=str, required=True, help="Test audio file")
+    parser.add_argument("--hotwords", type=str, default="", help="Hotwords (comma-separated)")
+    args = parser.parse_args()
+
+    success = asyncio.run(run_tests(args))
+    sys.exit(0 if success else 1)
+
+
+if __name__ == "__main__":
+    main()