Whisper Medusaを試してみた：高速で正確な音声認識モデル

音声認識技術の進歩は目覚ましく、最近では精度と速度の両面で大きな改善が見られています。今回は、OpenAIのWhisperモデルをベースに、さらなる高速化を実現した「Whisper Medusa」を実際に試してみました。

「Whisper-Medusa」使ってみた！
普通にCPUでも起動できた！
まだ速度を体感できてないのでGPU環境でも触ってみます！ https://t.co/o3pIK9H3Fw pic.twitter.com/R1HN87hm2g

— Maki@Sunwood AI Labs. (@hAru_mAki_ch) August 5, 2024

Whisper Medusaとは
インストールと準備
使用方法
実行結果と感想
まとめ
リポジトリ
1. 関連

Whisper Medusaとは

Whisper Medusaは、音声の転写と翻訳のための高度なエンコーダー・デコーダーモデルであるWhisperを改良したものです。通常のWhisperモデルに比べて、1回の推論で複数のトークンを予測することで、大幅な速度向上を実現しています。

開発元によると、平均で通常のWhisperの1.5倍の速度で生成を行うことができ、しかもWord Error Rate（WER）はほぼ同等（Whisper Medusa：4.2% vs Whisper：4%）とのことです。

インストールと準備

Whisper Medusaを使用するには、まず環境を整える必要があります。以下の手順でインストールを行いました：

Conda環境の作成と有効化： conda create -n whisper-medusa python=3.11 -y conda activate whisper-medusa
PyTorchのインストール： pip install torch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 --index-url https://download.pytorch.org/whl/cu118
Whisper Medusaのインストール： git clone https://github.com/aiola-lab/whisper-medusa.git cd whisper-medusa pip install -e .

使用方法

Whisper Medusaを使用するためのPythonスクリプトを作成しました。以下がその内容です：

import torch
import torchaudio

from whisper_medusa import WhisperMedusaModel
from transformers import WhisperProcessor

model_name = "aiola/whisper-medusa-v1"
model = WhisperMedusaModel.from_pretrained(model_name)
processor = WhisperProcessor.from_pretrained(model_name)

path_to_audio = "path/to/audio.wav"
SAMPLING_RATE = 16000
language = "en"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

input_speech, sr = torchaudio.load(path_to_audio)
if input_speech.shape[0] > 1:  # ステレオの場合、チャンネルを平均化
    input_speech = input_speech.mean(dim=0, keepdim=True)

if sr != SAMPLING_RATE:
    input_speech = torchaudio.transforms.Resample(sr, SAMPLING_RATE)(input_speech)

input_features = processor(input_speech.squeeze(), return_tensors="pt", sampling_rate=SAMPLING_RATE).input_features
input_features = input_features.to(device)

model = model.to(device)
model_output = model.generate(
    input_features,
    language=language,
)
predict_ids = model_output[0]
pred = processor.decode(predict_ids, skip_special_tokens=True)
print(pred)