音声認識技術の進歩は目覚ましく、最近では精度と速度の両面で大きな改善が見られています。今回は、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とは

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)

実行結果と感想

実際に上記のスクリプトを実行してみたところ、普通にCPU環境でも問題なく起動することができました。これは、Whisper Medusaが比較的軽量で、特別な環境がなくても動作可能であることを示しています。

ただし、CPU環境での実行では、まだ通常のWhisperモデルと比較して顕著な速度向上を体感するまでには至りませんでした。これは、Whisper Medusaの真の力を発揮するには、より高性能なハードウェア、特にGPUが必要である可能性を示唆しています。

今後は、GPU環境でも試してみる予定です。GPU環境での実行により、開発元が主張する1.5倍の速度向上を実際に体感できるのではないかと期待しています。

この初期の結果は、Whisper Medusaの潜在的な可能性を示唆していますが、同時に、最適なパフォーマンスを得るための適切な環境設定の重要性も浮き彫りにしています。今後のテストでは、特にGPU環境での性能に注目し、Whisper Medusaの真の力を探っていきたいと思います。

まとめ

Whisper Medusaは、高速かつ高精度な音声認識を実現する優れたモデルです。特に英語の音声認識タスクにおいては、速度と精度のバランスが取れた選択肢となるでしょう。ただし、他言語での使用や長時間の音声処理には制限があるため、使用目的に応じて適切なモデルを選択することが重要です。

今後のアップデートで、これらの制限が改善されることを期待しています。音声認識技術の進歩は目覚ましく、Whisper Medusaのような革新的なモデルが登場することで、さまざまな分野での応用が広がっていくことでしょう。