Initial commit: FunASR Speech Recognition Toolkit

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

examples/industrial_data_pretraining/fun_asr_nano/serve_vllm.py

@@ -0,0 +1,432 @@
+#!/usr/bin/env python3
+"""Fun-ASR-Nano vLLM Inference Server.
+
+Unified server with three interfaces:
+- HTTP REST: POST /asr (file upload)
+- WebSocket: ws://host:port/ws (streaming audio)
+- OpenAI API: POST /v1/audio/transcriptions (Whisper-compatible)
+
+All endpoints share the same vLLM engine + dynamic VAD + SPK + timestamps.
+
+Usage:
+    CUDA_VISIBLE_DEVICES=0 python serve_vllm.py --port 8000
+    CUDA_VISIBLE_DEVICES=0 python serve_vllm.py --port 8000 --model FunAudioLLM/Fun-ASR-Nano-2512
+"""
+
+import asyncio
+import argparse
+import io
+import json
+import logging
+import os
+import re
+import time
+import tempfile
+
+import numpy as np
+import soundfile as sf
+import torch
+import warnings
+
+warnings.filterwarnings('ignore')
+logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s')
+logger = logging.getLogger(__name__)
+
+def truncate_repetition(text, min_repeat_len=3, max_repeats=3):
+    """Detect and truncate repetitive patterns in ASR output."""
+    if not text or len(text) < 20:
+        return text
+    n = len(text)
+    for length in range(min_repeat_len, min(n // max_repeats, 30)):
+        for start in range(n - length * max_repeats):
+            chunk = text[start:start + length]
+            if text[start:start + length * max_repeats] == chunk * max_repeats:
+                return text[:start + length]
+    return text
+
+
+
+try:
+    from fastapi import FastAPI, File, UploadFile, Form, WebSocket, WebSocketDisconnect
+    from fastapi.responses import JSONResponse
+    import uvicorn
+except ImportError:
+    raise ImportError("pip install fastapi uvicorn python-multipart")
+
+from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM
+from funasr.models.fsmn_vad_streaming.dynamic_vad import DynamicStreamingVAD
+from funasr import AutoModel
+
+
+# ============================================================
+# Global state
+# ============================================================
+_engine = None
+_vad_model = None
+_spk_model = None
+_args = None
+
+
+def prepare_audio_for_inference(audio_data, sr, target_sr=16000):
+    """Return mono float32 audio at target_sr for ASR inference."""
+    audio_data = np.asarray(audio_data)
+    if audio_data.ndim > 1:
+        channel_axis = -1 if audio_data.shape[-1] <= audio_data.shape[0] else 0
+        audio_data = audio_data.mean(axis=channel_axis)
+
+    if sr != target_sr:
+        import librosa
+        audio_data = librosa.resample(audio_data, orig_sr=sr, target_sr=target_sr)
+        sr = target_sr
+
+    return audio_data.astype(np.float32), sr
+
+
+def load_engine(args):
+    global _engine, _vad_model, _spk_model, _args
+    _args = args
+    if _engine is None:
+        logger.info(f"Loading vLLM engine: {args.model}")
+        _engine = FunASRNanoVLLM.from_pretrained(
+            model=args.model, hub=args.hub, device=args.device, dtype=args.dtype,
+            max_model_len=args.max_model_len,
+            gpu_memory_utilization=args.gpu_memory_utilization,
+        )
+        logger.info(f"Loading VAD: {args.vad_model}")
+        _vad_model = AutoModel(model=args.vad_model, device=args.device, disable_update=True)
+        if args.spk_model:
+            logger.info(f"Loading SPK: {args.spk_model}")
+            _spk_model = AutoModel(model=args.spk_model, device=args.device, disable_update=True)
+        else:
+            logger.info("SPK disabled")
+        logger.info("All models ready!")
+
+
+def process_audio(audio_data, sr=16000, language=None, hotwords=None, 
+                  use_vad=True, use_spk=False, use_timestamp=True):
+    """Core processing: VAD segment → vLLM ASR → timestamps → SPK."""
+    audio_data, sr = prepare_audio_for_inference(audio_data, sr)
+
+    # VAD segmentation
+    if use_vad and len(audio_data) > sr * 1:
+        vad_res = _vad_model.generate(input=audio_data, fs=sr)
+        segments = vad_res[0]["value"]
+    else:
+        segments = [[0, int(len(audio_data) * 1000 / sr)]]
+
+    if not segments:
+        return {"text": "", "segments": [], "duration": len(audio_data) / sr}
+
+    # Extract segment audio
+    seg_audios = []
+    seg_times = []
+    for seg in segments:
+        s0 = int(seg[0] * sr / 1000)
+        s1 = int(seg[1] * sr / 1000)
+        seg_audio = audio_data[s0:s1]
+        if len(seg_audio) > sr * 0.3:
+            seg_audios.append(seg_audio)
+            seg_times.append((seg[0], seg[1]))
+
+    if not seg_audios:
+        return {"text": "", "segments": [], "duration": len(audio_data) / sr}
+
+    # vLLM batch ASR
+    gen_kwargs = {"max_new_tokens": 500}
+    if language:
+        gen_kwargs["language"] = language
+    if hotwords:
+        gen_kwargs["hotwords"] = hotwords
+
+    results = _engine.generate(inputs=seg_audios, **gen_kwargs)
+
+    # Build segments with timestamps
+    output_segments = []
+    full_text_parts = []
+
+    for i, (r, (start_ms, end_ms)) in enumerate(zip(results, seg_times)):
+        r["text"] = truncate_repetition(r["text"])
+        seg_info = {
+            "text": r["text"],
+            "start": start_ms / 1000,
+            "end": end_ms / 1000,
+        }
+        if use_timestamp and "timestamps" in r:
+            # Offset timestamps by segment start
+            offset = start_ms / 1000
+            seg_info["words"] = [
+                {"word": ts["token"], "start": ts["start_time"] + offset, "end": ts["end_time"] + offset}
+                for ts in r["timestamps"]
+            ]
+        output_segments.append(seg_info)
+        full_text_parts.append(r["text"])
+
+    # SPK diarization
+    if use_spk and _spk_model is not None:
+        from funasr.models.campplus.utils import sv_chunk, postprocess, distribute_spk
+        from funasr.models.campplus.cluster_backend import ClusterBackend
+
+        vad_segs = [[st, et, audio_data[int(st*sr):int(et*sr)]] 
+                    for st, et in [(s["start"], s["end"]) for s in output_segments]]
+        chunks = sv_chunk(vad_segs)
+        if chunks:
+            speech_list = [ch[2] for ch in chunks]
+            spk_res = _spk_model.generate(input=speech_list, cache={}, is_final=True)
+            embs = torch.cat([r["spk_embedding"] for r in spk_res], dim=0)
+            cluster = ClusterBackend(merge_thr=0.78).to(_args.device)
+            labels = cluster(embs.cpu(), oracle_num=None)
+            if not isinstance(labels, np.ndarray):
+                labels = np.array(labels)
+            all_sorted = sorted(chunks, key=lambda x: x[0])
+            sv_output = postprocess(all_sorted, None, labels, embs.cpu())
+            sentences = [{"text": s["text"], "start": int(s["start"]*1000), "end": int(s["end"]*1000)} 
+                        for s in output_segments]
+            distribute_spk(sentences, sv_output)
+            for i, s in enumerate(sentences):
+                output_segments[i]["speaker"] = f"SPK{s.get('spk', 0)}"
+
+    return {
+        "text": " ".join(full_text_parts),
+        "segments": output_segments,
+        "duration": len(audio_data) / sr,
+    }
+
+
+def build_openai_verbose_json(result, language=None):
+    """Build OpenAI-compatible verbose_json while preserving FunASR extensions."""
+    segments = []
+    for i, seg in enumerate(result["segments"]):
+        item = {
+            "id": i,
+            "start": seg["start"],
+            "end": seg["end"],
+            "text": seg["text"],
+            "words": seg.get("words", []),
+        }
+        if "speaker" in seg:
+            item["speaker"] = seg["speaker"]
+        segments.append(item)
+
+    return {
+        "task": "transcribe",
+        "language": language or "zh",
+        "duration": result["duration"],
+        "text": result["text"],
+        "segments": segments,
+    }
+
+
+# ============================================================
+# FastAPI App
+# ============================================================
+app = FastAPI(title="Fun-ASR-Nano vLLM Server", version="1.0")
+
+
+@app.on_event("startup")
+async def startup():
+    load_engine(_args)
+
+
+# --- HTTP REST: POST /asr ---
+@app.post("/asr")
+async def asr_endpoint(
+    file: UploadFile = File(...),
+    language: str = Form(default=None),
+    hotwords: str = Form(default=""),
+    spk: bool = Form(default=False),
+    timestamp: bool = Form(default=True),
+):
+    """ASR with file upload. Returns text + segments + timestamps + speaker."""
+    content = await file.read()
+    audio_data, sr = sf.read(io.BytesIO(content))
+
+    hw_list = [w.strip() for w in hotwords.split(",") if w.strip()] if hotwords else None
+
+    t0 = time.perf_counter()
+    result = process_audio(audio_data, sr=sr, language=language, 
+                          hotwords=hw_list, use_spk=spk, use_timestamp=timestamp)
+    t1 = time.perf_counter()
+
+    result["processing_time"] = round(t1 - t0, 3)
+    result["rtf"] = round((t1 - t0) / result["duration"], 4) if result["duration"] > 0 else 0
+    return JSONResponse(content=result)
+
+
+# --- OpenAI API: POST /v1/audio/transcriptions ---
+@app.post("/v1/audio/transcriptions")
+async def openai_transcriptions(
+    file: UploadFile = File(...),
+    model: str = Form(default="fun-asr-nano"),
+    language: str = Form(default=None),
+    response_format: str = Form(default="json"),
+    timestamp_granularities: str = Form(default="word"),
+    spk: bool = Form(default=False),
+):
+    """OpenAI Whisper-compatible transcription API (extended with spk support)."""
+    content = await file.read()
+    audio_data, sr = sf.read(io.BytesIO(content))
+
+    use_ts = "word" in timestamp_granularities or "segment" in timestamp_granularities
+    result = process_audio(audio_data, sr=sr, language=language, use_spk=spk, use_timestamp=use_ts)
+
+    if response_format == "text":
+        return JSONResponse(content=result["text"])
+    elif response_format == "verbose_json":
+        return JSONResponse(content=build_openai_verbose_json(result, language=language))
+    else:
+        return JSONResponse(content={"text": result["text"]})
+
+
+# --- WebSocket: ws://host:port/ws ---
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    """Streaming WebSocket ASR with dynamic VAD + SPK."""
+    await websocket.accept()
+    logger.info(f"WebSocket connected: {websocket.client}")
+
+    vad = DynamicStreamingVAD(_vad_model)
+    audio_buffer = np.array([], dtype=np.float32)
+    locked_sentences = []
+    language = None
+    hotwords = None
+    use_spk = False
+    is_active = False
+
+    try:
+        while True:
+            message = await websocket.receive()
+
+            if "text" in message:
+                cmd = message["text"].strip()
+                if cmd.upper() == "START":
+                    vad.reset()
+                    audio_buffer = np.array([], dtype=np.float32)
+                    locked_sentences = []
+                    is_active = True
+                    await websocket.send_json({"event": "started"})
+                elif cmd.upper().startswith("LANGUAGE:"):
+                    language = cmd[9:].strip() or None
+                    await websocket.send_json({"event": "language_set", "language": language})
+                elif cmd.upper().startswith("HOTWORDS:"):
+                    hotwords = [w.strip() for w in cmd[9:].split(",") if w.strip()]
+                    await websocket.send_json({"event": "hotwords_set", "hotwords": hotwords})
+                elif cmd.upper().startswith("SPK:"):
+                    use_spk = cmd[4:].strip().lower() in ("true", "1", "on", "yes")
+                    await websocket.send_json({"event": "spk_set", "spk": use_spk})
+                elif cmd.upper() == "STOP":
+                    if is_active and len(audio_buffer) > 0:
+                        # Final: process remaining audio
+                        final_segs = vad.finalize()
+                        for seg in final_segs:
+                            seg_audio = audio_buffer[int(seg[0]*16):int(seg[1]*16)]
+                            if len(seg_audio) > 8000:
+                                gen_kw = {"max_new_tokens": 500}
+                                if language: gen_kw["language"] = language
+                                if hotwords: gen_kw["hotwords"] = hotwords
+                                res = _engine.generate(inputs=[seg_audio], **gen_kw)
+                                if res[0]["text"].strip():
+                                    locked_sentences.append({
+                                        "text": res[0]["text"], "start": seg[0], "end": seg[1]
+                                    })
+
+                        # Handle ongoing speech
+                        if vad.is_speaking:
+                            end_ms = int(len(audio_buffer) * 1000 / 16000)
+                            start_ms = int(vad.current_speech_start) if hasattr(vad, 'current_speech_start') and vad.current_speech_start else 0
+                            seg_audio = audio_buffer[int(start_ms*16):]
+                            if len(seg_audio) > 8000:
+                                gen_kw = {"max_new_tokens": 500}
+                                if language: gen_kw["language"] = language
+                                if hotwords: gen_kw["hotwords"] = hotwords
+                                res = _engine.generate(inputs=[seg_audio], **gen_kw)
+                                if res[0]["text"].strip():
+                                    locked_sentences.append({
+                                        "text": res[0]["text"], "start": start_ms, "end": end_ms
+                                    })
+
+                        # SPK: run full clustering on all sentences (only if enabled)
+                        if use_spk and locked_sentences and _spk_model is not None:
+                            try:
+                                from funasr.models.campplus.utils import sv_chunk, postprocess, distribute_spk
+                                from funasr.models.campplus.cluster_backend import ClusterBackend
+                                vad_segs = [[s["start"]/1000, s["end"]/1000, 
+                                            audio_buffer[int(s["start"]*16):int(s["end"]*16)]]
+                                           for s in locked_sentences]
+                                chunks = sv_chunk(vad_segs)
+                                if chunks:
+                                    speech_list = [ch[2] for ch in chunks]
+                                    spk_res = _spk_model.generate(input=speech_list, cache={}, is_final=True)
+                                    import torch as _torch
+                                    embs = _torch.cat([r["spk_embedding"] for r in spk_res], dim=0)
+                                    cluster = ClusterBackend(merge_thr=0.78).to(_args.device)
+                                    labels = cluster(embs.cpu(), oracle_num=None)
+                                    if not isinstance(labels, np.ndarray):
+                                        labels = np.array(labels)
+                                    all_sorted = sorted(chunks, key=lambda x: x[0])
+                                    sv_output = postprocess(all_sorted, None, labels, embs.cpu())
+                                    spk_sents = [{"text": s["text"], "start": int(s["start"]), "end": int(s["end"])}
+                                                for s in locked_sentences]
+                                    distribute_spk(spk_sents, sv_output)
+                                    for i, ss in enumerate(spk_sents):
+                                        locked_sentences[i]["spk"] = ss.get("spk", 0)
+                            except Exception as e:
+                                logger.warning(f"SPK failed: {e}")
+
+                        await websocket.send_json({
+                            "sentences": locked_sentences,
+                            "is_final": True,
+                            "duration_ms": int(len(audio_buffer) * 1000 / 16000),
+                        })
+                        is_active = False
+                    await websocket.send_json({"event": "stopped"})
+
+            elif "bytes" in message and is_active:
+                pcm = np.frombuffer(message["bytes"], dtype=np.int16).astype(np.float32) / 32768.0
+                audio_buffer = np.concatenate([audio_buffer, pcm])
+
+                # Feed VAD
+                new_confirmed = vad.feed(torch.from_numpy(pcm).float())
+                for seg in new_confirmed:
+                    seg_audio = audio_buffer[int(seg[0]*16):int(seg[1]*16)]
+                    if len(seg_audio) > 8000:
+                        gen_kw = {"max_new_tokens": 500}
+                        if language: gen_kw["language"] = language
+                        if hotwords: gen_kw["hotwords"] = hotwords
+                        res = _engine.generate(inputs=[seg_audio], **gen_kw)
+                        if res[0]["text"].strip():
+                            locked_sentences.append({
+                                "text": res[0]["text"], "start": seg[0], "end": seg[1]
+                            })
+
+                # Send partial update
+                await websocket.send_json({
+                    "sentences": locked_sentences,
+                    "is_final": False,
+                    "duration_ms": int(len(audio_buffer) * 1000 / 16000),
+                })
+
+    except WebSocketDisconnect:
+        logger.info("WebSocket disconnected")
+    except Exception as e:
+        logger.error(f"WebSocket error: {e}", exc_info=True)
+
+
+# ============================================================
+# Main
+# ============================================================
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Fun-ASR-Nano vLLM Server")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    parser.add_argument("--model", type=str, default="FunAudioLLM/Fun-ASR-Nano-2512")
+    parser.add_argument("--hub", type=str, default="ms")
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--dtype", type=str, default="bf16")
+    parser.add_argument("--max-model-len", type=int, default=4096)
+    parser.add_argument("--gpu-memory-utilization", type=float, default=0.5)
+    parser.add_argument("--vad-model", type=str, default="fsmn-vad", help="VAD model name or local path")
+    parser.add_argument("--spk-model", type=str, default="iic/speech_eres2netv2_sv_zh-cn_16k-common", help="Speaker model name or local path (set empty to disable)")
+    _args = parser.parse_args()
+
+    load_engine(_args)
+    uvicorn.run(app, host=_args.host, port=_args.port)