diff --git a/docs/vllm_guide.md b/docs/vllm_guide.md new file mode 100644 index 000000000..62950051a --- /dev/null +++ b/docs/vllm_guide.md @@ -0,0 +1,602 @@ +# FunASR vLLM 推理引擎指南 + +FunASR 集成 vLLM 高吞吐推理引擎，用于加速 LLM-based ASR 模型的自回归解码。支持离线批量推理、SDK 流式推理、生产级 WebSocket 实时服务三种模式。 + +--- + +## 目录 + +1. [概述](#1-概述) +2. [安装与环境](#2-安装与环境) +3. [离线批量推理](#3-离线批量推理) +4. [流式 SDK 推理](#4-流式-sdk-推理) +5. [实时 WebSocket 服务](#5-实时-websocket-服务) +6. [性能对比](#6-性能对比) +7. [API 参考](#7-api-参考) +8. [FAQ](#8-faq) + +--- + +## 1. 概述 + +### 适用模型 + +| 模型 | vLLM 支持 | 说明 | +|------|-----------|------| +| **FunASRNano** | ✓ | Qwen3-0.6B LLM，支持离线和流式 | +| **LLMASR** | ✓ | Whisper + Qwen/Vicuna/LLaMA | +| **GLMASR** | ✓ | GLM-ASR-Nano | +| **QwenAudioWarp** | ✓ | Qwen-Audio | +| Paraformer | ✗ | 非自回归模型（CIF predictor），无 LLM 解码 | +| SenseVoice | ✗ | Whisper-like encoder-decoder，非 LLM | +| Conformer/Transformer | ✗ | CTC/attention 解码，非 LLM | + +### 架构 + +``` +┌────────────────────────────────────────────────────────┐ +│ FunASR + vLLM │ +├────────────────────────────────────────────────────────┤ +│ │ +│ Audio ──→ Frontend ──→ Audio Encoder ──→ Adaptor │ +│ (fbank) (SenseVoice) (Transformer) │ +│ │ │ +│ ▼ │ +│ Audio Embeddings │ +│ │ │ +│ Text Prompt ──→ Tokenize ──→ Embed ──→ [Concat] │ +│ (hotwords/language/itn) │ │ +│ ▼ │ +│ ┌──────────────┐ │ +│ │ vLLM Engine │ │ +│ │ (Qwen3 etc) │ │ +│ │ TP 并行加速 │ │ +│ └──────┬───────┘ │ +│ ▼ │ +│ Generated Text │ +│ │ +│ (可选) CTC Forced Alignment ──→ 字级别时间戳 │ +└────────────────────────────────────────────────────────┘ +``` + +### 三种使用模式 + +| 模式 | 入口 | 适用场景 | +|------|------|---------| +| [离线批量推理](#3-离线批量推理) | `AutoModelVLLM` / `FunASRNanoVLLM` | 大规模转写、批量处理 | +| [流式 SDK 推理](#4-流式-sdk-推理) | `FunASRNanoStreamingVLLM` | 实时字幕展示（SDK 集成） | +| [WebSocket 实时服务](#5-实时-websocket-服务) | `serve_realtime_ws.py` | 生产部署、多端接入 | + +--- + +## 2. 安装与环境 + +### 依赖 + +```bash +pip install funasr>=1.3.0 +pip install vllm>=0.12.0 +pip install safetensors tiktoken websockets regex + +# FunASR 开发模式安装（需要 auto_model_vllm） +cd /path/to/FunASR && pip install -e . +``` + +### 硬件要求 + +| 配置 | 最低要求 | 推荐 | +|------|---------|------| +| GPU 显存 | 8GB | 16GB+ | +| CUDA | 11.8 | 12.0+ | +| GPU 数量 | 1 | 2+（tensor parallel） | + +--- + +## 3. 离线批量推理 + +适用于大规模音频转写、离线批量处理。vLLM 的批处理能力在此场景优势最大。 + +### 通用接口（推荐） + +```python +from funasr.auto.auto_model_vllm import AutoModelVLLM + +model = AutoModelVLLM( + model="FunAudioLLM/Fun-ASR-Nano-2512", + hub="ms", # 或 "hf" + tensor_parallel_size=2, # 多卡并行 + gpu_memory_utilization=0.8, +) + +results = model.generate( + ["audio1.wav", "audio2.wav"], + language="中文", + hotwords=["张三", "北京"], +) +for r in results: + print(f"[{r['key']}] {r['text']}") +``` + +### 直接接口 + +```python +from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM + +engine = FunASRNanoVLLM.from_pretrained( + model="FunAudioLLM/Fun-ASR-Nano-2512", + tensor_parallel_size=4, +) + +results = engine.generate( + inputs="wav.scp", # 支持 scp/jsonl/文件列表 + hotwords=["开放时间"], + language="中文", + max_new_tokens=512, +) +``` + +### 命令行 + +```bash +cd examples/industrial_data_pretraining/fun_asr_nano + +# 单文件 +python demo_vllm.py --input audio.wav --language 中文 + +# 批量 + 多卡 +python demo_vllm.py --input wav.scp --tensor-parallel-size 4 --batch-size 32 + +# 带热词 + 保存结果 +python demo_vllm.py --input audio.wav --hotwords 张三 北京 --output results.jsonl +``` + +--- + +## 4. 流式 SDK 推理 + +将音频按 720ms chunk 逐步处理，输出逐步稳定的识别结果。适用于 SDK 集成实时字幕场景。 + +### 设计原理 + +``` +音频流（720ms chunks） + │ 累积重编码（每个 chunk 包含从头到当前的全部音频） + ▼ +┌──────────────────────┐ +│ Stage 1: 前 10 chunk │ ← 无 prev_text，批量生成 +│ 找到稳定输出 │ +└──────────┬───────────┘ + ▼ +┌──────────────────────┐ +│ Stage 2: 后续 chunk │ ← 用稳定输出作 prev_text +└──────────┬───────────┘ + ▼ +每个 chunk: [fixed 区域（确认）] + [8字 unfixed（可能变）] +``` + +### 用法 + +```python +from funasr.models.fun_asr_nano.inference_vllm_streaming import FunASRNanoStreamingVLLM + +engine = FunASRNanoStreamingVLLM.from_pretrained( + model="FunAudioLLM/Fun-ASR-Nano-2512", + chunk_ms=720, + rollback_chars=8, +) + +for result in engine.streaming_generate("audio.wav", language="中文"): + if result["is_final"]: + print(f"最终: {result['text']}") + else: + print(f"[{result['audio_duration_ms']:.0f}ms] 确认: {result['fixed_text']}") +``` + +### 输出特性 + +| 累积音频 | 输出质量 | +|---------|---------| +| < 1.5s | 空或噪声 | +| 1.5-3.0s | 部分正确 | +| > 3.0s | 准确输出 | + +> Note: `repetition_penalty=1.3` 内部硬编码，防止短 chunk 重复退化。 + +--- + +## 5. 实时 WebSocket 服务 + +生产级实时语音识别服务，集成 VAD 分句 + vLLM 推理 + 说话人分离 + 热词。 + +### 推理逻辑 + +**核心设计：基于 VAD 端点逐段解码（非固定 chunk）** + +``` +音频流 ──→ StreamingVAD (60ms) ──→ 检测到端点 ──→ vLLM 解码整段 + │ │ + │ ▼ + │ locked_sentences（锁定） + │ │ + │ ▼ + │ SPK assign（说话人分配） + │ + └──→ 未结束 ──→ vLLM partial decode（每0.48s） + │ + ▼ + partial text（预览，会覆盖） +``` + +**两条推理路径：** + +| 路径 | 触发条件 | 输出 | +|------|---------|------| +| 确认段解码 | VAD 检测到静音端点 | 锁定到 sentences，永不改变 | +| Partial 预览 | 每 0.48s + 新音频 ≥ 960ms | 临时文字，随时覆盖 | + +**动态 VAD 阈值：** + +| 累积时长 | 静音阈值 | 效果 | +|---------|---------|------| +| ≤ 5s | 2.0s | 短句不切碎 | +| 5-10s | 1.5s | 正常分句 | +| 10-15s | 1.0s | 开始收紧 | +| 15-30s | 0.8s | 较快切分 | +| 30-45s | 0.4s | 防止过长 | +| > 45s | 0.1s | 强制切分 | + +**STOP 最终处理：** +1. 剩余音频喂给 VAD（is_final=True） +2. 强制结束当前在说话的段 +3. 全局 SPK 重聚类（修正说话人 ID） +4. 返回 `is_final: true` + +### 部署 + +```bash +cd examples/industrial_data_pretraining/fun_asr_nano + +# 单卡 +CUDA_VISIBLE_DEVICES=0 python serve_realtime_ws.py --port 10095 --language 中文 + +# 多卡 +CUDA_VISIBLE_DEVICES=0,1 python serve_realtime_ws.py \ + --port 10095 --tensor-parallel-size 2 --language 中文 + +# 完整参数 +python serve_realtime_ws.py \ + --port 10095 \ + --model FunAudioLLM/Fun-ASR-Nano-2512 \ + --hub ms \ + --device cuda:0 \ + --decode-interval 0.48 \ + --hotword-file 热词列表 \ + --language 中文 \ + --dtype bf16 \ + --tensor-parallel-size 1 \ + --gpu-memory-utilization 0.8 \ + --max-model-len 2048 +``` + +### 客户端 + +| 客户端 | 用法 | +|--------|------| +| 浏览器 | 打开 `client_mic.html`（麦克风/文件/热词/说话人） | +| Python CLI | `python client_python.py --server ws://localhost:10095 --mic` | +| 测试脚本 | `python client_test.py --server ws://localhost:10095 --file audio.wav` | + +远程访问需 SSH 端口转发：`ssh -L 10095:localhost:10095 ` + +### WebSocket 协议 + +**客户端 → 服务端：** + +| 消息 | 格式 | 说明 | +|------|------|------| +| 开始会话 | `"START"` | 初始化 session | +| 设置热词 | `"HOTWORDS:词1,词2"` | 可选 | +| 设置语种 | `"LANGUAGE:中文"` | 可选 | +| 音频数据 | `bytes` | PCM16, 16kHz, mono | +| 结束会话 | `"STOP"` | 触发最终解码 | + +**服务端 → 客户端：** + +```json +// 事件 +{"event": "started"} +{"event": "hotwords_set", "hotwords": ["词1", "词2"]} +{"event": "language_set", "language": "中文"} +{"event": "stopped"} + +// 实时结果 +{ + "sentences": [{"text": "已确认", "start": 1700, "end": 5500, "spk": 0}], + "partial": "正在说...", + "partial_start_ms": 5800, + "duration_ms": 7200, + "is_final": false +} + +// 最终结果 +{ + "sentences": [{"text": "...", "start": ..., "end": ..., "spk": ...}, ...], + "partial": "", + "partial_start_ms": 0, + "duration_ms": 10000, + "is_final": true +} +``` + +**交互时序：** + +``` +Client Server + │── "START" ─────────────────→│ + │←─ {"event":"started"} ──────│ + │── "HOTWORDS:张三,北京" ────→│ + │←─ {"event":"hotwords_set"} ─│ + │── "LANGUAGE:中文" ─────────→│ + │←─ {"event":"language_set"} ─│ + │── [audio bytes] ───────────→│ + │←─ {sentences,partial} ──────│ (实时推送) + │── [audio bytes] ───────────→│ + │←─ {sentences,partial} ──────│ + │── "STOP" ──────────────────→│ + │←─ {sentences,is_final:true} │ (最终结果) + │←─ {"event":"stopped"} ──────│ +``` + +### 热词文件 + +默认文件名 `热词列表`（`--hotword-file` 指定），一行一个词： + +``` +张三 +李四 +北京大学 +``` + +也可通过 WebSocket 动态设置：`HOTWORDS:词1,词2,词3` + +--- + +## 6. 性能对比 + +### 离线推理 + +| 配置 | 5.6s 音频延迟 | 相对加速 | +|------|-------------|---------| +| PyTorch (baseline) | 0.89s | 1x | +| vLLM 1-GPU | 0.30s | **3x** | +| vLLM 2-GPU TP | ~0.20s | **4.5x** | + +### 批量吞吐 + +| Batch Size | 1-GPU | 2-GPU | 4-GPU | +|-----------|-------|-------|-------| +| 1 | ~1.5x | ~2x | ~2.5x | +| 16 | ~4x | ~7x | ~12x | +| 32 | ~5x | ~9x | ~15x | + +### WebSocket 实时服务 + +| 指标 | 数值 | +|------|------| +| RTF | < 0.08 | +| 首字延迟 | ~480ms | +| 30s 音频总耗时 | ~2.3s | +| 并发 | 多 WebSocket 连接 | + +### VAD + vLLM Pipeline + +| 场景 | PyTorch 串行 | vLLM 批量 | 加速 | +|------|-------------|-----------|------| +| 10 段 x 5s | ~9s | ~1.5s | **6x** | +| 20 段 x 5s | ~18s | ~2.5s | **7x** | + +--- + +## 7. API 参考 + +### AutoModelVLLM + +```python +from funasr.auto.auto_model_vllm import AutoModelVLLM +``` + +通用入口，自动检测模型类型并选择对应的 vLLM 实现。 + +| 参数 | 默认值 | 说明 | +|------|--------|------| +| `model` | - | 模型名（hub）或本地路径 | +| `hub` | `"ms"` | `"ms"` / `"hf"` | +| `device` | `"cuda:0"` | 音频编码器设备 | +| `dtype` | `"bf16"` | 精度 | +| `tensor_parallel_size` | `1` | vLLM GPU 并行数 | +| `gpu_memory_utilization` | `0.8` | KV Cache 显存比例 | +| `max_model_len` | `4096` | 最大序列长度 | + +### AutoModelVLLM.generate() + +| 参数 | 默认值 | 说明 | +|------|--------|------| +| `inputs` | - | 音频路径/列表/numpy/tensor | +| `language` | `None` | 语种提示 | +| `hotwords` | `None` | 热词列表 | +| `itn` | `True` | 逆文本正则化 | +| `max_new_tokens` | `512` | 最大生成 token | +| `temperature` | `0.0` | 采样温度 | +| `repetition_penalty` | `1.0` | 重复惩罚 | + +**返回**: `[{"key": str, "text": str, "timestamps": [...]}]` + +### FunASRNanoStreamingVLLM.from_pretrained() + +| 参数 | 默认值 | 说明 | +|------|--------|------| +| `model` | `"FunAudioLLM/Fun-ASR-Nano-2512"` | 模型名或路径 | +| `hub` | `"ms"` | 模型来源 | +| `device` | `"cuda:0"` | 设备 | +| `dtype` | `"bf16"` | 精度 | +| `tensor_parallel_size` | `1` | GPU 并行数 | +| `gpu_memory_utilization` | `0.8` | 显存比例 | +| `max_model_len` | `2048` | 序列长度 | +| `chunk_ms` | `720` | chunk 时长 | +| `rollback_chars` | `8` | 回退字符数 | + +### streaming_generate() + +| 参数 | 默认值 | 说明 | +|------|--------|------| +| `audio_input` | - | 音频路径/数据 | +| `chunk_ms` | `720` | chunk 大小 | +| `rollback_chars` | `8` | 回退字符 | +| `hotwords` | `None` | 热词 | +| `language` | `None` | 语种 | +| `max_new_tokens` | `200` | 每 chunk 最大 token | +| `temperature` | `0.0` | 采样温度 | + +**Yields**: `{"text", "fixed_text", "is_final", "chunk_idx", "audio_duration_ms"}` + +> `repetition_penalty=1.3` 内部硬编码。 + +### serve_realtime_ws.py 参数 + +| 参数 | 默认值 | 说明 | +|------|--------|------| +| `--port` | `10095` | WebSocket 端口 | +| `--model` | `FunAudioLLM/Fun-ASR-Nano-2512` | 模型 | +| `--hub` | `ms` | 来源 | +| `--device` | `cuda:0` | 设备 | +| `--decode-interval` | `0.48` | partial 解码间隔（秒） | +| `--hotword-file` | `热词列表` | 热词文件 | +| `--language` | `None` | 语种 | +| `--dtype` | `bf16` | 精度 | +| `--tensor-parallel-size` | `1` | GPU 并行 | +| `--gpu-memory-utilization` | `0.8` | 显存比例 | +| `--max-model-len` | `2048` | 序列长度 | + +--- + +## 8. FAQ + +### Q: 首次启动为什么慢？ +vLLM 需要初始化 KV Cache 和 CUDA Graph warmup，约 60-90 秒。后续推理即时响应。 + +### Q: CUDA OOM 怎么办？ +- 减小 `gpu_memory_utilization`（如 0.6） +- 增加 `tensor_parallel_size` 分摊到多卡 +- 减小 `max_model_len` + +### Q: Paraformer 能用 vLLM 吗？ +不能。Paraformer 是非自回归模型（CIF predictor），所有 token 并行生成，不使用 KV-cache。vLLM 只加速自回归 LLM 解码。 + +### Q: WebSocket 服务和 streaming_generate 有什么区别？ + +| | WebSocket 服务 | streaming_generate | +|---|---|---| +| 分句 | VAD 自然端点 | 固定 720ms chunk | +| 推理 | 每个 VAD 段整体解码 | 累积重编码全部音频 | +| 准确率 | 更高 | 前 3s 较低 | +| 场景 | 生产部署 | SDK 集成 | + +### Q: 流式推理前几秒输出为空？ +正常。模型需要 ~3 秒累积音频才能产生有意义输出，这是模型特性而非 vLLM 限制。 + +### Q: 支持哪些音频格式？ +wav、mp3、flac 等主流格式，采样率自动转为 16kHz。 + +### Q: 浏览器无法使用麦克风？ +Chrome 要求 HTTPS 或 localhost。远程服务器用 SSH 端口转发：`ssh -L 10095:localhost:10095 ` + +### Q: 多个并发连接会互相影响吗？ +不会。每个 WebSocket 连接有独立的 session（VAD/ASR 状态隔离）。vLLM 内部会自动调度。 + +--- + +## 附录：DynamicStreamingVAD + +`funasr.models.fsmn_vad_streaming.dynamic_vad.DynamicStreamingVAD` 是通用的动态阈值流式 VAD 封装， +在 fsmn-vad 基础上根据当前语音段的累积时长动态调整静音切分阈值。 + +### 设计动机 + +fsmn-vad 默认使用固定静音阈值（800ms）。实际场景中： +- 短句（如"好的"）需要等更长的静音才切，否则会把一句话切碎 +- 长段（如会议发言 30s+）需要更快切分，否则 ASR 输入过长导致质量下降 + +### 用法 + +```python +from funasr import AutoModel +from funasr.models.fsmn_vad_streaming.dynamic_vad import DynamicStreamingVAD + +vad_model = AutoModel(model="fsmn-vad", device="cuda:0") + +# 使用默认阈值配置 +vad = DynamicStreamingVAD(vad_model) + +# 或自定义阈值 +vad = DynamicStreamingVAD( + vad_model, + silence_schedule=[ + (3000, 1500), # 累积 <=3s: 等 1.5s 静音 + (10000, 800), # 累积 3-10s: 等 0.8s + (float('inf'), 300), # 累积 >10s: 等 0.3s + ], + speech_noise_thres=0.5, +) +``` + +#### 流式调用 + +```python +import torch + +for audio_chunk in audio_stream: # 实时音频流 + segments = vad.feed(torch.from_numpy(audio_chunk).float()) + for seg in segments: + print(f"Speech: {seg[0]}-{seg[1]}ms") + +# 结束时 +final_segments = vad.finalize() +``` + +#### 非流式调用 + +```python +segments = vad.process(full_audio_tensor) +for seg in segments: + print(f"Speech: {seg[0]}-{seg[1]}ms") +``` + +### 默认阈值配置 + +| 累积时长 | 静音阈值 | 说明 | +|---------|---------|------| +| ≤ 5s | 2.0s | 短句不切碎 | +| 5-10s | 1.5s | 正常分句 | +| 10-15s | 1.0s | 开始收紧 | +| 15-30s | 0.8s | 较快切分 | +| 30-45s | 0.4s | 防止过长 | +| > 45s | 0.1s | 强制切分 | + +### 参数 + +| 参数 | 默认值 | 说明 | +|------|--------|------| +| `vad_model` | - | FunASR AutoModel 加载的 fsmn-vad 实例 | +| `chunk_size_ms` | 60 | VAD 内部处理 chunk 大小 | +| `speech_noise_thres` | 0.5 | 语音/噪声判别阈值 | +| `speech_to_sil_thres_ms` | 150 | 语音转静音基础时间 | +| `silence_schedule` | 见上表 | 动态阈值配置 `[(上限ms, 静音ms), ...]` | +| `sample_rate` | 16000 | 采样率 | + +### 属性 + +| 属性 | 说明 | +|------|------| +| `vad.is_speaking` | 当前是否在语音状态中 | +| `vad.current_duration_ms` | 当前段已累积时长 | +| `vad.current_threshold_ms` | 当前使用的静音阈值 | diff --git a/examples/industrial_data_pretraining/fun_asr_nano/client_mic.html b/examples/industrial_data_pretraining/fun_asr_nano/client_mic.html new file mode 100644 index 000000000..911d1a1fe --- /dev/null +++ b/examples/industrial_data_pretraining/fun_asr_nano/client_mic.html @@ -0,0 +1,160 @@ + + + + +Fun-ASR-Nano · Streaming ASR with Speaker Diarization + + + +
+ +
+
+ +
+ GitHub + ModelScope + HuggingFace +
+
+
+ Streaming ASR + Speaker Diarization (Beta) + Hotword Customization + 31 Languages · 7 Dialects +
+

