faster-whisper+Dockerで音声からテキスト変換(STT)のAPIを実装する

AI Python

初めに

音声認識をしたい場合whisperを使うことが多いですが、より速くより使いやすくしたいと思ってたので実装をしてみました！

DockerでCUDAのver管理やGPUも使えるようにして、whisperモデルのキャッシュもしているのでかなり使いやすくなりました

OSSとして以下で公開しています

github.com

デモ

実際に動かすと以下のように動きます

開発環境

  • Docker Hub

実装

環境を作る

まずは、動かすための環境を作ります ver等を合わせるために、Dockerを使います

# 指定されたNVIDIA CUDAイメージ
FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04

# 環境変数の設定
ENV LANG C.UTF-8
ENV LC_ALL C.UTF-8
ENV PATH /usr/local/nvidia/bin:/usr/local/cuda/bin:${PATH}
ENV LD_LIBRARY_PATH /usr/local/nvidia/lib:/usr/local/nvidia/lib64

# 必要なパッケージのインストール
RUN apt-get update && apt-get install -y \
    python3-pip \
    python3-dev

# 作業ディレクトリの設定
WORKDIR /app

# 必要なPythonライブラリをインストール
COPY requirements.txt /app/
RUN pip3 install --no-cache-dir -r requirements.txt

# アプリケーションのコードをコピー
COPY . /app

# ポートを公開
EXPOSE 8000

# FastAPIサーバーを起動
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

毎回モデルをダウンロードするのは大変なので、Dockerのvolumesにキャッシュできるようにしていきます。またGPUで動かしたいので設定します。これらを満たすために、docker-composeを使ってDockerfileや環境を管理します

version: '3.8'
services:
  app:
    build: .
    ports:
      - "8000:8000"
    volumes:
      - faster_whisper:/models
    command: uvicorn main:app --host 0.0.0.0 --port 8000
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [ gpu ]
    tty: true

volumes:
  faster_whisper:

モデルのロードと初期化

環境ができたので、以下を実装します 1. GPUが使えるどうかの確認。使えない場合はCPU処理 2. モデルをdockerのvaluesにキャッシュしてそこからロード

以下の関数で初期化等を行います

ef initialize_model():
    model_path = "/models/whisper-large-v3"
    if torch.cuda.is_available():
        print("CUDA is available")
        return WhisperModel("large-v3", device="cuda", compute_type="float16", download_root=model_path)
    else:
        print("CUDA is not available or not enabled")
        cpu_threads = os.cpu_count()
        return WhisperModel("large-v3", device="cpu", compute_type="int8", cpu_threads=cpu_threads, download_root=model_path)

STTのAPIの作成

音声データはバイナリーデータでwhisper側に渡したいので、file形式でAPIを実装します

ただfileからバイナリーIOの変換は、こちらでやっていきます

@app.post("/transcribe")
async def transcribe_audio(file: UploadFile = Form(...)):
    try:
        # ファイルの内容をバイナリデータとして読み込む
        file_content = await file.read()

        # バイナリデータをBinaryIOオブジェクトに変換
        file_stream = io.BytesIO(file_content)

        # 音声ファイルの文字起こし
        segments, info = model.transcribe(
            audio=file_stream,  # BinaryIOオブジェクトを渡す
            beam_size=5,
            vad_filter=True,
            without_timestamps=True,
        )

        # 結果のフォーマット
        result = "Detected language '%s' with probability %f\n" % (info.language, info.language_probability)
        for segment in segments:
            result += "[%.2fs -> %.2fs] %s\n" % (segment.start, segment.end, segment.text)

        return {"transcription": result}
    except Exception as e:
        return {"error": str(e)}