本地部署 whisper-medusa

0. 引言

Whisper 是一种用于语音转录和翻译的高级编码器-解码器模型，通过编码和解码阶段处理音频。鉴于其尺寸大和推理速度慢，人们提出了各种优化策略（例如 Faster-Whisper 和 Speculative Decoding）来提高性能。我们的 Medusa 模型建立在 Whisper 的基础上，通过每次迭代预测多个标记，这显着提高了速度，同时 WER 略有下降。我们在 LibriSpeech 数据集上训练和评估我们的模型，与普通 Whisper 模型相比，展示了强大的性能速度改进和同等准确度。

1. 本地部署

1-1. 创建虚拟环境

conda create -n whisper-medusa python=3.11 -y
conda activate whisper-medusa

1-2. 克隆代码

git clone https://github.com/aiola-lab/whisper-medusa.git
cd whisper-medusa

1-3. 安装依赖模块

pip install torch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 --index-url https://download.pytorch.org/whl/cu118
pip install -e .
conda install matplotlib
pip install gradio

1-4. 创建 Web UI

# webui.py
import torch
import torchaudio
import gradio as gr
from whisper_medusa import WhisperMedusaModel
from transformers import WhisperProcessor# Load model and processor
model_name = "aiola/whisper-medusa-v1"
model = WhisperMedusaModel.from_pretrained(model_name)
processor = WhisperProcessor.from_pretrained(model_name)# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)# Constants
SAMPLING_RATE = 16000def transcribe_audio(audio_file, language):# Load and preprocess audioinput_speech, sr = torchaudio.load(audio_file)if input_speech.shape[0] > 1:  # If stereo, average the channelsinput_speech = input_speech.mean(dim=0, keepdim=True)if sr != SAMPLING_RATE:input_speech = torchaudio.transforms.Resample(sr, SAMPLING_RATE)(input_speech)# Process inputinput_features = processor(input_speech.squeeze(), return_tensors="pt", sampling_rate=SAMPLING_RATE).input_featuresinput_features = input_features.to(device)# Generate transcriptionmodel_output = model.generate(input_features,language=language,)predict_ids = model_output[0]transcription = processor.decode(predict_ids, skip_special_tokens=True)return transcription# Define Gradio interface
iface = gr.Interface(fn=transcribe_audio,inputs=[gr.Audio(type="filepath", label="Upload Audio"),gr.Dropdown(["en", "zh", "ja"], label="Select Language", value="en")],outputs="text",title="Audio Transcription with Whisper Medusa",description="Upload an audio file and select the language to transcribe the audio to text."
)# Launch the interface
iface.launch()

1-5. 启动 Web UI

python webui.py

1-5. 访问 Web UI

使用浏览器访问 http://localhost:7860