+ 基于 FunASR 的 vLLM 推理引擎，实现流式语音识别服务。支持实时 VAD 分句、说话人分离 (Beta)、热词定制化（加载人名、地名等实体词列表，提升专有名词识别准确率）、31种语言及中文方言。所有代码与模型已全部开源。
+ Streaming ASR with vLLM engine, real-time VAD, speaker diarization (Beta), hotword customization (names, places, entities), 31 languages & Chinese dialects. Fully open-sourced. +  ·  www.funasr.com +

+
+ +
+
+
+ + + + + + + + + + +
Ready
+
+
+
Click Mic for real-time recognition or Audio File to transcribe a file
Load a Hotwords file (.txt, one word per line) to boost recognition of names, places & entities
+
+
+
+ + + +
+ diff --git a/examples/industrial_data_pretraining/fun_asr_nano/client_python.py b/examples/industrial_data_pretraining/fun_asr_nano/client_python.py new file mode 100644 index 000000000..570fd86cb --- /dev/null +++ b/examples/industrial_data_pretraining/fun_asr_nano/client_python.py @@ -0,0 +1,233 @@ +#!/usr/bin/env python3 +"""Fun-ASR-Nano Python WebSocket Client. + +Supports real-time microphone recording and audio file streaming. + +Usage: + # Microphone mode + python client_python.py --server ws://localhost:10095 --mic + + # File mode + python client_python.py --server ws://localhost:10095 --file audio.wav + + # With hotwords + python client_python.py --server ws://localhost:10095 --file audio.wav --hotwords "张三,李四,北京" + + # Disable speaker diarization display + python client_python.py --server ws://localhost:10095 --mic --no-spk +""" + +import asyncio +import argparse +import json +import sys +import numpy as np + +try: + import websockets +except ImportError: + print("Please install websockets: pip install websockets") + sys.exit(1) + + +SAMPLE_RATE = 16000 +CHUNK_DURATION_MS = 100 +CHUNK_SAMPLES = int(SAMPLE_RATE * CHUNK_DURATION_MS / 1000) + +SPK_COLORS = [ + "\033[36m", "\033[35m", "\033[33m", "\033[32m", + "\033[34m", "\033[91m", "\033[96m", "\033[95m", +] +RESET = "\033[0m" +GRAY = "\033[90m" +GREEN = "\033[92m" + + +def format_time(ms): + s = ms / 1000 + return f"{s:.1f}s" + + +def print_result(data, show_spk=True): + """Print ASR result to terminal.""" + sentences = data.get("sentences", []) + partial = data.get("partial", "") + partial_start = data.get("partial_start_ms", 0) + is_final = data.get("is_final", False) + + sys.stdout.write("\033[2J\033[H") + + print(f"{GREEN}Fun-ASR-Nano Streaming ASR{RESET}") + print(f"{GRAY}{'─' * 60}{RESET}") + + for s in sentences: + start = s.get("start", s.get("start_ms", 0)) + end = s.get("end", s.get("end_ms", 0)) + spk = s.get("spk", -1) + text = s["text"] + + time_str = f"{GRAY}[{format_time(start)}-{format_time(end)}]{RESET}" + spk_str = "" + if show_spk and spk >= 0: + color = SPK_COLORS[spk % len(SPK_COLORS)] + spk_str = f" {color}SPK{spk}{RESET}" + + print(f" {time_str}{spk_str} {text}") + + if partial: + print(f" {GRAY}[{format_time(partial_start)}-...] {partial}{RESET}") + + if is_final: + print(f"\n{GRAY}{'─' * 60}{RESET}") + print(f"{GREEN}Done.{RESET} {len(sentences)} sentences") + else: + print(f"\n{GRAY}Recording... Press Ctrl+C to stop{RESET}") + + sys.stdout.flush() + + +async def run_mic(args): + """Stream from microphone.""" + try: + import sounddevice as sd + except ImportError: + print("Please install sounddevice: pip install sounddevice") + sys.exit(1) + + print(f"Connecting to {args.server}...") + async with websockets.connect(args.server, ping_interval=None) as ws: + await ws.send("START") + resp = await ws.recv() + event = json.loads(resp) + if event.get("event") != "started": + print(f"Unexpected response: {resp}") + return + + if args.hotwords: + await ws.send(f"HOTWORDS:{args.hotwords}") + await ws.recv() + + print("Recording... Press Ctrl+C to stop\n") + + audio_queue = asyncio.Queue() + + def audio_callback(indata, frames, time_info, status): + audio_queue.put_nowait(indata.copy()) + + stream = sd.InputStream( + samplerate=SAMPLE_RATE, channels=1, dtype='int16', + blocksize=CHUNK_SAMPLES, callback=audio_callback, + ) + + async def send_audio(): + with stream: + while True: + chunk = await audio_queue.get() + await ws.send(chunk.tobytes()) + + async def recv_results(): + async for msg in ws: + data = json.loads(msg) + if "sentences" in data: + print_result(data, show_spk=args.spk) + if data.get("is_final") or data.get("event") == "stopped": + break + + send_task = asyncio.create_task(send_audio()) + recv_task = asyncio.create_task(recv_results()) + + try: + await asyncio.gather(send_task, recv_task) + except (KeyboardInterrupt, asyncio.CancelledError): + pass + finally: + send_task.cancel() + if ws.open: + await ws.send("STOP") + async for msg in ws: + data = json.loads(msg) + if "sentences" in data: + print_result(data, show_spk=args.spk) + if data.get("is_final") or data.get("event") == "stopped": + break + + +async def run_file(args): + """Stream an audio file.""" + try: + import soundfile as sf + except ImportError: + print("Please install soundfile: pip install soundfile") + sys.exit(1) + + audio, sr = sf.read(args.file) + if sr != SAMPLE_RATE: + try: + import librosa + audio = librosa.resample(audio, orig_sr=sr, target_sr=SAMPLE_RATE) + except ImportError: + print(f"Audio is {sr}Hz, need 16kHz. Install librosa: pip install librosa") + sys.exit(1) + if audio.ndim > 1: + audio = audio[:, 0] + audio = audio.astype(np.float32) + + duration = len(audio) / SAMPLE_RATE + print(f"File: {args.file} ({duration:.1f}s)") + print(f"Connecting to {args.server}...") + + async with websockets.connect(args.server, ping_interval=None) as ws: + await ws.send("START") + await ws.recv() + + if args.hotwords: + await ws.send(f"HOTWORDS:{args.hotwords}") + await ws.recv() + + int16 = (audio * 32768).clip(-32768, 32767).astype(np.int16) + + chunk_size = CHUNK_SAMPLES + total_chunks = (len(int16) + chunk_size - 1) // chunk_size + + async def send_audio(): + for i in range(0, len(int16), chunk_size): + chunk = int16[i:i+chunk_size] + await ws.send(chunk.tobytes()) + await asyncio.sleep(CHUNK_DURATION_MS / 1000 * 0.5) + await ws.send("STOP") + + async def recv_results(): + async for msg in ws: + data = json.loads(msg) + if "sentences" in data: + print_result(data, show_spk=args.spk) + if data.get("is_final") or data.get("event") == "stopped": + break + + await asyncio.gather(send_audio(), recv_results()) + + +def main(): + parser = argparse.ArgumentParser(description="Fun-ASR-Nano Python Client") + parser.add_argument("--server", type=str, default="ws://localhost:10095") + parser.add_argument("--mic", action="store_true", help="Use microphone input") + parser.add_argument("--file", type=str, help="Audio file to transcribe") + parser.add_argument("--hotwords", type=str, default="", help="Hotwords (comma-separated)") + parser.add_argument("--spk", action="store_true", default=True, help="Show speaker IDs") + parser.add_argument("--no-spk", dest="spk", action="store_false") + args = parser.parse_args() + + if not args.mic and not args.file: + parser.error("Specify --mic or --file") + + try: + if args.mic: + asyncio.run(run_mic(args)) + else: + asyncio.run(run_file(args)) + except KeyboardInterrupt: + print(f"\n{RESET}Interrupted.") + + +if __name__ == "__main__": + main() diff --git a/examples/industrial_data_pretraining/fun_asr_nano/client_test.py b/examples/industrial_data_pretraining/fun_asr_nano/client_test.py new file mode 100644 index 000000000..9efc9f260 --- /dev/null +++ b/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() diff --git a/examples/industrial_data_pretraining/fun_asr_nano/demo_vllm.py b/examples/industrial_data_pretraining/fun_asr_nano/demo_vllm.py new file mode 100644 index 000000000..52e5eef30 --- /dev/null +++ b/examples/industrial_data_pretraining/fun_asr_nano/demo_vllm.py @@ -0,0 +1,171 @@ +#!/usr/bin/env python3 +"""Demo: Fun-ASR-Nano with vLLM inference backend. + +Usage: + # Single GPU (greedy decoding) + python demo_vllm.py + + # Multi-GPU tensor parallel + python demo_vllm.py --tensor-parallel-size 2 + + # Batch inference from wav.scp + python demo_vllm.py --input wav.scp --tensor-parallel-size 4 --batch-size 32 + + # With hotwords and language + python demo_vllm.py --input audio.wav --language 中文 --hotwords 开放时间 周一 +""" + +import argparse +import os +import time + +import torch + + +def main(): + parser = argparse.ArgumentParser(description="Fun-ASR-Nano vLLM Inference Demo") + parser.add_argument( + "--model-dir", + type=str, + default="FunAudioLLM/Fun-ASR-Nano-2512", + help="Model name (from hub) or local directory path", + ) + parser.add_argument("--input", type=str, default=None, help="Audio file, wav.scp, or jsonl") + parser.add_argument("--hub", type=str, default="ms", choices=["ms", "hf"]) + parser.add_argument("--device", type=str, default="cuda:0", help="Device for audio encoder") + parser.add_argument("--dtype", type=str, default="bf16", choices=["bf16", "fp16", "fp32"]) + parser.add_argument( + "--tensor-parallel-size", type=int, default=1, help="Number of GPUs for vLLM" + ) + parser.add_argument("--gpu-memory-utilization", type=float, default=0.8) + parser.add_argument("--max-model-len", type=int, default=2048) + parser.add_argument("--max-new-tokens", type=int, default=512) + parser.add_argument("--language", type=str, default="中文", help="Language hint") + parser.add_argument("--hotwords", type=str, nargs="*", default=[], help="Hotwords list") + parser.add_argument("--no-itn", action="store_true", help="Disable inverse text normalization") + parser.add_argument("--batch-size", type=int, default=16, help="Batch size for inference") + parser.add_argument("--output", type=str, default=None, help="Output file for results") + args = parser.parse_args() + + from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM + + print(f"=" * 60) + print(f"Fun-ASR-Nano vLLM Inference") + print(f"=" * 60) + print(f" Model: {args.model_dir}") + print(f" Tensor Parallel: {args.tensor_parallel_size} GPU(s)") + print(f" Dtype: {args.dtype}") + print(f" Language: {args.language}") + print(f" Hotwords: {args.hotwords or '(none)'}") + print() + + t_load = time.perf_counter() + engine = FunASRNanoVLLM.from_pretrained( + model=args.model_dir, + hub=args.hub, + device=args.device, + dtype=args.dtype, + tensor_parallel_size=args.tensor_parallel_size, + gpu_memory_utilization=args.gpu_memory_utilization, + max_model_len=args.max_model_len, + ) + print(f"Model loaded in {time.perf_counter() - t_load:.1f}s\n") + + # Determine input files + if args.input is None: + # Use default example audio + example_dir = os.path.join(engine.model_dir, "example") + if os.path.isdir(example_dir): + wav_files = [ + os.path.join(example_dir, f) + for f in sorted(os.listdir(example_dir)) + if f.endswith((".wav", ".mp3", ".flac")) + ] + else: + print("No --input specified and no example/ directory found.") + print("Usage: python demo_vllm.py --input ") + return + if not wav_files: + print("No audio files found in example/ directory.") + return + audio_files = wav_files + print(f"Using example audio: {audio_files}") + elif args.input.endswith(".scp"): + audio_files = [] + with open(args.input, "r") as f: + for line in f: + parts = line.strip().split(maxsplit=1) + if len(parts) == 2: + audio_files.append(parts[1]) + elif len(parts) == 1: + audio_files.append(parts[0]) + print(f"Loaded {len(audio_files)} files from {args.input}") + elif args.input.endswith(".jsonl"): + import json + + audio_files = [] + with open(args.input, "r") as f: + for line in f: + item = json.loads(line.strip()) + audio_files.append(item["source"]) + print(f"Loaded {len(audio_files)} files from {args.input}") + else: + audio_files = [args.input] + + # Run inference in batches + all_results = [] + total_audio_time = 0 + total_infer_time = 0 + + print(f"\nProcessing {len(audio_files)} audio file(s)...") + for i in range(0, len(audio_files), args.batch_size): + batch = audio_files[i : i + args.batch_size] + t0 = time.perf_counter() + results = engine.generate( + inputs=batch, + hotwords=args.hotwords if args.hotwords else None, + language=args.language, + itn=not args.no_itn, + max_new_tokens=args.max_new_tokens, + ) + t1 = time.perf_counter() + batch_time = t1 - t0 + total_infer_time += batch_time + all_results.extend(results) + + batch_num = i // args.batch_size + 1 + total_batches = (len(audio_files) + args.batch_size - 1) // args.batch_size + print(f" Batch {batch_num}/{total_batches}: {len(batch)} files in {batch_time:.2f}s") + + # Print results + print(f"\n{'=' * 60}") + print(f"Results: {len(all_results)} samples, total inference time: {total_infer_time:.2f}s") + print(f"{'=' * 60}") + for r in all_results: + print(f"\n[{r['key']}]") + print(f" Text: {r['text']}") + if "timestamps" in r and r["timestamps"]: + ts_preview = r["timestamps"][:5] + ts_str = " | ".join( + [f"{t['token']}({t['start_time']:.2f}-{t['end_time']:.2f}s)" for t in ts_preview] + ) + if len(r["timestamps"]) > 5: + ts_str += f" ... ({len(r['timestamps'])} total)" + print(f" Timestamps: {ts_str}") + + # Save results to file + if args.output: + import json + + with open(args.output, "w", encoding="utf-8") as f: + for r in all_results: + # Remove non-serializable fields + out = {k: v for k, v in r.items() if k != "timestamps"} + if "timestamps" in r: + out["timestamps"] = r["timestamps"] + f.write(json.dumps(out, ensure_ascii=False) + "\n") + print(f"\nResults saved to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/examples/industrial_data_pretraining/fun_asr_nano/docs/realtime_demo.md b/examples/industrial_data_pretraining/fun_asr_nano/docs/realtime_demo.md new file mode 100644 index 000000000..53d4e90bf --- /dev/null +++ b/examples/industrial_data_pretraining/fun_asr_nano/docs/realtime_demo.md @@ -0,0 +1,63 @@ +# Fun-ASR-Nano 实时 WebSocket 服务 — 快速上手 + +> 完整文档请参见：[FunASR vLLM 推理引擎指南](../../../../docs/vllm_guide.md) + +## 30 秒启动 + +```bash +cd examples/industrial_data_pretraining/fun_asr_nano + +# 安装依赖 +pip install -r requirements.txt +pip install vllm>=0.12.0 + +# 启动服务 +CUDA_VISIBLE_DEVICES=0 python serve_realtime_ws.py --port 10095 --language 中文 +``` + +## 客户端 + +```bash +# 浏览器 +open client_mic.html + +# Python 麦克风 +python client_python.py --server ws://localhost:10095 --mic + +# Python 文件 +python client_python.py --server ws://localhost:10095 --file audio.wav + +# 自动化测试 +python client_test.py --server ws://localhost:10095 --file audio.wav +``` + +## 远程访问 + +```bash +ssh -L 10095:localhost:10095 +# 然后本地打开 client_mic.html +``` + +## 功能 + +- **vLLM 推理引擎**：RTF < 0.08，支持 tensor parallel 多卡加速 +- **流式 VAD 分句**：动态静音阈值，自然断句 +- **说话人分离 (Beta)**：流式 ID 分配 + 最终重聚类 +- **热词定制化**：加载人名、地名等实体词文件 +- **语种指定**：31 种语言 + 中文方言 +- **幻觉检测**：自动检测重复模式并截断 + +## 文件列表 + +| 文件 | 说明 | +|------|------| +| `serve_realtime_ws.py` | WebSocket 服务端 | +| `client_mic.html` | 浏览器客户端 | +| `client_python.py` | Python CLI 客户端 | +| `client_test.py` | 自动化测试脚本 | +| `热词列表` | 热词文件示例 | +| `demo_vllm.py` | 离线 vLLM 推理 demo | + +## 详细文档 + +- [FunASR vLLM 推理引擎指南](../../../../docs/vllm_guide.md) — 完整文档（离线/流式/WebSocket/API/FAQ） diff --git a/examples/industrial_data_pretraining/fun_asr_nano/model.py b/examples/industrial_data_pretraining/fun_asr_nano/model.py index c2830e815..491fd0ecb 100644 --- a/examples/industrial_data_pretraining/fun_asr_nano/model.py +++ b/examples/industrial_data_pretraining/fun_asr_nano/model.py @@ -470,7 +470,7 @@ class FunASRNano(nn.Module): ): meta_data = {} - if kwargs.get("batch_size", 1) > 1: + if isinstance(data_in, (list, tuple)) and len(data_in) > 1: raise NotImplementedError("batch decoding is not implemented") contents = self.data_template(data_in[0]) diff --git a/examples/industrial_data_pretraining/fun_asr_nano/requirements.txt b/examples/industrial_data_pretraining/fun_asr_nano/requirements.txt index eac7d4a42..8bb51b590 100644 --- a/examples/industrial_data_pretraining/fun_asr_nano/requirements.txt +++ b/examples/industrial_data_pretraining/fun_asr_nano/requirements.txt @@ -1,9 +1,10 @@ -torch>=2.9.0 -torchaudio>=2.9.0 -transformers>=4.51.3 -funasr>=1.3.0 -zhconv -whisper_normalizer -pyopenjtalk-plus -compute-wer -openai-whisper +torch>=2.0 +funasr +websockets>=12.0 +regex +numpy +soundfile + +# Optional: for Python client microphone mode +# sounddevice +# librosa diff --git a/examples/industrial_data_pretraining/fun_asr_nano/serve_realtime_ws.py b/examples/industrial_data_pretraining/fun_asr_nano/serve_realtime_ws.py new file mode 100644 index 000000000..5c58bc5ab --- /dev/null +++ b/examples/industrial_data_pretraining/fun_asr_nano/serve_realtime_ws.py @@ -0,0 +1,518 @@ +#!/usr/bin/env python3 +"""Fun-ASR-Nano Streaming WebSocket Server. + +Features: +- Streaming VAD segmentation (fsmn-vad) +- Per-segment ASR decoding (Fun-ASR-Nano via vLLM) +- Speaker diarization (eres2netv2 + ClusterBackend) +- Hotword customization +- Hallucination detection & prevention +""" + +import asyncio +import json +import logging +import os +import time +import argparse +import numpy as np +import torch +import warnings +import regex +import websockets + +warnings.filterwarnings('ignore') +logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s') +logger = logging.getLogger(__name__) + + +def detect_and_fix_hallucination(text, max_ngram_length=12, max_occurrences=3): + """Detect repeated patterns (hallucination) and truncate to keep one occurrence.""" + if not text or len(text) < max_ngram_length * 2: + return text, False + + cleaned = regex.sub(r'\p{P}+', '', text) + + word_pattern = rf'(?= 0: + end_pos = text.find(repeated, pos + len(repeated)) + if end_pos >= 0: + return text[:end_pos + len(repeated)], True + return text[:len(text)//2], True + + for length in range(1, max_ngram_length): + pattern = rf'(?= 0: + end_pos = text.find(repeated, pos + len(repeated)) + if end_pos >= 0: + return text[:end_pos + len(repeated)], True + return text[:len(text)//2], True + + return text, False + + +def _clean_asr_text(text): + """Remove timestamp tags and artifacts from vLLM output.""" + import re + text = re.sub(r'<[^>]*>', '', text) + text = re.sub(r'\[.*?\]', '', text) + text = re.sub(r'[Ｏ\[\]&＆|｜]', '', text) + text = re.sub(r'/sil|endofbreak|FFFF', '', text) + text = re.sub(r'\s+', ' ', text) + return text.strip() + + +from funasr.models.fsmn_vad_streaming.dynamic_vad import DynamicStreamingVAD + + +class HybridSpeakerTracker: + """Speaker diarization: streaming ClusterBackend + final re-clustering.""" + + def __init__(self, spk_model, device, threshold=0.6): + self.spk_model = spk_model + self.device = device + self.threshold = threshold + self.speaker_centers = [] + from funasr.models.campplus.utils import sv_chunk, postprocess, distribute_spk + from funasr.models.campplus.cluster_backend import ClusterBackend + self.sv_chunk = sv_chunk + self.postprocess = postprocess + self.distribute_spk = distribute_spk + self.cluster_backend = ClusterBackend(merge_thr=0.78).to(device) + self.all_chunks = [] + self.all_embeddings = [] + self.display_map = {} + self.next_display_id = 0 + + @torch.no_grad() + def assign_streaming(self, audio_samples, seg_start_s, seg_end_s, sentence): + """Assign speaker ID during streaming using ClusterBackend.""" + vad_seg = [[seg_start_s, seg_end_s, audio_samples]] + chunks = self.sv_chunk(vad_seg) + if not chunks: + sentence["spk"] = self.next_display_id + self.next_display_id += 1 + return + + self.all_chunks.extend(chunks) + speech_list = [ch[2] for ch in chunks] + spk_res = self.spk_model.generate(input=speech_list, cache={}, is_final=True) + embs = torch.cat([r["spk_embedding"] for r in spk_res], dim=0) + self.all_embeddings.append(embs) + + all_embs = torch.cat(self.all_embeddings, dim=0) + labels = self.cluster_backend(all_embs.cpu(), oracle_num=None) + if not isinstance(labels, np.ndarray): + labels = np.array(labels) + + all_sorted = sorted(self.all_chunks, key=lambda x: x[0]) + sv_output = self.postprocess(all_sorted, None, labels, all_embs.cpu()) + temp = [{"start": int(seg_start_s*1000), "end": int(seg_end_s*1000), "text": sentence["text"]}] + self.distribute_spk(temp, sv_output) + raw_spk = temp[0].get("spk", 0) + + if raw_spk not in self.display_map: + self.display_map[raw_spk] = self.next_display_id + self.next_display_id += 1 + sentence["spk"] = self.display_map[raw_spk] + + @torch.no_grad() + def finalize(self, sentences, min_split_s=3.0): + """Final re-clustering for accurate speaker assignment.""" + if not self.all_embeddings or not sentences: + return sentences + + all_embs = torch.cat(self.all_embeddings, dim=0) + labels = self.cluster_backend(all_embs.cpu(), oracle_num=None) + if not isinstance(labels, np.ndarray): + labels = np.array(labels) + + all_sorted = sorted(self.all_chunks, key=lambda x: x[0]) + sv_output = self.postprocess(all_sorted, None, labels, all_embs.cpu()) + + for s in sentences: + s.pop("spk", None) + self.distribute_spk(sentences, sv_output) + + id_map = {} + next_id = 0 + for s in sentences: + raw = s.get("spk", 0) + if raw not in id_map: + id_map[raw] = next_id + next_id += 1 + s["spk"] = id_map[raw] + + final_sentences = [] + for s in sentences: + sub = self._try_split(s, sv_output, id_map, min_split_s) + final_sentences.extend(sub) + + return final_sentences + + def _try_split(self, sentence, sv_output, id_map, min_split_s): + """Split a sentence if multiple speakers detected within its time range.""" + sent_start = sentence["start"] / 1000.0 + sent_end = sentence["end"] / 1000.0 + text = sentence["text"] + + overlapping = [] + for sv_start, sv_end, sv_spk in sv_output: + o_start = max(sent_start, sv_start) + o_end = min(sent_end, sv_end) + if o_end > o_start: + mapped_spk = id_map.get(int(sv_spk), int(sv_spk)) + overlapping.append([o_start, o_end, mapped_spk]) + + if len(overlapping) <= 1: + return [sentence] + + filtered = [overlapping[0]] + for i in range(1, len(overlapping)): + cur = overlapping[i] + prev = filtered[-1] + if cur[2] == prev[2]: + filtered[-1] = [prev[0], cur[1], prev[2]] + elif (cur[1] - cur[0]) < min_split_s: + filtered[-1] = [prev[0], cur[1], prev[2]] + else: + filtered.append(cur) + + merged = [filtered[0]] + for i in range(1, len(filtered)): + if (merged[-1][1] - merged[-1][0]) < min_split_s: + merged[-1] = [merged[-1][0], filtered[i][1], filtered[i][2]] + else: + merged.append(filtered[i]) + if len(merged) > 1 and (merged[-1][1] - merged[-1][0]) < min_split_s: + merged[-2] = [merged[-2][0], merged[-1][1], merged[-2][2]] + merged.pop() + + if len(merged) <= 1: + return [sentence] + + total_dur = sum(m[1] - m[0] for m in merged) + sub_sentences = [] + char_pos = 0 + for i, (m_start, m_end, m_spk) in enumerate(merged): + if i == len(merged) - 1: + sub_text = text[char_pos:] + else: + n_chars = max(1, int(len(text) * (m_end - m_start) / total_dur)) + sub_text = text[char_pos:char_pos + n_chars] + char_pos += n_chars + if sub_text.strip(): + sub_sentences.append({"text": sub_text.strip(), "start": int(m_start*1000), "end": int(m_end*1000), "spk": m_spk}) + + return sub_sentences if sub_sentences else [sentence] + + def reset(self): + self.speaker_centers = [] + self.all_chunks = [] + self.all_embeddings = [] + self.display_map = {} + self.next_display_id = 0 + + +class RealtimeASRSession: + """Manages a single streaming ASR session.""" + + def __init__(self, vllm_engine, asr_kwargs, vad, spk_tracker=None, sample_rate=16000, chunk_ms=960): + self.vllm_engine = vllm_engine + self.asr_kwargs = asr_kwargs + self.vad = vad + self.sample_rate = sample_rate + self.chunk_samples = int(sample_rate * chunk_ms / 1000) + self.first_chunk_samples = int(sample_rate * 480 / 1000) + self.first_decode_done = False + + self.audio_buffer = np.array([], dtype=np.float32) + self.vad_fed_samples = 0 + self.prev_text = "" + self.last_partial_text = "" + self.last_decode_samples = 0 + self.locked_sentences = [] + self.prev_seg_text = "" + self.spk_tracker = spk_tracker + self.use_context = True + self.is_active = False + + def add_audio(self, pcm_bytes): + audio_int16 = np.frombuffer(pcm_bytes, dtype=np.int16) + audio_float = audio_int16.astype(np.float32) / 32768.0 + self.audio_buffer = np.concatenate([self.audio_buffer, audio_float]) + + new_audio = self.audio_buffer[self.vad_fed_samples:] + if len(new_audio) > 0: + new_confirmed = self.vad.feed(torch.from_numpy(new_audio).float(), is_final=False) + self.vad_fed_samples = len(self.audio_buffer) + + for seg in new_confirmed: + seg_text = self._decode_segment(seg) + self.prev_text = "" + if not seg_text.strip(): + continue + self.locked_sentences.append({"text": seg_text, "start": int(seg[0]), "end": int(seg[1])}) + if self.spk_tracker: + s0 = int(seg[0] * self.sample_rate / 1000) + s1 = min(int(seg[1] * self.sample_rate / 1000), len(self.audio_buffer)) + self.spk_tracker.assign_streaming(self.audio_buffer[s0:s1], seg[0]/1000, seg[1]/1000, self.locked_sentences[-1]) + logger.info(f"Locked: [{seg[0]}-{seg[1]}ms] \"{seg_text[:40]}\"") + + def should_decode(self): + threshold = self.first_chunk_samples if not self.first_decode_done else self.chunk_samples + return (len(self.audio_buffer) - self.last_decode_samples) >= threshold + + @torch.no_grad() + def decode(self, is_final=False): + if len(self.audio_buffer) < self.chunk_samples: + return self._build_response(is_final) + + if is_final: + remaining = self.audio_buffer[self.vad_fed_samples:] + if len(remaining) > 0: + new_confirmed = self.vad.feed(torch.from_numpy(remaining).float(), is_final=True) + self.vad_fed_samples = len(self.audio_buffer) + for seg in new_confirmed: + seg_text = self._decode_segment(seg) + if not seg_text.strip(): + continue + self.locked_sentences.append({"text": seg_text, "start": int(seg[0]), "end": int(seg[1])}) + + if self.vad.current_speech_start is not None: + end_ms = int(len(self.audio_buffer) * 1000 / self.sample_rate) + seg = [self.vad.current_speech_start, end_ms] + seg_text = self._decode_segment(seg) + if seg_text.strip(): + self.locked_sentences.append({"text": seg_text, "start": int(seg[0]), "end": int(seg[1])}) + self.vad.current_speech_start = None + + if self.spk_tracker and self.locked_sentences: + self.locked_sentences = self.spk_tracker.finalize(self.locked_sentences) + return self._build_response(is_final) + + if self.vad.current_speech_start is not None: + seg_start_sample = int(self.vad.current_speech_start * self.sample_rate / 1000) + seg_audio = self.audio_buffer[seg_start_sample:] + else: + self.last_decode_samples = len(self.audio_buffer) + self.last_partial_text = "" + return self._build_response(is_final) + + if len(seg_audio) < self.chunk_samples // 2: + return self._build_response(is_final) + + audio_tensor = torch.from_numpy(seg_audio).float() + try: + results = self.vllm_engine.generate( + inputs=[audio_tensor], + hotwords=self.asr_kwargs.get("hotwords"), + language=self.asr_kwargs.get("language"), + max_new_tokens=200, + ) + text = results[0]["text"] if results else "" + text = _clean_asr_text(text) + except Exception as e: + logger.error(f"ASR error: {e}") + return self._build_response(is_final) + + text, hallucinated = detect_and_fix_hallucination(text) + if hallucinated: + self.prev_text = "" + + self.last_decode_samples = len(self.audio_buffer) + self.last_partial_text = text + if text.strip() and not self.first_decode_done: + self.first_decode_done = True + + tokenizer = self.vllm_engine._engine.tokenizer + encoded = tokenizer.encode(text) + if len(encoded) > 5: + self.prev_text = tokenizer.decode(encoded[:-5], skip_special_tokens=True) + else: + self.prev_text = "" + + return self._build_response(is_final) + + @torch.no_grad() + def _decode_segment(self, seg): + """Decode a completed VAD segment via vLLM.""" + start_sample = int(seg[0] * self.sample_rate / 1000) + end_sample = min(int(seg[1] * self.sample_rate / 1000), len(self.audio_buffer)) + seg_audio = self.audio_buffer[start_sample:end_sample] + if len(seg_audio) < 1600: + return "" + audio_tensor = torch.from_numpy(seg_audio).float() + try: + results = self.vllm_engine.generate( + inputs=[audio_tensor], + hotwords=self.asr_kwargs.get("hotwords"), + language=self.asr_kwargs.get("language"), + max_new_tokens=512, + ) + text = results[0]["text"] if results else "" + text = _clean_asr_text(text) + self.prev_seg_text = text + return text + except Exception as e: + logger.error(f"Segment decode error: {e}") + return "" + + def _build_response(self, is_final): + duration_ms = int(len(self.audio_buffer) * 1000 / self.sample_rate) + sentences = list(self.locked_sentences) + partial = self.last_partial_text + partial_start = self.vad.current_speech_start or duration_ms + + if is_final: + return {"sentences": sentences, "partial": "", "partial_start_ms": 0, + "duration_ms": duration_ms, "is_final": True} + return {"sentences": sentences, "partial": partial, + "partial_start_ms": partial_start, + "duration_ms": duration_ms, "is_final": False} + + def reset(self): + self.audio_buffer = np.array([], dtype=np.float32) + self.vad_fed_samples = 0 + self.first_decode_done = False + self.vad.reset() + self.prev_text = "" + self.last_partial_text = "" + self.last_decode_samples = 0 + self.locked_sentences = [] + if self.spk_tracker: + self.spk_tracker.reset() + + +_vllm_engine = None +_asr_kwargs = None +_vad_model = None +_spk_model = None + + +def load_models(args): + global _vllm_engine, _asr_kwargs, _vad_model, _spk_model + if _vllm_engine is None: + from funasr import AutoModel + from funasr.auto.auto_model_vllm import AutoModelVLLM + + logger.info(f"Loading ASR (vLLM): {args.model}") + _vllm_engine = AutoModelVLLM( + model=args.model, hub=args.hub, device=args.device, + dtype=getattr(args, 'dtype', 'bf16'), + tensor_parallel_size=getattr(args, 'tensor_parallel_size', 1), + gpu_memory_utilization=getattr(args, 'gpu_memory_utilization', 0.8), + max_model_len=getattr(args, 'max_model_len', 2048), + ) + + _asr_kwargs = {} + hw_file = getattr(args, 'hotword_file', '热词列表') + if hw_file and os.path.isfile(hw_file): + with open(hw_file, "r", encoding="utf-8") as hf: + hotwords = [line.strip() for line in hf if line.strip()] + _asr_kwargs["hotwords"] = hotwords + logger.info(f"Loaded {len(hotwords)} hotwords from '{hw_file}'") + + if getattr(args, 'language', None): + _asr_kwargs["language"] = args.language + logger.info(f"Language: {args.language}") + + logger.info("Loading VAD: fsmn-vad (streaming)") + _vad_model = AutoModel(model="fsmn-vad", device=args.device, disable_update=True) + + logger.info("Loading SPK: eres2netv2") + _spk_model = AutoModel(model="iic/speech_eres2netv2_sv_zh-cn_16k-common", device=args.device, disable_update=True) + + logger.info("All models ready!") + return _vllm_engine, _asr_kwargs, _vad_model, _spk_model + + +async def handle_client(websocket, args): + vllm_engine, asr_kwargs, vad_model, spk_model = load_models(args) + vad = DynamicStreamingVAD(vad_model) + spk_tracker = HybridSpeakerTracker(spk_model, args.device) + session = RealtimeASRSession(vllm_engine, asr_kwargs, vad, spk_tracker=spk_tracker) + logger.info(f"Client connected: {websocket.remote_address}") + + decode_interval = args.decode_interval + last_decode_time = 0 + + try: + async for message in websocket: + if isinstance(message, str): + cmd = message.strip() + if cmd.upper() == "START": + session.reset() + session.is_active = True + await websocket.send(json.dumps({"event": "started"})) + logger.info("Session started") + elif cmd.upper().startswith("HOTWORDS:"): + hw_str = cmd[9:] + hotwords = [w.strip() for w in hw_str.split(",") if w.strip()] + session.asr_kwargs = dict(session.asr_kwargs) + session.asr_kwargs["hotwords"] = hotwords + await websocket.send(json.dumps({"event": "hotwords_set", "hotwords": hotwords})) + logger.info(f"Hotwords set: {len(hotwords)} words") + elif cmd.upper().startswith("LANGUAGE:"): + lang = cmd[9:].strip() + session.asr_kwargs = dict(session.asr_kwargs) + session.asr_kwargs["language"] = lang if lang else None + await websocket.send(json.dumps({"event": "language_set", "language": lang})) + logger.info(f"Language set: {lang}") + elif cmd.upper() == "STOP": + if session.is_active and len(session.audio_buffer) > 0: + result = session.decode(is_final=True) + await websocket.send(json.dumps(result)) + logger.info(f"Final: {len(result['sentences'])} sentences") + session.is_active = False + await websocket.send(json.dumps({"event": "stopped"})) + elif isinstance(message, bytes) and session.is_active: + session.add_audio(message) + now = time.time() + if now - last_decode_time >= decode_interval and session.should_decode(): + result = session.decode(is_final=False) + await websocket.send(json.dumps(result)) + last_decode_time = now + + except websockets.exceptions.ConnectionClosed: + logger.info("Client disconnected") + except Exception as e: + logger.error(f"Error: {e}", exc_info=True) + + +async def main(args): + load_models(args) + logger.info(f"Server on ws://0.0.0.0:{args.port}") + async with websockets.serve( + lambda ws: handle_client(ws, args), "0.0.0.0", args.port, + max_size=10*1024*1024, + ): + await asyncio.Future() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Fun-ASR-Nano Streaming WebSocket Server") + parser.add_argument("--port", type=int, default=10095) + parser.add_argument("--model", type=str, default="FunAudioLLM/Fun-ASR-Nano-2512") + parser.add_argument("--hub", type=str, default="ms", choices=["ms", "hf"]) + parser.add_argument("--device", type=str, default="cuda:0") + parser.add_argument("--use-context", action="store_true", default=True) + parser.add_argument("--no-context", dest="use_context", action="store_false") + parser.add_argument("--decode-interval", type=float, default=0.48) + parser.add_argument("--hotword-file", type=str, default="热词列表") + parser.add_argument("--language", type=str, default=None, help="Language hint (e.g. 中文, English, 日本語)") + parser.add_argument("--dtype", type=str, default="bf16", choices=["bf16", "fp16", "fp32"]) + parser.add_argument("--tensor-parallel-size", type=int, default=1) + parser.add_argument("--gpu-memory-utilization", type=float, default=0.8) + parser.add_argument("--max-model-len", type=int, default=2048) + args = parser.parse_args() + asyncio.run(main(args)) diff --git a/examples/industrial_data_pretraining/fun_asr_nano/热词列表 b/examples/industrial_data_pretraining/fun_asr_nano/热词列表 new file mode 100644 index 000000000..d29477dde --- /dev/null +++ b/examples/industrial_data_pretraining/fun_asr_nano/热词列表 @@ -0,0 +1,4 @@ +兔子 +窝边草 +土匪 +师爷 diff --git a/funasr/auto/auto_model_vllm.py b/funasr/auto/auto_model_vllm.py new file mode 100644 index 000000000..70bdec94f --- /dev/null +++ b/funasr/auto/auto_model_vllm.py @@ -0,0 +1,300 @@ +#!/usr/bin/env python3 +# -*- encoding: utf-8 -*- +# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved. +# MIT License (https://opensource.org/licenses/MIT) + +""" +Generic vLLM inference wrapper for ALL LLM-based ASR models in FunASR. + +Applicable models (any model with audio_encoder + adaptor + LLM architecture): + - FunASRNano (Fun-ASR-Nano-2512, Fun-ASR-MLT-Nano-2512) + - LLMASR (Whisper + Qwen/Vicuna/LLaMA) + - GLMASR (GLM-ASR-Nano) + +NOT applicable (these models don't use autoregressive LLM decoding): + - Paraformer (non-autoregressive CIF predictor + attention decoder) + - SenseVoice (Whisper-like encoder-decoder, not LLM-based) + - Conformer/Transformer ASR (CTC/attention, no LLM) + - CT-Transformer (punctuation model, small transformer) + - Qwen3-ASR (uses external qwen-asr package with its own optimized inference) + +Usage: + from funasr.auto.auto_model_vllm import AutoModelVLLM + + # Works for any LLM-based ASR model + model = AutoModelVLLM( + model="FunAudioLLM/Fun-ASR-Nano-2512", + tensor_parallel_size=2, + ) + results = model.generate(["audio.wav"]) + + # Also works for LLMASR models + model = AutoModelVLLM( + model="/path/to/llm_asr_model", + tensor_parallel_size=4, + ) +""" + +import glob +import json +import logging +import os +import re +import shutil +import time +from typing import List, Optional, Union + +import numpy as np +import torch +import torch.nn as nn + +logger = logging.getLogger(__name__) + +dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32} + +# Models that use LLM and can benefit from vLLM +_LLM_BASED_MODELS = {"FunASRNano", "LLMASR", "LLMASRNAR", "GLMASR", "QwenAudioWarp"} + +# Models that CANNOT use vLLM (no autoregressive LLM) +_NON_LLM_MODELS = { + "Paraformer": "Non-autoregressive model using CIF predictor. No LLM decoding.", + "SenseVoice": "Whisper-like encoder-decoder. Not LLM-based.", + "CTTransformer": "Small punctuation model. No benefit from vLLM.", + "Conformer": "CTC/attention encoder-decoder. No LLM.", + "Qwen3ASR": "Uses external qwen-asr package with optimized inference.", +} + + +def check_vllm_applicable(model_name: str) -> bool: + """Check if a model can use vLLM inference. + + Args: + model_name: The model class name from config.yaml. + + Returns: + True if vLLM is applicable. + + Raises: + ValueError: If model explicitly cannot use vLLM, with explanation. + """ + if model_name in _LLM_BASED_MODELS: + return True + for non_llm, reason in _NON_LLM_MODELS.items(): + if non_llm in model_name: + raise ValueError( + f"Model '{model_name}' cannot use vLLM: {reason}\n" + f"vLLM only accelerates autoregressive LLM decoding. " + f"Use the standard FunASR AutoModel for this model." + ) + return False + + +def prepare_vllm_weights(model_dir: str, output_dir: str = None) -> str: + """Extract LLM weights from model.pt into vLLM-compatible format. + + Works for any model that stores LLM weights with 'llm.' prefix in model.pt + and has a config directory (e.g., Qwen3-0.6B/) with model config and tokenizer. + + Args: + model_dir: Path to the FunASR model directory. + output_dir: Where to save extracted weights. Auto-detected if None. + + Returns: + Path to vLLM-ready model directory. + """ + if output_dir is None: + # Find the LLM config directory + from omegaconf import OmegaConf + config_path = os.path.join(model_dir, "config.yaml") + if os.path.exists(config_path): + config = OmegaConf.load(config_path) + llm_conf = OmegaConf.to_container(config.get("llm_conf", {}), resolve=True) + llm_path = llm_conf.get("init_param_path", "") + if llm_path and not os.path.isabs(llm_path): + llm_path = os.path.join(model_dir, llm_path) + if os.path.isdir(llm_path): + output_dir = llm_path + "-vllm" + else: + output_dir = os.path.join(model_dir, "llm-vllm") + else: + output_dir = os.path.join(model_dir, "llm-vllm") + + # Check if already prepared + if glob.glob(os.path.join(output_dir, "*.safetensors")) or glob.glob( + os.path.join(output_dir, "model*.bin") + ): + logger.info(f"vLLM weights already at {output_dir}") + return output_dir + + os.makedirs(output_dir, exist_ok=True) + + # Find and copy LLM config/tokenizer files + from omegaconf import OmegaConf + config = OmegaConf.load(os.path.join(model_dir, "config.yaml")) + llm_conf = OmegaConf.to_container(config.get("llm_conf", {}), resolve=True) + llm_config_dir = llm_conf.get("init_param_path", "") + if llm_config_dir and not os.path.isabs(llm_config_dir): + llm_config_dir = os.path.join(model_dir, llm_config_dir) + + if os.path.isdir(llm_config_dir): + for fname in os.listdir(llm_config_dir): + src = os.path.join(llm_config_dir, fname) + dst = os.path.join(output_dir, fname) + if os.path.isfile(src) and not os.path.exists(dst): + shutil.copy2(src, dst) + + # Extract LLM weights from model.pt + model_pt = os.path.join(model_dir, "model.pt") + if not os.path.exists(model_pt): + raise FileNotFoundError(f"model.pt not found at {model_pt}") + + logger.info(f"Extracting LLM weights from {model_pt}...") + checkpoint = torch.load(model_pt, map_location="cpu") + state_dict = checkpoint.get("state_dict", checkpoint) + + llm_state = {} + for key, value in state_dict.items(): + if key.startswith("llm."): + llm_state[key[4:]] = value # Remove 'llm.' prefix + + if not llm_state: + raise RuntimeError("No LLM weights found (expected 'llm.*' prefix)") + + logger.info(f"Extracted {len(llm_state)} LLM weight tensors") + + try: + from safetensors.torch import save_file + save_path = os.path.join(output_dir, "model.safetensors") + save_file(llm_state, save_path) + index = { + "metadata": {"total_size": sum(v.numel() * v.element_size() for v in llm_state.values())}, + "weight_map": {k: "model.safetensors" for k in llm_state.keys()}, + } + with open(os.path.join(output_dir, "model.safetensors.index.json"), "w") as f: + json.dump(index, f, indent=2) + except ImportError: + torch.save(llm_state, os.path.join(output_dir, "model.bin")) + + return output_dir + + +class AutoModelVLLM: + """Generic vLLM wrapper for LLM-based ASR models. + + Automatically detects model architecture, extracts LLM weights, + loads audio components in PyTorch, and uses vLLM for generation. + + Works for: FunASRNano, LLMASR, GLMASR, and any model with + audio_encoder + audio_adaptor + LLM architecture. + + Args: + model: Model name (hub) or local directory path. + hub: "ms" (ModelScope) or "hf" (HuggingFace). + device: Device for audio encoder/adaptor. + dtype: Compute dtype ("bf16", "fp16", "fp32"). + tensor_parallel_size: GPUs for vLLM tensor parallelism. + gpu_memory_utilization: GPU memory fraction for vLLM. + max_model_len: Maximum sequence length. + + Example: + >>> model = AutoModelVLLM(model="FunAudioLLM/Fun-ASR-Nano-2512") + >>> results = model.generate(["audio.wav"], language="中文") + """ + + def __init__( + self, + model: str, + hub: str = "ms", + device: str = "cuda:0", + dtype: str = "bf16", + tensor_parallel_size: int = 1, + gpu_memory_utilization: float = 0.8, + max_model_len: int = 4096, + enforce_eager: bool = False, + **kwargs, + ): + # Resolve model directory + if os.path.isdir(model): + self.model_dir = model + else: + if hub in ("ms", "modelscope"): + from modelscope.hub.snapshot_download import snapshot_download + self.model_dir = snapshot_download(model, revision=kwargs.get("revision", "master")) + elif hub in ("hf", "huggingface"): + from huggingface_hub import snapshot_download + self.model_dir = snapshot_download(model) + else: + raise ValueError(f"Unsupported hub: {hub}") + + # Check model type + from omegaconf import OmegaConf + config = OmegaConf.load(os.path.join(self.model_dir, "config.yaml")) + self.model_type = config.get("model", "unknown") + check_vllm_applicable(self.model_type) + + self.device = device + self.dtype = dtype + self.torch_dtype = dtype_map.get(dtype, torch.bfloat16) + + # Use the specialized implementation if available + if self.model_type == "FunASRNano": + from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM + self._engine = FunASRNanoVLLM( + model_dir=self.model_dir, device=device, dtype=dtype, + tensor_parallel_size=tensor_parallel_size, + gpu_memory_utilization=gpu_memory_utilization, + max_model_len=max_model_len, enforce_eager=enforce_eager, + **kwargs, + ) + elif self.model_type in ("LLMASR", "LLMASRNAR"): + self._engine = self._build_llmasr_engine( + config, tensor_parallel_size, gpu_memory_utilization, + max_model_len, enforce_eager, **kwargs, + ) + else: + # Generic fallback using the FunASRNano approach + # (works for any model with audio_encoder + adaptor + LLM) + from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM + self._engine = FunASRNanoVLLM( + model_dir=self.model_dir, device=device, dtype=dtype, + tensor_parallel_size=tensor_parallel_size, + gpu_memory_utilization=gpu_memory_utilization, + max_model_len=max_model_len, enforce_eager=enforce_eager, + **kwargs, + ) + + def _build_llmasr_engine(self, config, tensor_parallel_size, gpu_memory_utilization, + max_model_len, enforce_eager, **kwargs): + """Build vLLM engine for LLMASR models (Whisper + Qwen/Vicuna).""" + from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM + + # LLMASR follows same pattern as FunASRNano + return FunASRNanoVLLM( + model_dir=self.model_dir, device=self.device, dtype=self.dtype, + tensor_parallel_size=tensor_parallel_size, + gpu_memory_utilization=gpu_memory_utilization, + max_model_len=max_model_len, enforce_eager=enforce_eager, + **kwargs, + ) + + def generate(self, inputs, **kwargs): + """Run ASR inference. + + Args: + inputs: Audio file path(s), numpy arrays, or tensors. + **kwargs: Model-specific parameters (language, hotwords, etc.) + + Returns: + List of result dicts with "key" and "text" fields. + """ + return self._engine.generate(inputs, **kwargs) + + @classmethod + def supported_models(cls): + """Return dict of model types and their vLLM support status.""" + info = {} + for m in _LLM_BASED_MODELS: + info[m] = {"supported": True, "reason": "LLM-based, autoregressive generation"} + for m, reason in _NON_LLM_MODELS.items(): + info[m] = {"supported": False, "reason": reason} + return info diff --git a/funasr/models/fsmn_vad_streaming/dynamic_vad.py b/funasr/models/fsmn_vad_streaming/dynamic_vad.py new file mode 100644 index 000000000..bd14515c9 --- /dev/null +++ b/funasr/models/fsmn_vad_streaming/dynamic_vad.py @@ -0,0 +1,222 @@ +#!/usr/bin/env python3 +"""DynamicStreamingVAD — 动态阈值流式 VAD 封装。 + +在 fsmn-vad 基础上，根据当前语音段的累积时长动态调整静音切分阈值： +短句等待更长静音（避免切碎），长句快速切分（避免堆积）。 + +支持流式（逐帧喂入）和非流式（一次性处理完整音频）两种调用方式。 + +Usage (流式): + from funasr import AutoModel + from funasr.models.fsmn_vad_streaming.dynamic_vad import DynamicStreamingVAD + + vad_model = AutoModel(model="fsmn-vad", device="cuda:0") + vad = DynamicStreamingVAD(vad_model) + + for audio_chunk in audio_stream: + segments = vad.feed(audio_chunk) + for seg in segments: + print(f"Speech: {seg[0]}-{seg[1]}ms") + + # 结束时 + final_segments = vad.finalize() + +Usage (非流式): + segments = vad.process(full_audio_tensor) + for seg in segments: + print(f"Speech: {seg[0]}-{seg[1]}ms") +""" + +from typing import List, Optional, Tuple + +import torch +import numpy as np + + +# 默认动态阈值配置：(累积时长上限ms, 静音阈值ms) +DEFAULT_SILENCE_SCHEDULE = [ + (5000, 2000), + (10000, 1500), + (15000, 1000), + (30000, 800), + (45000, 400), + (float('inf'), 100), +] + + +class DynamicStreamingVAD: + """动态阈值流式 VAD。 + + 在 fsmn-vad 的流式推理基础上，根据当前语音段已累积的时长 + 动态调整静音切分阈值，实现「短句不切碎、长句快切分」。 + + Args: + vad_model: FunASR AutoModel 加载的 fsmn-vad 模型实例。 + chunk_size_ms: 每次喂入 VAD 的 chunk 大小（毫秒），默认 60。 + speech_noise_thres: 语音/噪声判别阈值，默认 0.5。 + speech_to_sil_thres_ms: 语音转静音的基础时间（毫秒），默认 150。 + silence_schedule: 动态阈值配置表，格式为 + [(累积时长上限ms, 对应的静音阈值ms), ...]。 + 当累积时长 <= 上限时，使用对应的静音阈值。 + 默认值适合实时对话场景。设为 None 禁用动态调整（使用固定阈值）。 + sample_rate: 采样率，默认 16000。 + + Example: + # 自定义阈值：更激进的切分 + vad = DynamicStreamingVAD( + vad_model, + silence_schedule=[ + (3000, 1500), + (8000, 800), + (15000, 400), + (float('inf'), 200), + ], + ) + """ + + def __init__( + self, + vad_model, + chunk_size_ms: int = 60, + speech_noise_thres: float = 0.5, + speech_to_sil_thres_ms: int = 150, + silence_schedule: Optional[List[Tuple[float, int]]] = None, + sample_rate: int = 16000, + ): + self.model = vad_model + self.chunk_size_ms = chunk_size_ms + self.speech_noise_thres = speech_noise_thres + self.speech_to_sil_thres_ms = speech_to_sil_thres_ms + self.silence_schedule = silence_schedule if silence_schedule is not None else DEFAULT_SILENCE_SCHEDULE + self.sample_rate = sample_rate + + self.cache = {} + self.confirmed_segments: List[List[int]] = [] + self.current_speech_start: Optional[int] = None + self.accumulated_since_cut_ms: int = 0 + + def _get_silence_threshold(self) -> int: + """根据当前累积时长，从 schedule 中查询静音阈值。""" + for limit_ms, silence_ms in self.silence_schedule: + if self.accumulated_since_cut_ms <= limit_ms: + return silence_ms + return self.silence_schedule[-1][1] + + def _apply_dynamic_threshold(self): + """将动态阈值应用到 VAD 内部 cache。""" + if "stats" not in self.cache: + return + stats = self.cache["stats"] + stats.speech_noise_thres = self.speech_noise_thres + desired_silence_ms = self._get_silence_threshold() + stats.max_end_sil_frame_cnt_thresh = max(desired_silence_ms - self.speech_to_sil_thres_ms, 0) + + def feed(self, audio_chunk: torch.Tensor, is_final: bool = False) -> List[List[int]]: + """喂入一段音频，返回新确认的语音段。 + + Args: + audio_chunk: 音频数据（float32 tensor，16kHz）。 + 可以是任意长度，内部按 chunk_size_ms 处理。 + is_final: 是否为最后一段音频。设为 True 时会强制结束当前语音段。 + + Returns: + 新确认的语音段列表，每段为 [start_ms, end_ms]。 + 仅在检测到语音结束时返回非空列表。 + """ + if audio_chunk.dim() > 1: + audio_chunk = audio_chunk.squeeze() + + chunk_samples = len(audio_chunk) + self.accumulated_since_cut_ms += int(chunk_samples * 1000 / self.sample_rate) + + self._apply_dynamic_threshold() + + res = self.model.generate( + input=[audio_chunk], cache=self.cache, + is_final=is_final, chunk_size=self.chunk_size_ms, + ) + + signals = res[0].get("value", []) + new_confirmed = [] + + for sig in signals: + if sig[0] >= 0 and sig[1] == -1: + self.current_speech_start = sig[0] + elif sig[0] == -1 and sig[1] >= 0: + start = self.current_speech_start if self.current_speech_start is not None else 0 + seg = [start, sig[1]] + self.confirmed_segments.append(seg) + new_confirmed.append(seg) + self.current_speech_start = None + self.accumulated_since_cut_ms = 0 + elif sig[0] >= 0 and sig[1] >= 0: + self.confirmed_segments.append(sig) + new_confirmed.append(sig) + self.current_speech_start = None + self.accumulated_since_cut_ms = 0 + + return new_confirmed + + def finalize(self) -> List[List[int]]: + """结束流式处理，返回最后可能未结束的语音段。 + + 调用此方法后，VAD 状态会被重置。 + 如果当前有正在进行的语音段，会被强制结束。 + + Returns: + 最后确认的语音段列表。 + """ + # Feed empty with is_final=True to flush + empty = torch.zeros(int(self.sample_rate * 0.01), dtype=torch.float32) + return self.feed(empty, is_final=True) + + def process(self, audio: torch.Tensor) -> List[List[int]]: + """非流式接口：一次性处理完整音频，返回所有语音段。 + + Args: + audio: 完整音频（float32 tensor，16kHz）。 + + Returns: + 所有检测到的语音段 [[start_ms, end_ms], ...]。 + """ + self.reset() + + if isinstance(audio, np.ndarray): + audio = torch.from_numpy(audio).float() + if audio.dim() > 1: + audio = audio.squeeze() + + # 分 chunk 喂入 + chunk_samples = int(self.sample_rate * self.chunk_size_ms / 1000) + total = len(audio) + all_segments = [] + + for i in range(0, total, chunk_samples): + chunk = audio[i:i + chunk_samples] + is_last = (i + chunk_samples >= total) + segs = self.feed(chunk, is_final=is_last) + all_segments.extend(segs) + + return all_segments + + @property + def is_speaking(self) -> bool: + """当前是否在语音状态中。""" + return self.current_speech_start is not None + + @property + def current_duration_ms(self) -> int: + """当前段已累积的时长（毫秒）。""" + return self.accumulated_since_cut_ms + + @property + def current_threshold_ms(self) -> int: + """当前使用的静音阈值（毫秒）。""" + return self._get_silence_threshold() + + def reset(self): + """重置所有状态，开始新一轮检测。""" + self.cache = {} + self.confirmed_segments = [] + self.current_speech_start = None + self.accumulated_since_cut_ms = 0 diff --git a/funasr/models/fsmn_vad_streaming/model.py b/funasr/models/fsmn_vad_streaming/model.py index 3acd481e1..dc0c394c1 100644 --- a/funasr/models/fsmn_vad_streaming/model.py +++ b/funasr/models/fsmn_vad_streaming/model.py @@ -18,6 +18,26 @@ from typing import List, Tuple, Dict, Any, Optional from funasr.utils.datadir_writer import DatadirWriter from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank +# Dynamic silence threshold schedule: (accumulated_speech_ms, silence_threshold_ms) +# Streaming mode: longer initial patience (waiting for user to continue) +STREAMING_SILENCE_SCHEDULE = [ + (5000, 2000), + (10000, 1500), + (15000, 1000), + (30000, 800), + (45000, 400), + (float("inf"), 100), +] + +# Non-streaming mode: shorter patience (audio already complete, cut by natural pauses) +DEFAULT_SILENCE_SCHEDULE = [ + (5000, 800), + (10000, 600), + (20000, 500), + (30000, 400), + (float("inf"), 300), +] + class VadStateMachine(Enum): kVadInStateStartPointNotDetected = 1 @@ -914,10 +934,32 @@ class FsmnVADStreaming(nn.Module): n = int(len(audio_sample) // chunk_stride_samples + int(_is_final)) m = int(len(audio_sample) % chunk_stride_samples * (1 - int(_is_final))) segments = [] + # Dynamic silence threshold + dynamic_silence = kwargs.get("dynamic_silence", True) + silence_schedule = kwargs.get("silence_schedule", DEFAULT_SILENCE_SCHEDULE) + speech_to_sil_ms = self.vad_opts.speech_to_sil_time_thres + accumulated_ms = cache.get("_dynamic_accumulated_ms", 0) + in_speech = cache.get("_dynamic_in_speech", False) + for i in range(n): kwargs["is_final"] = _is_final and i == n - 1 audio_sample_i = audio_sample[i * chunk_stride_samples : (i + 1) * chunk_stride_samples] + # Apply dynamic silence threshold (only accumulate while in speech) + if dynamic_silence and "stats" in cache: + vad_state = cache["stats"].vad_state_machine + # VadStateMachine: 2=kVadInStateStartPointNotDetected, 3=kVadInStateInSpeechSegment + if vad_state.value == 2 or in_speech: # kVadInStateInSpeechSegment + accumulated_ms += chunk_size + in_speech = True + for limit_ms, silence_ms in silence_schedule: + if accumulated_ms <= limit_ms: + cache["stats"].max_end_sil_frame_cnt_thresh = max(silence_ms - speech_to_sil_ms, 0) + cache["stats"].speech_noise_thres = 0.5 + break + cache["_dynamic_accumulated_ms"] = accumulated_ms + cache["_dynamic_in_speech"] = in_speech + # extract fbank feats speech, speech_lengths = extract_fbank( [audio_sample_i], @@ -944,6 +986,11 @@ class FsmnVADStreaming(nn.Module): segments_i = self.forward(**batch) if len(segments_i) > 0: segments.extend(*segments_i) + if dynamic_silence: + accumulated_ms = 0 + in_speech = False + cache["_dynamic_accumulated_ms"] = 0 + cache["_dynamic_in_speech"] = False cache["prev_samples"] = audio_sample[-m:] if m > 0 else torch.empty(0) if _is_final: diff --git a/funasr/models/fun_asr_nano/inference_vllm.py b/funasr/models/fun_asr_nano/inference_vllm.py new file mode 100644 index 000000000..927da7b5c --- /dev/null +++ b/funasr/models/fun_asr_nano/inference_vllm.py @@ -0,0 +1,650 @@ +#!/usr/bin/env python3 +# -*- encoding: utf-8 -*- +# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved. +# MIT License (https://opensource.org/licenses/MIT) + +""" +Fun-ASR-Nano vLLM inference engine. + +Uses vLLM for high-throughput LLM decoding while keeping the audio encoder +and adaptor in PyTorch. Supports batch inference and tensor-parallel for +multi-GPU acceleration. + +Usage: + from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM + + engine = FunASRNanoVLLM.from_pretrained( + model="FunAudioLLM/Fun-ASR-Nano-2512", + tensor_parallel_size=2, + ) + results = engine.generate(["audio1.wav", "audio2.wav"]) +""" + +import glob +import json +import logging +import os +import re +import shutil +import time +from typing import List, Optional, Union + +import numpy as np +import torch +import torch.nn as nn + +logger = logging.getLogger(__name__) + +dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32} + + +def prepare_vllm_model_dir(model_dir: str, output_dir: str = None) -> str: + """Extract LLM weights from Fun-ASR-Nano model.pt and save in HuggingFace format. + + Fun-ASR-Nano stores all weights (audio encoder + adaptor + LLM) in a single + model.pt file. vLLM needs the LLM weights in standard HuggingFace format. + This function extracts LLM weights and saves them alongside the config/tokenizer + files from the Qwen3-0.6B subdirectory. + + Args: + model_dir: Path to the Fun-ASR-Nano model directory. + output_dir: Where to save the extracted LLM. Defaults to model_dir/Qwen3-0.6B-vllm. + + Returns: + Path to the directory containing the vLLM-ready LLM model. + """ + if output_dir is None: + output_dir = os.path.join(model_dir, "Qwen3-0.6B-vllm") + + # Check if already prepared + safetensors_files = glob.glob(os.path.join(output_dir, "*.safetensors")) + bin_files = glob.glob(os.path.join(output_dir, "model*.bin")) + if safetensors_files or bin_files: + logger.info(f"vLLM model already prepared at {output_dir}") + return output_dir + + os.makedirs(output_dir, exist_ok=True) + + # Copy config and tokenizer from Qwen3-0.6B + qwen_dir = os.path.join(model_dir, "Qwen3-0.6B") + if not os.path.isdir(qwen_dir): + raise FileNotFoundError(f"Qwen3-0.6B config directory not found at {qwen_dir}") + + for fname in os.listdir(qwen_dir): + src = os.path.join(qwen_dir, fname) + dst = os.path.join(output_dir, fname) + if os.path.isfile(src) and not os.path.exists(dst): + shutil.copy2(src, dst) + + # Load model.pt and extract LLM weights + model_pt = os.path.join(model_dir, "model.pt") + if not os.path.exists(model_pt): + raise FileNotFoundError( + f"model.pt not found at {model_pt}. Make sure the model is fully downloaded." + ) + + logger.info(f"Loading model.pt from {model_pt}...") + checkpoint = torch.load(model_pt, map_location="cpu") + if "state_dict" in checkpoint: + state_dict = checkpoint["state_dict"] + else: + state_dict = checkpoint + + # Extract LLM weights (prefixed with "llm.") + llm_state = {} + for key, value in state_dict.items(): + if key.startswith("llm."): + new_key = key[len("llm."):] + llm_state[new_key] = value + + if not llm_state: + raise RuntimeError("No LLM weights found in model.pt (expected prefix 'llm.')") + + logger.info(f"Extracted {len(llm_state)} LLM weight tensors") + + # Save in safetensors format (preferred by vLLM) + try: + from safetensors.torch import save_file + + save_path = os.path.join(output_dir, "model.safetensors") + save_file(llm_state, save_path) + logger.info(f"Saved LLM weights to {save_path}") + + # Create model index + index = { + "metadata": {"total_size": sum(v.numel() * v.element_size() for v in llm_state.values())}, + "weight_map": {k: "model.safetensors" for k in llm_state.keys()}, + } + with open(os.path.join(output_dir, "model.safetensors.index.json"), "w") as f: + json.dump(index, f, indent=2) + except ImportError: + save_path = os.path.join(output_dir, "model.bin") + torch.save(llm_state, save_path) + logger.info(f"Saved LLM weights to {save_path} (install safetensors for faster loading)") + + return output_dir + + +class FunASRNanoVLLM: + """Fun-ASR-Nano with vLLM backend for high-throughput inference. + + Architecture: + Audio -> WavFrontend -> SenseVoiceEncoder -> AudioAdaptor -> audio embeddings + Text tokens -> LLM embedding layer -> text embeddings + Combined embeddings -> vLLM (Qwen3-0.6B) -> generated text + + The audio encoder and adaptor run in PyTorch on a single GPU, + while vLLM handles the LLM inference with optional tensor parallelism. + + Args: + model_dir: Path to the Fun-ASR-Nano model directory. + device: Device for audio encoder/adaptor (e.g. "cuda:0"). + dtype: Dtype for audio processing ("bf16", "fp16", "fp32"). + tensor_parallel_size: Number of GPUs for vLLM tensor parallelism. + gpu_memory_utilization: Fraction of GPU memory for vLLM KV cache. + max_model_len: Maximum sequence length for vLLM. + enforce_eager: Disable CUDA graph for debugging. + + Example: + >>> engine = FunASRNanoVLLM( + ... model_dir="/path/to/Fun-ASR-Nano-2512", + ... tensor_parallel_size=2, + ... ) + >>> results = engine.generate(["audio1.wav", "audio2.wav"]) + >>> for r in results: + ... print(r["text"]) + """ + + def __init__( + self, + model_dir: str, + device: str = "cuda:0", + dtype: str = "bf16", + tensor_parallel_size: int = 1, + gpu_memory_utilization: float = 0.8, + max_model_len: int = 2048, + enforce_eager: bool = False, + **kwargs, + ): + from vllm import LLM, SamplingParams + from vllm.inputs.data import EmbedsPrompt + + self.device = device + self.dtype = dtype + self.torch_dtype = dtype_map.get(dtype, torch.bfloat16) + self.model_dir = model_dir + + # Step 1: Prepare LLM weights for vLLM (extract from model.pt if needed) + vllm_model_dir = prepare_vllm_model_dir(model_dir) + + # Step 2: Load audio components (encoder + adaptor + frontend) + self._load_audio_components(model_dir, **kwargs) + + # Step 3: Initialize vLLM engine + logger.info(f"Initializing vLLM with model: {vllm_model_dir}") + logger.info(f" tensor_parallel_size={tensor_parallel_size}") + logger.info(f" gpu_memory_utilization={gpu_memory_utilization}") + + vllm_kwargs = kwargs.get("vllm_kwargs", {}) + self.vllm_engine = LLM( + enable_prompt_embeds=True, + model=vllm_model_dir, + tensor_parallel_size=tensor_parallel_size, + gpu_memory_utilization=gpu_memory_utilization, + max_model_len=max_model_len, + enforce_eager=enforce_eager, + dtype={"bf16": "bfloat16", "fp16": "float16", "fp32": "auto"}.get(dtype, dtype), + trust_remote_code=True, + **vllm_kwargs, + ) + + # Step 4: Get tokenizer and LLM embedding layer + self.tokenizer = self.vllm_engine.get_tokenizer() + self._load_embedding_layer(model_dir) + + def _load_audio_components(self, model_dir: str, **kwargs): + """Load audio encoder, adaptor, frontend, and CTC from checkpoint.""" + from omegaconf import OmegaConf + from funasr.register import tables + + config_path = os.path.join(model_dir, "config.yaml") + config = OmegaConf.load(config_path) + self._config = OmegaConf.to_container(config, resolve=True) + + # --- Frontend --- + frontend_class = tables.frontend_classes.get(config["frontend"]) + frontend_conf = OmegaConf.to_container(config.get("frontend_conf", {}), resolve=True) + cmvn_file = frontend_conf.get("cmvn_file") + if cmvn_file and not os.path.isabs(cmvn_file): + frontend_conf["cmvn_file"] = os.path.join(model_dir, cmvn_file) + self.frontend = frontend_class(**frontend_conf) + self.frontend.eval() + + # --- Audio Encoder --- + encoder_conf = OmegaConf.to_container(config.get("audio_encoder_conf", {}), resolve=True) + hub = encoder_conf.get("hub", None) + if hub == "ms": + from funasr import AutoModel as FunAutoModel + + enc_model = FunAutoModel( + model=config["audio_encoder"], model_revision="master", disable_update=True + ) + self.audio_encoder_output_size = ( + enc_model.model.encoder_output_size + if hasattr(enc_model.model, "encoder_output_size") + else -1 + ) + self.audio_encoder = ( + enc_model.model.model.encoder + if hasattr(enc_model.model, "model") + else enc_model.model.encoder + ) + else: + encoder_class = tables.encoder_classes.get(config["audio_encoder"]) + input_size = self.frontend.output_size() + self.audio_encoder = encoder_class(input_size=input_size, **encoder_conf) + self.audio_encoder_output_size = self.audio_encoder.output_size() + + self.audio_encoder.eval() + for p in self.audio_encoder.parameters(): + p.requires_grad = False + + # --- Audio Adaptor --- + adaptor_conf = OmegaConf.to_container(config.get("audio_adaptor_conf", {}), resolve=True) + adaptor_class = tables.adaptor_classes.get(config["audio_adaptor"]) + if self.audio_encoder_output_size > 0: + adaptor_conf["encoder_dim"] = self.audio_encoder_output_size + self.audio_adaptor = adaptor_class(**adaptor_conf) + self.audio_adaptor.eval() + for p in self.audio_adaptor.parameters(): + p.requires_grad = False + self.use_low_frame_rate = adaptor_conf.get("use_low_frame_rate", False) + + # --- CTC Decoder (optional, for timestamps) --- + self.ctc_decoder = None + self.ctc = None + self.ctc_tokenizer = None + self.blank_id = None + + ctc_decoder_name = self._config.get("ctc_decoder", None) + if ctc_decoder_name: + ctc_decoder_class = tables.adaptor_classes.get(ctc_decoder_name) + ctc_decoder_conf = self._config.get("ctc_decoder_conf", {}) + if self.audio_encoder_output_size > 0: + ctc_decoder_conf["encoder_dim"] = self.audio_encoder_output_size + self.ctc_decoder = ctc_decoder_class(**ctc_decoder_conf) + self.ctc_decoder.eval() + for p in self.ctc_decoder.parameters(): + p.requires_grad = False + + from funasr.models.fun_asr_nano.ctc import CTC + + ctc_conf = self._config.get("ctc_conf", {}) + ctc_vocab_size = self._config.get("ctc_vocab_size", 60515) + self.blank_id = ctc_conf.get("blank_id", ctc_vocab_size - 1) + self.ctc = CTC( + odim=ctc_vocab_size, + encoder_output_size=self.audio_encoder_output_size, + blank_id=self.blank_id, + **ctc_conf, + ) + + # CTC tokenizer + ds_conf = self._config.get("dataset_conf", {}) + ctc_tokenizer_name = ds_conf.get("ctc_tokenizer", None) + ctc_tokenizer_conf = ds_conf.get("ctc_tokenizer_conf", {}) + if ctc_tokenizer_name: + ctc_tokenizer_class = tables.tokenizer_classes.get(ctc_tokenizer_name) + vocab_path = ctc_tokenizer_conf.get("vocab_path") + if vocab_path is None or not os.path.isabs(vocab_path): + multilingual_path = os.path.join(model_dir, "multilingual.tiktoken") + if os.path.exists(multilingual_path): + ctc_tokenizer_conf["vocab_path"] = multilingual_path + elif vocab_path and not os.path.isabs(vocab_path): + ctc_tokenizer_conf["vocab_path"] = os.path.join(model_dir, vocab_path) + self.ctc_tokenizer = ctc_tokenizer_class(**ctc_tokenizer_conf) + + # --- Load weights from model.pt --- + model_pt = os.path.join(model_dir, "model.pt") + if os.path.exists(model_pt): + logger.info(f"Loading audio component weights from {model_pt}") + checkpoint = torch.load(model_pt, map_location="cpu") + state_dict = checkpoint.get("state_dict", checkpoint) + + # Audio encoder + enc_state = { + k[len("audio_encoder."):]: v + for k, v in state_dict.items() + if k.startswith("audio_encoder.") + } + if enc_state: + self.audio_encoder.load_state_dict(enc_state, strict=False) + logger.info(f" Loaded audio_encoder: {len(enc_state)} params") + + # Audio adaptor + adp_state = { + k[len("audio_adaptor."):]: v + for k, v in state_dict.items() + if k.startswith("audio_adaptor.") + } + if adp_state: + self.audio_adaptor.load_state_dict(adp_state, strict=False) + logger.info(f" Loaded audio_adaptor: {len(adp_state)} params") + + # CTC decoder + if self.ctc_decoder is not None: + ctc_dec_state = { + k[len("ctc_decoder."):]: v + for k, v in state_dict.items() + if k.startswith("ctc_decoder.") + } + if ctc_dec_state: + self.ctc_decoder.load_state_dict(ctc_dec_state, strict=False) + ctc_state = { + k[len("ctc."):]: v + for k, v in state_dict.items() + if k.startswith("ctc.") and not k.startswith("ctc_decoder.") + } + if ctc_state: + self.ctc.load_state_dict(ctc_state, strict=False) + + # Move to device + self.audio_encoder = self.audio_encoder.to(self.device, dtype=torch.float32) + self.audio_adaptor = self.audio_adaptor.to(self.device, dtype=self.torch_dtype) + if self.ctc_decoder is not None: + self.ctc_decoder = self.ctc_decoder.to(self.device, dtype=torch.float32) + self.ctc = self.ctc.to(self.device, dtype=torch.float32) + + def _load_embedding_layer(self, model_dir: str): + """Load the LLM embedding layer for text token embedding computation.""" + model_pt = os.path.join(model_dir, "model.pt") + checkpoint = torch.load(model_pt, map_location="cpu") + state_dict = checkpoint.get("state_dict", checkpoint) + + # Look for embedding weights + embed_key = None + for key in state_dict.keys(): + if "embed_tokens.weight" in key and key.startswith("llm."): + embed_key = key + break + + if embed_key is None: + raise RuntimeError("Could not find LLM embedding weights in model.pt") + + embed_weight = state_dict[embed_key] + self.embed_tokens = nn.Embedding.from_pretrained(embed_weight, freeze=True) + self.embed_tokens = self.embed_tokens.to(self.device, dtype=self.torch_dtype) + logger.info(f"Loaded embedding layer: {embed_weight.shape}") + + @torch.no_grad() + def _encode_audio(self, audio_input: Union[str, torch.Tensor, np.ndarray]): + """Encode audio through frontend -> encoder -> adaptor. + + Returns: + adaptor_out: (1, T', D_llm) audio embeddings for LLM input + adaptor_out_lens: (1,) lengths + encoder_out: (1, T, D_enc) encoder output for CTC + encoder_out_lens: (1,) encoder output lengths + """ + from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank + + if isinstance(audio_input, str): + data_src = load_audio_text_image_video(audio_input, fs=self.frontend.fs) + elif isinstance(audio_input, np.ndarray): + data_src = torch.from_numpy(audio_input).float() + elif isinstance(audio_input, torch.Tensor): + data_src = audio_input.float() + else: + raise ValueError(f"Unsupported audio input type: {type(audio_input)}") + + speech, speech_lengths = extract_fbank( + data_src, data_type="sound", frontend=self.frontend, is_final=True + ) + speech = speech.to(self.device, dtype=torch.float32) + speech_lengths = speech_lengths.to(self.device) + + encoder_out, encoder_out_lens = self.audio_encoder(speech, speech_lengths) + encoder_out_for_adaptor = encoder_out.to(dtype=self.torch_dtype) + adaptor_out, adaptor_out_lens = self.audio_adaptor(encoder_out_for_adaptor, encoder_out_lens) + return adaptor_out, adaptor_out_lens, encoder_out, encoder_out_lens + + def _build_prompt_text( + self, + hotwords: List[str] = None, + language: str = None, + itn: bool = True, + ) -> str: + """Build the ASR prompt string.""" + hotwords = hotwords or [] + if len(hotwords) > 0: + hotwords_str = ", ".join(hotwords) + prompt = ( + "请结合上下文信息，更加准确地完成语音转写任务。" + "如果没有相关信息，我们会留空。\n\n\n**上下文信息：**\n\n\n" + ) + prompt += f"热词列表：[{hotwords_str}]\n" + else: + prompt = "" + if language is None: + prompt += "语音转写" + else: + prompt += f"语音转写成{language}" + if not itn: + prompt += "，不进行文本规整" + return prompt + "：" + + @torch.no_grad() + def _build_input_embeds( + self, + audio_embeds: torch.Tensor, + audio_embed_lens: torch.Tensor, + hotwords: List[str] = None, + language: str = None, + itn: bool = True, + system_prompt: str = "You are a helpful assistant.", + ) -> torch.Tensor: + """Build the full input embedding sequence with audio inserted. + + Returns: + Tensor of shape (seq_len, D_llm) + """ + prompt = self._build_prompt_text(hotwords, language, itn) + + # ChatML format + prefix_text = ( + f"<|im_start|>system\n{system_prompt}<|im_end|>\n" + f"<|im_start|>user\n{prompt}<|startofspeech|>" + ) + suffix_text = "<|endofspeech|><|im_end|>\n<|im_start|>assistant\n\n\n\n\n" + + # Tokenize + prefix_ids = self.tokenizer.encode(prefix_text, add_special_tokens=False) + suffix_ids = self.tokenizer.encode(suffix_text, add_special_tokens=False) + + # Embed text tokens + prefix_tensor = torch.tensor(prefix_ids, dtype=torch.long, device=self.device) + suffix_tensor = torch.tensor(suffix_ids, dtype=torch.long, device=self.device) + prefix_embeds = self.embed_tokens(prefix_tensor) + suffix_embeds = self.embed_tokens(suffix_tensor) + + # Audio embeddings + audio_len = audio_embed_lens[0].item() + audio_emb = audio_embeds[0, :audio_len, :] + + # Concat: [prefix_text_emb | audio_emb | suffix_text_emb] + inputs_embeds = torch.cat([prefix_embeds, audio_emb, suffix_embeds], dim=0) + return inputs_embeds + + def generate( + self, + inputs: Union[str, List[str], np.ndarray, torch.Tensor, List], + hotwords: List[str] = None, + language: str = None, + itn: bool = True, + max_new_tokens: int = 512, + temperature: float = 0.0, + top_p: float = 1.0, + top_k: int = -1, + repetition_penalty: float = 1.0, + **kwargs, + ) -> List[dict]: + """Run batch ASR inference using vLLM. + + Args: + inputs: Audio input(s). Accepts: + - str: single file path + - List[str]: batch of file paths + - np.ndarray / torch.Tensor: raw audio samples (16kHz) + hotwords: Keywords to boost recognition accuracy. + language: Language hint (e.g. "中文", "英文", "日文"). + itn: Apply inverse text normalization (default True). + max_new_tokens: Maximum tokens to generate per sample. + temperature: Sampling temperature (0 = greedy decoding). + top_p: Nucleus sampling parameter. + top_k: Top-k sampling (-1 = disabled). + repetition_penalty: Repetition penalty factor. + + Returns: + List of result dicts: [{"key": str, "text": str, "timestamps": [...]}] + """ + from vllm import SamplingParams + from vllm.inputs.data import EmbedsPrompt + + if isinstance(inputs, (str, np.ndarray, torch.Tensor)): + inputs = [inputs] + + sampling_params = SamplingParams( + max_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k if top_k > 0 else -1, + repetition_penalty=repetition_penalty, + skip_special_tokens=True, + ) + + # Encode all audio samples and build embedding prompts + prompts = [] + encoder_outputs = [] + + t0 = time.perf_counter() + for audio_input in inputs: + adaptor_out, adaptor_out_lens, encoder_out, encoder_out_lens = self._encode_audio( + audio_input + ) + encoder_outputs.append((encoder_out, encoder_out_lens)) + + input_embeds = self._build_input_embeds( + adaptor_out, adaptor_out_lens, hotwords, language, itn + ) + # vLLM EmbedsPrompt expects float32 or the model's dtype + prompts.append(EmbedsPrompt(prompt_embeds=input_embeds.float())) + + t1 = time.perf_counter() + logger.info(f"Audio encoding: {len(inputs)} samples in {t1 - t0:.3f}s") + + # vLLM batch generation + outputs = self.vllm_engine.generate(prompts, sampling_params, use_tqdm=len(inputs) > 1) + + t2 = time.perf_counter() + logger.info(f"vLLM generation: {t2 - t1:.3f}s") + + # Process results + results = [] + for i, output in enumerate(outputs): + text = output.outputs[0].text + if not text and output.outputs[0].token_ids: + text = self.tokenizer.decode(list(output.outputs[0].token_ids), skip_special_tokens=True) + text_clean = re.sub(r"\s+", " ", text.replace("/sil", " ")).strip() + + key = ( + os.path.splitext(os.path.basename(inputs[i]))[0] + if isinstance(inputs[i], str) + else f"sample_{i}" + ) + result = {"key": key, "text": text_clean} + + # Timestamps via CTC forced alignment + if self.ctc_decoder is not None and self.ctc_tokenizer is not None: + try: + timestamps = self._compute_timestamps( + encoder_outputs[i][0], encoder_outputs[i][1], text_clean + ) + if timestamps: + result["timestamps"] = timestamps + except Exception as e: + logger.debug(f"Timestamp computation failed for {key}: {e}") + + results.append(result) + + return results + + @torch.no_grad() + def _compute_timestamps(self, encoder_out, encoder_out_lens, text): + """CTC forced alignment for character-level timestamps.""" + from funasr.models.fun_asr_nano.tools.utils import forced_align + + decoder_out, decoder_out_lens = self.ctc_decoder(encoder_out, encoder_out_lens) + ctc_logits = self.ctc.log_softmax(decoder_out) + x = ctc_logits[0, : encoder_out_lens[0].item(), :] + + target_ids = torch.tensor(self.ctc_tokenizer.encode(text), dtype=torch.int64) + if len(target_ids) == 0: + return [] + + timestamps = forced_align(x, target_ids, self.blank_id) + for ts in timestamps: + ts["token"] = self.ctc_tokenizer.decode([ts["token"]]) + ts["start_time"] = ts["start_time"] * 6 * 10 / 1000 + ts["end_time"] = ts["end_time"] * 6 * 10 / 1000 + return timestamps + + @classmethod + def from_pretrained( + cls, + model: str = "FunAudioLLM/Fun-ASR-Nano-2512", + hub: str = "ms", + device: str = "cuda:0", + dtype: str = "bf16", + tensor_parallel_size: int = 1, + gpu_memory_utilization: float = 0.8, + max_model_len: int = 2048, + **kwargs, + ) -> "FunASRNanoVLLM": + """Load model from hub or local path. + + Args: + model: Model name or local directory path. + hub: "ms" (ModelScope) or "hf" (HuggingFace). + device: Device for audio encoder/adaptor. + dtype: Compute dtype ("bf16", "fp16", "fp32"). + tensor_parallel_size: GPUs for vLLM tensor parallel. + gpu_memory_utilization: GPU memory fraction for vLLM. + max_model_len: Maximum sequence length. + + Returns: + Initialized FunASRNanoVLLM engine. + """ + if os.path.isdir(model): + model_dir = model + else: + if hub in ("ms", "modelscope"): + from modelscope.hub.snapshot_download import snapshot_download + + model_dir = snapshot_download(model, revision=kwargs.pop("revision", "master")) + elif hub in ("hf", "huggingface"): + from huggingface_hub import snapshot_download + + model_dir = snapshot_download(model) + else: + raise ValueError(f"Unsupported hub: {hub}. Use 'ms' or 'hf'.") + + logger.info(f"Model directory: {model_dir}") + return cls( + model_dir=model_dir, + device=device, + dtype=dtype, + tensor_parallel_size=tensor_parallel_size, + gpu_memory_utilization=gpu_memory_utilization, + max_model_len=max_model_len, + **kwargs, + ) diff --git a/funasr/models/fun_asr_nano/inference_vllm_pipeline.py b/funasr/models/fun_asr_nano/inference_vllm_pipeline.py new file mode 100644 index 000000000..d7afc2f6d --- /dev/null +++ b/funasr/models/fun_asr_nano/inference_vllm_pipeline.py @@ -0,0 +1,364 @@ +#!/usr/bin/env python3 +# -*- encoding: utf-8 -*- +# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved. +# MIT License (https://opensource.org/licenses/MIT) + +""" +Fun-ASR-Nano vLLM Pipeline: VAD + ASR(vLLM) + Speaker Diarization. + +Replicates AutoModel's inference_with_vad pipeline but uses vLLM for +the LLM decoding step, enabling batch processing of all VAD segments +in a single generate() call. + +Usage: + from funasr.models.fun_asr_nano.inference_vllm_pipeline import FunASRNanoVLLMPipeline + + model = FunASRNanoVLLMPipeline( + model="FunAudioLLM/Fun-ASR-Nano-2512", + vad_model="fsmn-vad", + spk_model="cam++", + tensor_parallel_size=2, + ) + results = model.generate("long_meeting.wav", language="中文") +""" + +import logging +import os +import re +import time +from typing import List, Optional, Union + +import numpy as np +import torch +import torch.nn as nn + +logger = logging.getLogger(__name__) + +dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32} + + +def _clean_text(text: str) -> str: + """Remove tags, fillers, and garbage from output.""" + text = re.sub(r"<[^>]*>|]*>", "", text) + text = re.sub(r"(>.{2,8}?)\1{3,}", "", text) + text = re.sub(r"\[breath\]|\[noise\]|/sil|endofbreak|FFFF", "", text) + text = re.sub(r"\s+", " ", text) + text = text.replace("�", "").lstrip(">") + return text.strip() + + +class FunASRNanoVLLMPipeline: + """VAD + ASR(vLLM) + Speaker pipeline. + + Pipeline: + 1. VAD: segment long audio into speech regions (torch) + 2. ASR: batch ALL segments through vLLM in single generate() call + 3. Speaker: extract embeddings per segment, cluster (torch) + 4. Combine: merge text + timestamps + speaker labels + + Args: + model: Fun-ASR-Nano model name or path. + vad_model: VAD model name (e.g. "fsmn-vad"). None to disable. + vad_kwargs: VAD config (e.g. {"max_single_segment_time": 30000}). + spk_model: Speaker model name (e.g. "cam++"). None to disable. + hub: "ms" or "hf". + device: Device for audio encoder + VAD + speaker. + dtype: Compute dtype for ASR. + tensor_parallel_size: GPUs for vLLM. + gpu_memory_utilization: GPU memory fraction for vLLM. + max_model_len: Maximum sequence length for vLLM. + """ + + def __init__( + self, + model: str = "FunAudioLLM/Fun-ASR-Nano-2512", + vad_model: str = None, + vad_kwargs: dict = None, + spk_model: str = None, + spk_kwargs: dict = None, + hub: str = "ms", + device: str = "cuda:0", + dtype: str = "bf16", + tensor_parallel_size: int = 1, + gpu_memory_utilization: float = 0.8, + max_model_len: int = 4096, + enforce_eager: bool = False, + **kwargs, + ): + from funasr.models.fun_asr_nano.inference_vllm import FunASRNanoVLLM + + # ASR engine (vLLM) + self.asr_engine = FunASRNanoVLLM.from_pretrained( + model=model, hub=hub, device=device, dtype=dtype, + tensor_parallel_size=tensor_parallel_size, + gpu_memory_utilization=gpu_memory_utilization, + max_model_len=max_model_len, enforce_eager=enforce_eager, + **kwargs, + ) + + # VAD model (torch) + self.vad_model = None + if vad_model is not None: + from funasr import AutoModel + vad_kw = vad_kwargs or {} + self.vad_model = AutoModel( + model=vad_model, device=device, disable_update=True, **vad_kw + ) + + # Speaker model (torch) + self.spk_model = None + self.cb_model = None + if spk_model is not None: + from funasr import AutoModel + from funasr.models.campplus.cluster_backend import ClusterBackend + spk_kw = spk_kwargs or {} + self.spk_model = AutoModel( + model=spk_model, device=device, disable_update=True, **spk_kw + ) + cb_kwargs = spk_kw.get("cb_kwargs", {}) + self.cb_model = ClusterBackend(**cb_kwargs).to(device) + + self.device = device + self.sample_rate = 16000 + + def generate( + self, + input: Union[str, List[str]], + hotwords: List[str] = None, + language: str = None, + itn: bool = True, + max_new_tokens: int = 512, + batch_size_s: int = 300, + return_spk_res: bool = True, + **kwargs, + ) -> List[dict]: + """Run the full pipeline: VAD → ASR(vLLM) → Speaker. + + Args: + input: Audio file path(s). + hotwords: Hotwords for ASR. + language: Language hint. + itn: Inverse text normalization. + max_new_tokens: Max tokens per segment. + batch_size_s: Max batch duration in seconds (for memory control). + return_spk_res: Whether to return speaker info. + + Returns: + List of dicts: [{"key", "text", "timestamp", "sentence_info"}] + """ + if isinstance(input, str): + input = [input] + + results_all = [] + for audio_path in input: + result = self._process_one( + audio_path, hotwords=hotwords, language=language, + itn=itn, max_new_tokens=max_new_tokens, + batch_size_s=batch_size_s, return_spk_res=return_spk_res, + **kwargs, + ) + results_all.append(result) + return results_all + + def _process_one(self, audio_path, **kwargs): + """Process a single audio file through the full pipeline.""" + from funasr.utils.load_utils import load_audio_text_image_video + from funasr.utils.vad_utils import slice_padding_audio_samples + + key = os.path.splitext(os.path.basename(audio_path))[0] + + # Load audio + audio_data = load_audio_text_image_video(audio_path, fs=self.sample_rate) + if isinstance(audio_data, torch.Tensor): + audio_np = audio_data.numpy() + else: + audio_np = np.array(audio_data) + speech_length = len(audio_np) + + # Step 1: VAD + if self.vad_model is not None: + vad_res = self.vad_model.generate(input=audio_path, cache={}, is_final=True) + vad_segments = vad_res[0]["value"] # [[start_ms, end_ms], ...] + else: + vad_segments = [[0, int(speech_length / self.sample_rate * 1000)]] + + if not vad_segments: + return {"key": key, "text": "", "timestamp": []} + + n_segments = len(vad_segments) + logger.info(f"VAD: {n_segments} segments for {key}") + + # Step 2: Slice audio by VAD segments and encode + segment_audios = [] + for seg in vad_segments: + start_sample = int(seg[0] * self.sample_rate / 1000) + end_sample = int(seg[1] * self.sample_rate / 1000) + end_sample = min(end_sample, speech_length) + segment_audios.append(audio_np[start_sample:end_sample]) + + # Step 3: Batch ASR via vLLM + # Encode all segments and build prompts + from vllm import SamplingParams + from vllm.inputs.data import EmbedsPrompt + + prompts = [] + for seg_audio in segment_audios: + seg_tensor = torch.from_numpy(seg_audio).float() + adaptor_out, adaptor_out_lens, _, _ = self.asr_engine._encode_audio(seg_tensor) + input_embeds = self.asr_engine._build_input_embeds( + adaptor_out, adaptor_out_lens, + hotwords=kwargs.get("hotwords"), + language=kwargs.get("language"), + itn=kwargs.get("itn", True), + ) + prompts.append(EmbedsPrompt(prompt_embeds=input_embeds.float())) + + params = SamplingParams( + max_tokens=kwargs.get("max_new_tokens", 512), + temperature=0.0, + repetition_penalty=1.3, + skip_special_tokens=True, + ) + + # Single batch generate for ALL segments + t0 = time.perf_counter() + outputs = self.asr_engine.vllm_engine.generate(prompts, params, use_tqdm=False) + t1 = time.perf_counter() + logger.info(f"vLLM batch ASR: {n_segments} segments in {t1-t0:.3f}s") + + # Decode results + asr_results = [] + for output in outputs: + text = output.outputs[0].text + if not text and output.outputs[0].token_ids: + text = self.asr_engine.tokenizer.decode( + list(output.outputs[0].token_ids), skip_special_tokens=True + ) + text = _clean_text(text) + asr_results.append(text) + + # Step 4: Speaker embeddings (if spk_model configured) + spk_embeddings = None + if self.spk_model is not None and kwargs.get("return_spk_res", True): + from funasr.models.campplus.utils import sv_chunk, postprocess, distribute_spk + + all_segments = [] + all_spk_embs = [] + for i, seg_audio in enumerate(segment_audios): + vad_seg = [ + [vad_segments[i][0] / 1000.0, vad_segments[i][1] / 1000.0, seg_audio] + ] + chunks = sv_chunk(vad_seg) + all_segments.extend(chunks) + speech_chunks = [c[2] for c in chunks] + spk_res = self.spk_model.generate(input=speech_chunks, cache={}, is_final=True) + embs = torch.cat([r["spk_embedding"] for r in spk_res], dim=0) + all_spk_embs.append(embs) + + if all_spk_embs: + spk_embeddings = torch.cat(all_spk_embs, dim=0) + + # Step 5: Combine results + # Merge text with timestamps + full_text = "" + all_timestamps = [] + for i, (seg, text) in enumerate(zip(vad_segments, asr_results)): + if text: + if full_text: + full_text += " " + full_text += text + # Simple word-level timestamp from VAD boundaries + all_timestamps.append([int(seg[0]), int(seg[1])]) + + result = {"key": key, "text": full_text} + + # Add timestamps if available from CTC + if self.asr_engine.ctc_decoder is not None: + try: + detailed_timestamps = self._compute_all_timestamps( + segment_audios, vad_segments, asr_results + ) + if detailed_timestamps: + result["timestamp"] = detailed_timestamps + except Exception as e: + logger.debug(f"Timestamp computation failed: {e}") + result["timestamp"] = all_timestamps + else: + result["timestamp"] = all_timestamps + + # Add speaker info + if spk_embeddings is not None and self.cb_model is not None: + from funasr.models.campplus.utils import postprocess, distribute_spk + + all_segments_sorted = sorted(all_segments, key=lambda x: x[0]) + labels = self.cb_model( + spk_embeddings.cpu(), + oracle_num=kwargs.get("preset_spk_num", None), + ) + sv_output = postprocess(all_segments_sorted, None, labels, spk_embeddings.cpu()) + + # Build sentence_info + sentence_list = [] + for i, (seg, text) in enumerate(zip(vad_segments, asr_results)): + if text: + sentence_list.append({ + "start": seg[0], + "end": seg[1], + "text": text, + "timestamp": [[int(seg[0]), int(seg[1])]], + }) + distribute_spk(sentence_list, sv_output) + result["sentence_info"] = sentence_list + + return result + + def _compute_all_timestamps(self, segment_audios, vad_segments, asr_results): + """Compute CTC timestamps for all segments with VAD offsets.""" + from funasr.models.fun_asr_nano.tools.utils import forced_align + + all_timestamps = [] + for seg_audio, vad_seg, text in zip(segment_audios, vad_segments, asr_results): + if not text: + continue + try: + seg_tensor = torch.from_numpy(seg_audio).float() + from funasr.utils.load_utils import extract_fbank + speech, speech_lengths = extract_fbank( + seg_tensor, data_type="sound", + frontend=self.asr_engine.frontend, is_final=True + ) + speech = speech.to(self.device, dtype=torch.float32) + speech_lengths = speech_lengths.to(self.device) + + with torch.no_grad(): + enc_out, enc_lens = self.asr_engine.audio_encoder(speech, speech_lengths) + dec_out, dec_lens = self.asr_engine.ctc_decoder(enc_out, enc_lens) + ctc_logits = self.asr_engine.ctc.log_softmax(dec_out) + + x = ctc_logits[0, :enc_lens[0].item(), :] + target_ids = torch.tensor( + self.asr_engine.ctc_tokenizer.encode(text), dtype=torch.int64 + ) + if len(target_ids) == 0: + continue + + timestamps = forced_align(x, target_ids, self.asr_engine.blank_id) + vad_offset_ms = int(vad_seg[0]) + for ts in timestamps: + ts["token"] = self.asr_engine.ctc_tokenizer.decode([ts["token"]]) + ts["start_time"] = ts["start_time"] * 6 * 10 / 1000 + vad_offset_ms / 1000 + ts["end_time"] = ts["end_time"] * 6 * 10 / 1000 + vad_offset_ms / 1000 + all_timestamps.extend(timestamps) + except Exception as e: + logger.debug(f"Timestamp failed for segment: {e}") + all_timestamps.append({ + "start_time": vad_seg[0] / 1000, + "end_time": vad_seg[1] / 1000, + "token": text, + }) + return all_timestamps + + @classmethod + def from_pretrained(cls, model="FunAudioLLM/Fun-ASR-Nano-2512", **kwargs): + """Convenience constructor.""" + return cls(model=model, **kwargs) diff --git a/funasr/models/fun_asr_nano/inference_vllm_streaming.py b/funasr/models/fun_asr_nano/inference_vllm_streaming.py new file mode 100644 index 000000000..cdc17ebf1 --- /dev/null +++ b/funasr/models/fun_asr_nano/inference_vllm_streaming.py @@ -0,0 +1,340 @@ +#!/usr/bin/env python3 +# -*- encoding: utf-8 -*- +# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved. +# MIT License (https://opensource.org/licenses/MIT) + +""" +Fun-ASR-Nano Streaming vLLM Inference Engine. + +Design: + - Audio split into 720ms chunks (cumulative re-encoding) + - ALL chunks batched into single vLLM generate call for correctness + - Fixed/Unfixed: last 8 chars are unfixed (may change on next chunk) + - Output stabilizes as more audio accumulates (~3s+) + +Note: vLLM processes all chunks in one batch for throughput. +For real-time streaming, use the torch-based inference in demo2.py. +""" + +import logging +import os +import re +from typing import Generator, List, Union + +import numpy as np +import torch +import torch.nn as nn + +logger = logging.getLogger(__name__) +dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32} + +_CJK_RE = re.compile(r"[一-鿿]") + + +def _clean_text(text: str) -> str: + """Remove tags, repetitive garbage, filler tokens, and invalid chars.""" + text = re.sub(r'<[^>]*>|]*>', '', text) + text = re.sub(r'(>.{2,8}?){3,}', '', text) + text = re.sub(r'\[breath\]|\[noise\]|/sil|endofbreak|FFFF', '', text) + text = re.sub(r'\s+', ' ', text) + text = text.replace('�', '').lstrip('>') + return text.strip() + + +def _is_meaningful(text: str) -> bool: + """Check if text has real ASR content.""" + return len(_CJK_RE.findall(text)) >= 2 + + +class FunASRNanoStreamingVLLM: + """Streaming ASR with vLLM backend (batch-all-chunks approach). + + Processes audio in 720ms chunks. All chunks are encoded and batched + into a single vLLM generate() call for correct and efficient inference. + Results are returned per-chunk with fixed/unfixed regions. + + Args: + model_dir: Path to Fun-ASR-Nano model directory. + device: Device for audio encoder/adaptor. + dtype: Compute dtype ("bf16", "fp16", "fp32"). + tensor_parallel_size: GPUs for vLLM tensor parallelism. + gpu_memory_utilization: GPU memory fraction for KV cache. + max_model_len: Maximum sequence length. + chunk_ms: Chunk duration in ms (default 720). + rollback_chars: Characters to rollback per chunk (default 8). + """ + + def __init__(self, model_dir, device="cuda:0", dtype="bf16", + tensor_parallel_size=1, gpu_memory_utilization=0.8, + max_model_len=2048, enforce_eager=False, + chunk_ms=720, rollback_chars=8, **kwargs): + from vllm import LLM + from funasr.models.fun_asr_nano.inference_vllm import prepare_vllm_model_dir + + self.device = device + self.dtype = dtype + self.torch_dtype = dtype_map.get(dtype, torch.bfloat16) + self.model_dir = model_dir + self.chunk_ms = chunk_ms + self.rollback_chars = rollback_chars + + vllm_model_dir = prepare_vllm_model_dir(model_dir) + self._load_audio_components(model_dir) + + vllm_kwargs = kwargs.get("vllm_kwargs", {}) + self.vllm_engine = LLM( + enable_prompt_embeds=True, model=vllm_model_dir, + tensor_parallel_size=tensor_parallel_size, + gpu_memory_utilization=gpu_memory_utilization, + max_model_len=max_model_len, enforce_eager=enforce_eager, + dtype={"bf16": "bfloat16", "fp16": "float16", "fp32": "auto"}.get(dtype, dtype), + trust_remote_code=True, **vllm_kwargs, + ) + self.tokenizer = self.vllm_engine.get_tokenizer() + self._load_embedding_layer(model_dir) + self.sample_rate = self.frontend.fs + self.chunk_samples = int(self.sample_rate * self.chunk_ms / 1000) + + def _load_audio_components(self, model_dir): + from omegaconf import OmegaConf + from funasr.register import tables + config = OmegaConf.load(os.path.join(model_dir, "config.yaml")) + self._config = OmegaConf.to_container(config, resolve=True) + + frontend_class = tables.frontend_classes.get(config["frontend"]) + frontend_conf = OmegaConf.to_container(config.get("frontend_conf", {}), resolve=True) + self.frontend = frontend_class(**frontend_conf) + self.frontend.eval() + + encoder_conf = OmegaConf.to_container(config.get("audio_encoder_conf", {}), resolve=True) + if encoder_conf.get("hub") == "ms": + from funasr import AutoModel as FAM + enc_m = FAM(model=config["audio_encoder"], model_revision="master", disable_update=True) + self.audio_encoder_output_size = getattr(enc_m.model, "encoder_output_size", -1) + self.audio_encoder = enc_m.model.model.encoder if hasattr(enc_m.model, "model") else enc_m.model.encoder + else: + encoder_class = tables.encoder_classes.get(config["audio_encoder"]) + self.audio_encoder = encoder_class(input_size=self.frontend.output_size(), **encoder_conf) + self.audio_encoder_output_size = self.audio_encoder.output_size() + self.audio_encoder.eval() + for p in self.audio_encoder.parameters(): p.requires_grad = False + + adaptor_conf = OmegaConf.to_container(config.get("audio_adaptor_conf", {}), resolve=True) + adaptor_class = tables.adaptor_classes.get(config["audio_adaptor"]) + if self.audio_encoder_output_size > 0: + adaptor_conf["encoder_dim"] = self.audio_encoder_output_size + self.audio_adaptor = adaptor_class(**adaptor_conf) + self.audio_adaptor.eval() + for p in self.audio_adaptor.parameters(): p.requires_grad = False + + model_pt = os.path.join(model_dir, "model.pt") + if os.path.exists(model_pt): + ckpt = torch.load(model_pt, map_location="cpu") + sd = ckpt.get("state_dict", ckpt) + enc_s = {k[len("audio_encoder."):]: v for k, v in sd.items() if k.startswith("audio_encoder.")} + if enc_s: self.audio_encoder.load_state_dict(enc_s, strict=False) + adp_s = {k[len("audio_adaptor."):]: v for k, v in sd.items() if k.startswith("audio_adaptor.")} + if adp_s: self.audio_adaptor.load_state_dict(adp_s, strict=False) + + self.audio_encoder = self.audio_encoder.to(self.device, dtype=torch.float32) + self.audio_adaptor = self.audio_adaptor.to(self.device, dtype=self.torch_dtype) + + def _load_embedding_layer(self, model_dir): + ckpt = torch.load(os.path.join(model_dir, "model.pt"), map_location="cpu") + sd = ckpt.get("state_dict", ckpt) + for key in sd: + if "embed_tokens.weight" in key and key.startswith("llm."): + self.embed_tokens = nn.Embedding.from_pretrained(sd[key], freeze=True) + self.embed_tokens = self.embed_tokens.to(self.device, dtype=self.torch_dtype) + return + raise RuntimeError("Could not find LLM embedding weights") + + @torch.no_grad() + def _encode_audio(self, audio_samples): + from funasr.utils.load_utils import extract_fbank + speech, speech_lengths = extract_fbank( + audio_samples, data_type="sound", frontend=self.frontend, is_final=True) + speech = speech.to(self.device, dtype=torch.float32) + speech_lengths = speech_lengths.to(self.device) + enc_out, enc_lens = self.audio_encoder(speech, speech_lengths) + adp_out, adp_lens = self.audio_adaptor(enc_out.to(dtype=self.torch_dtype), enc_lens) + return adp_out, adp_lens + + def _build_prompt_text(self, hotwords=None, language=None, itn=True): + hotwords = hotwords or [] + if hotwords: + prompt = "请结合上下文信息，更加准确地完成语音转写任务。如果没有相关信息，我们会留空。\n\n\n**上下文信息：**\n\n\n" + prompt += f"热词列表：[{', '.join(hotwords)}]\n" + else: + prompt = "" + prompt += f"语音转写成{language}" if language else "语音转写" + if not itn: prompt += "，不进行文本规整" + return prompt + "：" + + @torch.no_grad() + def _build_embeds(self, audio_embeds, audio_embed_lens, prev_text="", hotwords=None, language=None, itn=True): + """Build input embeddings. prev_text is appended as assistant prefix for continuation.""" + prompt = self._build_prompt_text(hotwords, language, itn) + prefix_text = f"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n{prompt}<|startofspeech|>" + suffix_text = "<|endofspeech|><|im_end|>\n<|im_start|>assistant\n\n\n\n\n" + if prev_text: + suffix_text += prev_text + + prefix_ids = self.tokenizer.encode(prefix_text, add_special_tokens=False) + suffix_ids = self.tokenizer.encode(suffix_text, add_special_tokens=False) + prefix_emb = self.embed_tokens(torch.tensor(prefix_ids, dtype=torch.long, device=self.device)) + suffix_emb = self.embed_tokens(torch.tensor(suffix_ids, dtype=torch.long, device=self.device)) + audio_emb = audio_embeds[0, :audio_embed_lens[0].item(), :] + return torch.cat([prefix_emb, audio_emb, suffix_emb], dim=0) + + def streaming_generate(self, audio_input, chunk_ms=None, rollback_chars=None, + hotwords=None, language=None, itn=True, + max_new_tokens=200, temperature=0.0, **kwargs): + """Streaming ASR: process all chunks and yield results per chunk. + + All chunks are batched into a single vLLM generate() call for + correct results. Yields incrementally improving transcriptions. + + Args: + audio_input: File path, numpy array, or tensor (16kHz). + chunk_ms: Chunk size in ms (default 720). + rollback_chars: Chars to rollback (default 8). + hotwords: Hotword list. + language: Language hint (e.g. "中文"). + itn: Inverse text normalization. + max_new_tokens: Max tokens per chunk generation. + temperature: Sampling temperature (0 = greedy). + + Yields: + {"text": full_text, "fixed_text": confirmed_text, + "is_final": bool, "chunk_idx": int, "audio_duration_ms": float} + """ + from vllm import SamplingParams + from vllm.inputs.data import EmbedsPrompt + from funasr.utils.load_utils import load_audio_text_image_video + + chunk_ms = chunk_ms or self.chunk_ms + rollback_chars = rollback_chars or self.rollback_chars + + if isinstance(audio_input, str): + audio_data = load_audio_text_image_video(audio_input, fs=self.sample_rate) + elif isinstance(audio_input, np.ndarray): + audio_data = torch.from_numpy(audio_input).float() + elif isinstance(audio_input, torch.Tensor): + audio_data = audio_input.float() + else: + raise ValueError(f"Unsupported audio type: {type(audio_input)}") + if audio_data.dim() > 1: + audio_data = audio_data.squeeze() + + total_samples = audio_data.shape[0] + chunk_samples = int(self.sample_rate * chunk_ms / 1000) + num_chunks = (total_samples + chunk_samples - 1) // chunk_samples + + params = SamplingParams(max_tokens=max_new_tokens, temperature=temperature, + repetition_penalty=1.3, skip_special_tokens=True) + + # Two-stage approach for long audio: + # Stage 1: batch first N chunks fresh (no prev_text) to find stable output + # Stage 2: batch remaining chunks WITH prev_text from stable output + stage1_count = min(10, num_chunks) # ~7.2s should be enough to stabilize + + # Stage 1: encode and batch first chunks + prompts_s1 = [] + chunk_infos_s1 = [] + for i in range(stage1_count): + end_sample = min((i + 1) * chunk_samples, total_samples) + adaptor_out, adaptor_out_lens = self._encode_audio(audio_data[:end_sample]) + embeds = self._build_embeds(adaptor_out, adaptor_out_lens, prev_text="", + hotwords=hotwords, language=language, itn=itn) + prompts_s1.append(EmbedsPrompt(prompt_embeds=embeds.float())) + chunk_infos_s1.append({ + "chunk_idx": i + 1, + "is_final": end_sample >= total_samples, + "audio_duration_ms": end_sample * 1000 / self.sample_rate, + }) + + outputs_s1 = self.vllm_engine.generate(prompts_s1, params, use_tqdm=False) + + # Find best stable output from stage 1 + best_text = "" + results_s1 = [] + for output in outputs_s1: + text = output.outputs[0].text + if not text and output.outputs[0].token_ids: + text = self.tokenizer.decode(list(output.outputs[0].token_ids), skip_special_tokens=True) + text = _clean_text(text) + results_s1.append(text) + if _is_meaningful(text) and len(text) > len(best_text): + best_text = text + + # Yield stage 1 results + for i, (text, info) in enumerate(zip(results_s1, chunk_infos_s1)): + if info["is_final"]: + fixed_text = text + elif _is_meaningful(text) and len(text) > rollback_chars: + fixed_text = text[:-rollback_chars] + else: + fixed_text = "" + yield {"text": text, "fixed_text": fixed_text, **info} + + # Stage 2: if more chunks remain, use prev_text from stable output + if stage1_count < num_chunks: + prev_text = best_text[:-rollback_chars] if len(best_text) > rollback_chars else best_text + + prompts_s2 = [] + chunk_infos_s2 = [] + for i in range(stage1_count, num_chunks): + end_sample = min((i + 1) * chunk_samples, total_samples) + adaptor_out, adaptor_out_lens = self._encode_audio(audio_data[:end_sample]) + embeds = self._build_embeds(adaptor_out, adaptor_out_lens, prev_text=prev_text, + hotwords=hotwords, language=language, itn=itn) + prompts_s2.append(EmbedsPrompt(prompt_embeds=embeds.float())) + chunk_infos_s2.append({ + "chunk_idx": i + 1, + "is_final": end_sample >= total_samples, + "audio_duration_ms": end_sample * 1000 / self.sample_rate, + }) + + outputs_s2 = self.vllm_engine.generate(prompts_s2, params, use_tqdm=False) + + for output, info in zip(outputs_s2, chunk_infos_s2): + text = output.outputs[0].text + if not text and output.outputs[0].token_ids: + text = self.tokenizer.decode(list(output.outputs[0].token_ids), skip_special_tokens=True) + text = _clean_text(text) + full_text = prev_text + text + + if info["is_final"]: + fixed_text = full_text + elif _is_meaningful(full_text) and len(full_text) > rollback_chars: + fixed_text = full_text[:-rollback_chars] + else: + fixed_text = prev_text + + yield {"text": full_text, "fixed_text": fixed_text, **info} + + def generate(self, audio_input, **kwargs): + """Run streaming and return all chunk results.""" + return list(self.streaming_generate(audio_input, **kwargs)) + + @classmethod + def from_pretrained(cls, model="FunAudioLLM/Fun-ASR-Nano-2512", hub="ms", + device="cuda:0", dtype="bf16", tensor_parallel_size=1, + gpu_memory_utilization=0.8, max_model_len=2048, + chunk_ms=720, rollback_chars=8, **kwargs): + """Load from hub or local path.""" + if os.path.isdir(model): + model_dir = model + else: + if hub in ("ms", "modelscope"): + from modelscope.hub.snapshot_download import snapshot_download + model_dir = snapshot_download(model, revision=kwargs.pop("revision", "master")) + else: + from huggingface_hub import snapshot_download + model_dir = snapshot_download(model) + return cls(model_dir=model_dir, device=device, dtype=dtype, + tensor_parallel_size=tensor_parallel_size, + gpu_memory_utilization=gpu_memory_utilization, + max_model_len=max_model_len, chunk_ms=chunk_ms, + rollback_chars=rollback_chars, **kwargs)