mirror of
https://github.com/modelscope/FunASR.git
synced 2026-10-12 02:19:18 +00:00
Merge branch 'vllm' - vLLM inference engine + streaming WebSocket service
Features: - AutoModelVLLM: generic vLLM inference for all LLM-based ASR models - serve_realtime_ws.py: streaming WebSocket ASR with VAD+SPK+hotwords - DynamicStreamingVAD: dynamic silence threshold (default on for all models) - Fun-ASR-Nano batch_size fix for VAD pipeline - Comprehensive vLLM documentation (docs/vllm_guide.md)
This commit is contained in:
游雁
@@ -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 <server>`
|
||||
|
||||
### 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 <server>`
|
||||
|
||||
### 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` | 当前使用的静音阈值 |
|
||||
@@ -0,0 +1,160 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="zh-CN">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<title>Fun-ASR-Nano · Streaming ASR with Speaker Diarization</title>
|
||||
<link href="https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800;900&display=swap" rel="stylesheet">
|
||||
<style>
|
||||
* { margin:0; padding:0; box-sizing:border-box; }
|
||||
html, body { height:100%; }
|
||||
body { font-family:'Inter',system-ui,sans-serif; background:#080810; color:#e0e0e0; display:flex; justify-content:center; height:100vh; overflow:hidden; }
|
||||
|
||||
/* === Hero Section === */
|
||||
.hero { padding:24px 28px 16px; flex-shrink:0; }
|
||||
.hero-top { display:flex; align-items:center; gap:16px; margin-bottom:14px; }
|
||||
.logo { font-size:30px; font-weight:900; letter-spacing:-1px; }
|
||||
.logo .g { background:linear-gradient(135deg,#64ffda,#00bfa5); -webkit-background-clip:text; -webkit-text-fill-color:transparent; }
|
||||
.logo .w { color:#fff; }
|
||||
.logo .s { color:rgba(255,255,255,0.3); font-weight:400; font-size:15px; margin-left:8px; }
|
||||
.hero-links { margin-left:auto; display:flex; gap:10px; }
|
||||
.hero-links a { font-size:12px; color:#aaa; text-decoration:none; padding:6px 14px; border-radius:8px; background:rgba(255,255,255,0.04); border:1px solid rgba(255,255,255,0.08); transition:all 0.2s; font-weight:500; }
|
||||
.hero-links a:hover { color:#64ffda; border-color:rgba(100,255,218,0.3); }
|
||||
|
||||
.features { display:flex; gap:10px; flex-wrap:wrap; margin-bottom:14px; }
|
||||
.feat { font-size:11px; padding:5px 12px; border-radius:20px; font-weight:500; }
|
||||
.feat-green { color:#64ffda; background:rgba(100,255,218,0.08); border:1px solid rgba(100,255,218,0.2); }
|
||||
.feat-yellow { color:#fbbf24; background:rgba(251,191,36,0.08); border:1px solid rgba(251,191,36,0.2); }
|
||||
.feat-purple { color:#a78bfa; background:rgba(167,139,250,0.08); border:1px solid rgba(167,139,250,0.2); }
|
||||
.feat-blue { color:#38bdf8; background:rgba(56,189,248,0.08); border:1px solid rgba(56,189,248,0.2); }
|
||||
|
||||
.hero-desc { font-size:13px; color:rgba(255,255,255,0.45); line-height:1.8; }
|
||||
.hero-desc b { color:rgba(255,255,255,0.75); font-weight:600; }
|
||||
.hero-desc a { color:#64ffda; text-decoration:none; font-weight:600; }
|
||||
.hero-desc a:hover { text-decoration:underline; }
|
||||
|
||||
/* === Main Panel === */
|
||||
.main { flex:1; display:flex; flex-direction:column; padding:0 28px 20px; min-height:0; }
|
||||
.panel { flex:1; display:flex; flex-direction:column; background:rgba(16,16,24,0.95); border:1px solid rgba(255,255,255,0.05); border-radius:16px; box-shadow:0 10px 40px rgba(0,0,0,0.4); min-height:0; }
|
||||
|
||||
.controls { padding:14px 20px; display:flex; align-items:center; gap:12px; border-bottom:1px solid rgba(255,255,255,0.04); flex-shrink:0; }
|
||||
.controls input[type=text] { background:rgba(255,255,255,0.05); border:1px solid rgba(255,255,255,0.1); border-radius:8px; padding:8px 12px; color:#ccc; font-size:12px; width:180px; outline:none; }
|
||||
.controls input[type=text]:focus { border-color:rgba(100,255,218,0.4); }
|
||||
.controls input[type=file] { display:none; }
|
||||
.controls label { font-size:11px; color:#888; display:flex; align-items:center; gap:5px; cursor:pointer; }
|
||||
.controls input[type=checkbox] { accent-color:#64ffda; }
|
||||
.btn { padding:8px 18px; border:none; border-radius:10px; font-size:12px; cursor:pointer; font-weight:600; transition:all 0.12s; }
|
||||
.btn:hover { transform:scale(1.04); }
|
||||
.btn:active { transform:scale(0.97); }
|
||||
.btn-mic { background:linear-gradient(135deg,#64ffda,#00bfa5); color:#080810; }
|
||||
.btn-file { background:linear-gradient(135deg,#fbbf24,#f59e0b); color:#1a1a2e; }
|
||||
.btn-hw { background:rgba(167,139,250,0.12); color:#a78bfa; border:1px solid rgba(167,139,250,0.3); }
|
||||
.btn-stop { background:linear-gradient(135deg,#ef4444,#dc2626); color:#fff; }
|
||||
.sta { margin-left:auto; display:flex; align-items:center; gap:6px; }
|
||||
.dot { width:7px; height:7px; border-radius:50%; background:#222; }
|
||||
.dot.on { background:#64ffda; box-shadow:0 0 10px rgba(100,255,218,0.6); animation:g 1.5s infinite; }
|
||||
@keyframes g { 0%,100%{opacity:1} 50%{opacity:0.3} }
|
||||
.sta span { font-size:10px; color:#555; }
|
||||
.hw-info { font-size:10px; color:#a78bfa; margin-left:2px; }
|
||||
|
||||
.result { flex:1; overflow-y:auto; padding:18px 22px; min-height:0; }
|
||||
.result::-webkit-scrollbar { width:4px; }
|
||||
.result::-webkit-scrollbar-thumb { background:rgba(255,255,255,0.06); border-radius:2px; }
|
||||
.line { padding:6px 0; border-bottom:1px solid rgba(255,255,255,0.02); display:flex; align-items:baseline; gap:0; word-break:break-all; }
|
||||
.time { color:rgba(100,255,218,0.5); font-size:11px; font-family:'SF Mono','Menlo',monospace; width:100px; min-width:100px; flex-shrink:0; padding-top:2px; }
|
||||
.spk { font-size:10px; font-weight:600; padding:2px 8px; border-radius:5px; flex-shrink:0; margin-right:10px; }
|
||||
.text { color:rgba(255,255,255,0.88); font-size:15px; line-height:1.7; flex:1; }
|
||||
.partial .text { color:rgba(255,255,255,0.25); font-style:italic; }
|
||||
.ph { color:rgba(255,255,255,0.1); text-align:center; padding:50px 20px; font-size:14px; line-height:2; }
|
||||
</style>
|
||||
</head>
|
||||
<body><div style="width:920px;height:100%;display:flex;flex-direction:column">
|
||||
|
||||
<div class="hero">
|
||||
<div class="hero-top">
|
||||
<div class="logo"><span class="g">Fun-ASR-Nano</span><span class="s">vLLM Engine</span></div>
|
||||
<div class="hero-links">
|
||||
<a href="https://github.com/modelscope/FunASR" target="_blank">GitHub</a>
|
||||
<a href="https://modelscope.cn/models/FunAudioLLM/Fun-ASR-Nano-2512" target="_blank">ModelScope</a>
|
||||
<a href="https://huggingface.co/FunAudioLLM/Fun-ASR-Nano-2512" target="_blank">HuggingFace</a>
|
||||
</div>
|
||||
</div>
|
||||
<div class="features">
|
||||
<span class="feat feat-green">Streaming ASR</span>
|
||||
<span class="feat feat-yellow">Speaker Diarization <span style="font-size:9px;opacity:0.7">(Beta)</span></span>
|
||||
<span class="feat feat-purple">Hotword Customization</span>
|
||||
<span class="feat feat-blue">31 Languages · 7 Dialects</span>
|
||||
</div>
|
||||
<p class="hero-desc">
|
||||
基于 <b>FunASR</b> 的 <b>vLLM 推理引擎</b>，实现流式语音识别服务。支持实时 VAD 分句、说话人分离 <span style="color:#f59e0b;font-size:11px">(Beta)</span>、<b>热词定制化</b>（加载人名、地名等实体词列表，提升专有名词识别准确率）、31种语言及中文方言。所有代码与模型已全部开源。<br>
|
||||
<span style="font-size:12px">Streaming ASR with vLLM engine, real-time VAD, speaker diarization <span style="color:#f59e0b">(Beta)</span>, <b>hotword customization</b> (names, places, entities), 31 languages & Chinese dialects. Fully open-sourced.</span>
|
||||
· <a href="https://www.funasr.com" target="_blank">www.funasr.com</a>
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<div class="main">
|
||||
<div class="panel">
|
||||
<div class="controls">
|
||||
<input type="text" id="serverUrl" value="ws://localhost:10095">
|
||||
<label><input type="checkbox" id="showSpk" checked> Speaker <span style="font-size:9px;color:#f59e0b;background:rgba(245,158,11,0.08);padding:1px 5px;border-radius:4px">Beta</span></label>
|
||||
<button class="btn btn-mic" id="btnMic" onclick="startMic()">Mic</button>
|
||||
<input type="file" id="fileInput" accept="audio/*,.wav,.mp3,.flac,.mp4,.m4a">
|
||||
<button class="btn btn-file" id="btnFile" onclick="document.getElementById('fileInput').click()">Audio File</button>
|
||||
<input type="file" id="hotwordFile" accept=".txt" style="display:none">
|
||||
<button class="btn btn-hw" id="btnHw" onclick="document.getElementById('hotwordFile').click()">Hotwords</button>
|
||||
<span class="hw-info" id="hwInfo"></span>
|
||||
<select id="langSelect" style="background:rgba(255,255,255,0.05);border:1px solid rgba(255,255,255,0.1);border-radius:8px;padding:6px 10px;color:#ccc;font-size:11px;outline:none">
|
||||
<option value="">Auto</option>
|
||||
<option value="中文">中文</option>
|
||||
<option value="English">English</option>
|
||||
<option value="日本語">日本語</option>
|
||||
<option value="한국어">한국어</option>
|
||||
<option value="Deutsch">Deutsch</option>
|
||||
<option value="Français">Français</option>
|
||||
<option value="Español">Español</option>
|
||||
<option value="Русский">Русский</option>
|
||||
<option value="العربية">العربية</option>
|
||||
<option value="Português">Português</option>
|
||||
<option value="Italiano">Italiano</option>
|
||||
</select>
|
||||
<button class="btn btn-stop" id="btnStop" onclick="stopAll()" style="display:none">Stop</button>
|
||||
<div class="sta"><div class="dot" id="dot"></div><span id="status">Ready</span></div>
|
||||
</div>
|
||||
<div class="result" id="resultBox">
|
||||
<div class="ph">Click <b>Mic</b> for real-time recognition or <b>Audio File</b> to transcribe a file<br>Load a <b>Hotwords</b> file (.txt, one word per line) to boost recognition of names, places & entities</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<input type="text" id="hotwords" style="display:none">
|
||||
<script>
|
||||
let ws=null,mediaStream=null,audioContext=null,processor=null,isRecording=false;
|
||||
const C=['#64ffda','#f472b6','#fbbf24','#34d399','#a78bfa','#fb923c','#67e8f9','#f87171','#38bdf8','#c084fc'];
|
||||
|
||||
document.getElementById('hotwordFile').addEventListener('change',function(e){
|
||||
const f=e.target.files[0];if(!f)return;
|
||||
const r=new FileReader();r.onload=ev=>{
|
||||
const words=ev.target.result.split('\n').map(l=>l.trim()).filter(l=>l);
|
||||
document.getElementById('hotwords').value=words.join(',');
|
||||
document.getElementById('hwInfo').textContent=words.length+' words loaded';
|
||||
setS('Hotwords: '+words.length+' loaded');
|
||||
};r.readAsText(f);
|
||||
});
|
||||
document.getElementById('fileInput').addEventListener('change',function(){if(this.files.length)startFile();});
|
||||
|
||||
function setS(m,on){document.getElementById('status').textContent=m;document.getElementById('dot').className=on?'dot on':'dot';}
|
||||
function sStop(){document.getElementById('btnMic').style.display='none';document.getElementById('btnFile').style.display='none';document.getElementById('btnHw').style.display='none';document.getElementById('btnStop').style.display='inline';}
|
||||
function sStart(){document.getElementById('btnMic').style.display='inline';document.getElementById('btnFile').style.display='inline';document.getElementById('btnHw').style.display='inline';document.getElementById('btnStop').style.display='none';}
|
||||
function render(ss,p,ps,d,f){
|
||||
const b=document.getElementById('resultBox'),sk=document.getElementById('showSpk').checked;let h='';
|
||||
ss.forEach(s=>{let st=s.start!==undefined?s.start:(s.start_ms||0),en=s.end!==undefined?s.end:(s.end_ms||0),sp=s.spk!==undefined?s.spk:-1,sh='';
|
||||
if(sk&&sp>=0){let c=C[sp%C.length];sh='<span class="spk" style="color:'+c+';background:'+c+'12;border:1px solid '+c+'30">SPK'+sp+'</span>';}
|
||||
h+='<div class="line"><span class="time">'+(st/1000).toFixed(1)+' - '+(en/1000).toFixed(1)+'s</span>'+sh+'<span class="text">'+s.text+'</span></div>';});
|
||||
if(p)h+='<div class="line partial"><span class="time">'+(ps/1000).toFixed(1)+'s ...</span><span class="text">'+p+'</span></div>';
|
||||
if(!h)h='<div class="ph">Listening...</div>';b.innerHTML=h;b.scrollTop=b.scrollHeight;}
|
||||
function con(cb){ws=new WebSocket(document.getElementById('serverUrl').value);ws.onopen=()=>{ws.send('START');var hw=document.getElementById('hotwords').value.trim();if(hw)ws.send('HOTWORDS:'+hw);var lang=document.getElementById('langSelect').value;if(lang)ws.send('LANGUAGE:'+lang);cb();};ws.onmessage=e=>{const d=JSON.parse(e.data);if(d.sentences!==undefined)render(d.sentences,d.partial,d.partial_start_ms,d.duration_ms,d.is_final);};ws.onerror=()=>setS('Error');ws.onclose=()=>{if(isRecording)stopAll();};}
|
||||
function startMic(){con(async()=>{setS('Recording',true);sStop();document.getElementById('resultBox').innerHTML='<div class="ph">Listening...</div>';try{mediaStream=await navigator.mediaDevices.getUserMedia({audio:{sampleRate:16000,channelCount:1,echoCancellation:true}});}catch(e){setS('Mic denied');ws.close();sStart();return;}audioContext=new AudioContext({sampleRate:16000});const s=audioContext.createMediaStreamSource(mediaStream);processor=audioContext.createScriptProcessor(4096,1,1);processor.onaudioprocess=e=>{if(!isRecording)return;const f=e.inputBuffer.getChannelData(0),i=new Int16Array(f.length);for(let j=0;j<f.length;j++)i[j]=Math.max(-32768,Math.min(32767,Math.round(f[j]*32768)));if(ws&&ws.readyState===1)ws.send(i.buffer);};s.connect(processor);processor.connect(audioContext.destination);isRecording=true;});}
|
||||
function startFile(){const fi=document.getElementById('fileInput');if(!fi.files.length){setS('Pick a file');return;}const file=fi.files[0];setS('Decoding...');document.getElementById('resultBox').innerHTML='<div class="ph">Decoding '+file.name+'...</div>';const r=new FileReader();r.onload=async e=>{const a=new AudioContext({sampleRate:16000});let buf;try{buf=await a.decodeAudioData(e.target.result);}catch(err){setS('Decode failed');return;}const p=buf.getChannelData(0),i=new Int16Array(p.length);for(let j=0;j<p.length;j++)i[j]=Math.max(-32768,Math.min(32767,Math.round(p[j]*32768)));a.close();con(async()=>{setS('Streaming '+(p.length/16000|0)+'s',true);sStop();isRecording=true;for(let j=0;j<i.length&&isRecording;j+=4096){if(ws&&ws.readyState===1)ws.send(i.slice(j,j+4096).buffer);await new Promise(r=>setTimeout(r,50));}if(isRecording)stopAll();});};r.readAsArrayBuffer(file);}
|
||||
function stopAll(){isRecording=false;if(ws&&ws.readyState===1)ws.send('STOP');if(processor){processor.disconnect();processor=null;}if(audioContext){audioContext.close();audioContext=null;}if(mediaStream){mediaStream.getTracks().forEach(t=>t.stop());mediaStream=null;}setTimeout(()=>{if(ws){ws.close();ws=null;}},2000);sStart();setS('Done');}
|
||||
</script>
|
||||
</div></body>
|
||||
</html>
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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 <audio_file_or_scp>")
|
||||
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()
|
||||
@@ -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 <server>
|
||||
# 然后本地打开 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）
|
||||
@@ -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])
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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'(?<!\S)(?!\d+$)(\w+)(?:\s+\1){{{max_occurrences - 1},}}(?!\S)'
|
||||
if regex.search(word_pattern, cleaned, regex.IGNORECASE):
|
||||
match = regex.search(word_pattern, cleaned, regex.IGNORECASE)
|
||||
repeated = match.group(1)
|
||||
pos = text.find(repeated)
|
||||
if pos >= 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'(?<!\d)(\S{{{length}}})\1{{{max_occurrences - 1},}}(?!\d)'
|
||||
combined = rf'(?=.*\D){pattern}'
|
||||
match = regex.search(combined, cleaned)
|
||||
if match:
|
||||
repeated = match.group(1)
|
||||
pos = text.find(repeated)
|
||||
if pos >= 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))
|
||||
@@ -0,0 +1,4 @@
|
||||
兔子
|
||||
窝边草
|
||||
土匪
|
||||
师爷
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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:
|
||||
|
||||
@@ -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<think>\n\n</think>\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,
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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<think>\n\n</think>\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)
|
||||
Reference in New Issue
Block a user