Hugging Face Inference Endpointsによる強力なASR + ダイアライゼーション + 投機的デコーディング

私たちはオープンソースとオープンサイエンスを通じて人工知能を推進し、民主化する旅を続けています。

Hugging Face Blog · May 1, 2024

Powerful ASR + diarization + speculative decoding with Hugging Face Inference Endpoints

この記事は中国語版简体中文でもご覧いただけます。

Whisperは最高のオープンソース音声認識モデルの1つであり、間違いなく最も広く使われているモデルです。Hugging FaceのInference Endpointsを使えば、任意のWhisperモデルをすぐに簡単にデプロイできます。ただし、話者を識別するダイアライゼーションパイプラインや、投機的デコーディングのためのアシスト生成などの追加機能を導入したい場合、作業はより複雑になります。その理由は、Whisperを他のモデルと組み合わせつつ、単一のAPIエンドポイントを公開する必要があるためです。

この課題は、カスタム推論ハンドラーを使用することで解決できます。これによりInference Endpoints上で自動音声認識（ASR）とダイアライゼーションパイプラインを実装し、投機的デコーディングもサポートします。ダイアライゼーションパイプラインの実装は、有名なInsanely Fast Whisperに着想を得ており、ダイアライゼーションにはPyannoteモデルを使用しています。

これはまた、Inference Endpointsの柔軟性を示すデモンストレーションでもあり、ほぼあらゆるものをホストできることを示しています。こちらが参照用のコードです。エンドポイントの初期化時にはリポジトリ全体がマウントされるため、すべてのロジックを単一のファイルに収めたくない場合は、handler.pyからリポジトリ内の他のファイルを参照できます。このケースでは、コードを整理するために複数のファイルに分割することにしました：

handler.py：初期化および推論コードを含む
diarization_utils.py：ダイアライゼーション関連の前処理・後処理をすべて含む
config.py：ModelSettingsおよびInferenceConfigを含む。ModelSettingsはパイプラインで使用するモデルを定義し（すべてを使用する必要はありません）、InferenceConfigはデフォルトの推論パラメータを定義します

PyTorch 2.2以降、SDPAはFlash Attention 2を標準でサポートしているため、高速な推論のためにこのバージョンを使用します。

主なモジュール

以下は、エンドポイントの内部構造を示す高レベルな図です：

ASRおよびダイアライゼーションパイプラインの実装は、幅広いユースケースに対応できるようにモジュール化されており、ダイアライゼーションパイプラインはASRの出力上で動作し、ダイアライゼーションが不要な場合はASR部分のみを使用できます。ダイアライゼーションには、現在SOTAのオープンソース実装であるPyannoteモデルの使用を提案します。

また、推論を高速化する方法として投機的デコーディングも追加します。高速化は、より小さく高速なモデルを使って生成候補を提案し、それをより大きなモデルで検証することで実現されます。Whisper特有の動作については、こちらの優れたブログ記事で詳しく解説されています。

投機的デコーディングには以下の制約があります：

少なくともアシスタントモデルのデコーダー部分がメインモデルと同じアーキテクチャであること
バッチサイズは1でなければならない

上記の点を考慮してください。プロダクションのユースケースによっては、より大きなバッチをサポートする方が投機的デコーディングよりも高速になる場合があります。アシスタントモデルを使用しない場合は、設定内のassistant_modelをNoneのままにしてください。

アシスタントモデルを使用する場合、Whisperには蒸留版が優れた選択肢となります。

独自のエンドポイントをセットアップする

最も簡単な方法は、リポジトリ複製ツールを使ってカスタムハンドラーリポジトリをクローンすることです。

以下はhandler.pyからのモデル読み込み部分です：

from pyannote.audio import Pipeline
from transformers import pipeline, AutoModelForCausalLM

...

self.asr_pipeline = pipeline(
      "automatic-speech-recognition",
      model=model_settings.asr_model,
      torch_dtype=torch_dtype,
      device=device
  )

  self.assistant_model = AutoModelForCausalLM.from_pretrained(
      model_settings.assistant_model,
      torch_dtype=torch_dtype,
      low_cpu_mem_usage=True,
      use_safetensors=True
  ) 
  
  ...

  self.diarization_pipeline = Pipeline.from_pretrained(
      checkpoint_path=model_settings.diarization_model,
      use_auth_token=model_settings.hf_token,
  ) 
  
  ...

必要に応じてパイプラインをカスタマイズできます。config.pyファイル内のModelSettingsは初期化に使用するパラメータを保持し、推論中に使用するモデルを定義します：

class ModelSettings(BaseSettings):
    asr_model: str
    assistant_model: Optional[str] = None
    diarization_model: Optional[str] = None
    hf_token: Optional[str] = None

これらのパラメータは、対応する名前の環境変数を渡すことで調整できます。これはカスタムコンテナと推論ハンドラーの両方で機能します。これはPydanticの機能です。ビルド時にコンテナへ環境変数を渡すには、APIコール経由でエンドポイントを作成する必要があります（インターフェース経由では作成できません）。

環境変数として渡す代わりにモデル名をハードコードすることも可能ですが、ダイアライゼーションパイプラインにはトークンを明示的に渡す必要がある（hf_token）ことに注意してください。セキュリティ上の理由からトークンをハードコードすることはできません。つまり、ダイアライゼーションモデルを使用するにはAPIコール経由でエンドポイントを作成する必要があります。

念のため、ダイアライゼーション関連の前処理・後処理ユーティリティはすべてdiarization_utils.pyにあります。

必須コンポーネントはASRモデルのみです。オプションで、投機的デコーディングに使用するアシスタントモデルや、話者ごとに書き起こしを分割するダイアライゼーションモデルを指定できます。

Inference Endpointsへのデプロイ

ASR部分のみが必要な場合は、config.pyでasr_model/assistant_modelを指定し、ワンクリックでデプロイできます：

Inference Endpointsでホストされるコンテナに環境変数を渡すには、提供されているAPIを使用してプログラム的にエンドポイントを作成する必要があります。以下はコール例です：

body = {
    "compute": {
        "accelerator": "gpu",
        "instanceSize": "medium",
        "instanceType": "g5.2xlarge",
        "scaling": {
            "maxReplica": 1,
            "minReplica": 0
        }
    },
    "model": {
        "framework": "pytorch",
        "image": {
            # デフォルトコンテナ
            "huggingface": {
                "env": {
            # ここにHubモデルがマウントされます
                    "HF_MODEL_DIR": "/repository", 
                    "DIARIZATION_MODEL": "pyannote/speaker-diarization-3.1",
                    "HF_TOKEN": "<your_token>",
                    "ASR_MODEL": "openai/whisper-large-v3",
                    "ASSISTANT_MODEL": "distil-whisper/distil-large-v3"
                }
            }
        },
        # Hub上のモデルリポジトリ
        "repository": "sergeipetrov/asrdiarization-handler",
        "task": "custom"
    },
    # エンドポイント名
    "name": "asr-diarization-1",
    "provider": {
        "region": "us-east-1",
        "vendor": "aws"
    },
    "type": "private"
}

アシスタントモデルを使用する場合

アシスタントモデルを使用するメリットをよりよく理解するため、k6で実施したベンチマークを以下に示します：

# セットアップ:
# GPU: A10
ASR_MODEL=openai/whisper-large-v3
ASSISTANT_MODEL=distil-whisper/distil-large-v3

# long: 60秒音声; short: 8秒音声
long_assisted..................: avg=4.15s    min=3.84s    med=3.95s    max=6.88s    p(90)=4.03s    p(95)=4.89s   
long_not_assisted..............: avg=3.48s    min=3.42s    med=3.46s    max=3.71s    p(90)=3.56s    p(95)=3.61s   
short_assisted.................: avg=326.96ms min=313.01ms med=319.41ms max=960.75ms p(90)=325.55ms p(95)=326.07ms
short_not_assisted.............: avg=784.35ms min=736.55ms med=747.67ms max=2s       p(90)=772.9ms  p(95)=774.1ms

ご覧のとおり、アシスト生成は音声が短い場合（バッチサイズが1）に劇的なパフォーマンス向上をもたらします。音声が長い場合、推論は自動的にバッチに分割するため、以前に述べた制約により投機的デコーディングが推論時間を悪化させる可能性があります。

推論パラメータ

すべての推論パラメータはconfig.pyにあります：

class InferenceConfig(BaseModel):
    task: Literal["transcribe", "translate"] = "transcribe"
    batch_size: int = 24
    assisted: bool = False
    chunk_length_s: int = 30
    sampling_rate: int = 16000
    language: Optional[str] = None
    num_speakers: Optional[int] = None
    min_speakers: Optional[int] = None
    max_speakers: Optional[int] = None

もちろん、必要に応じてパラメータを追加または削除できます。話者数に関連するパラメータはダイアライゼーションパイプラインに渡され、その他は主にASRパイプライン用です。sampling_rateは処理する音声のサンプリングレートを示し、前処理に使用されます。assistedフラグは投機的デコーディングを使用するかどうかをパイプラインに指示します。アシスト生成では、batch_sizeを1に設定する必要があることに注意してください。

ペイロード

デプロイ完了後、以下のように（Pythonで）音声と推論パラメータを推論エンドポイントに送信します：

import base64
import requests

API_URL = "<your endpoint URL>"
filepath = "/path/to/audio"

with open(filepath, "rb") as f:
    audio_encoded = base64.b64encode(f.read()).decode("utf-8")

data = {
    "inputs": audio_encoded,
    "parameters": {
        "batch_size": 24
    }
}

resp = requests.post(API_URL, json=data, headers={"Authorization": "Bearer <your token>"})
print(resp.json())

ここで"parameters"フィールドは、InferenceConfigから調整したいすべてのパラメータを含む辞書です。InferenceConfigで指定されていないパラメータは無視されることに注意してください。

またはInferenceClientを使用します（非同期版もあります）：

from huggingface_hub import InferenceClient

client = InferenceClient(model = "<your endpoint URL>", token="<your token>")

with open("/path/to/audio", "rb") as f:
    audio_encoded = base64.b64encode(f.read()).decode("utf-8")
data = {
    "inputs": audio_encoded,
    "parameters": {
        "batch_size": 24
    }
}

res = client.post(json=data)

まとめ

このブログでは、Hugging Face Inference Endpointsを使用してモジュール化されたASR + ダイアライゼーション + 投機的デコーディングパイプラインをセットアップする方法について説明しました。パイプラインを必要に応じて簡単に設定・調整できるように最善を尽くし、Inference Endpointsでのデプロイは常に非常に簡単です。実装に使用した、コミュニティにオープンに提供されている優れたモデルとツールに恵まれていることを幸運に思います：

OpenAIによるWhisperモデルファミリー
Pyannoteによるダイアライゼーションモデル
主な着想源となったInsanely Fast Whisperリポジトリ

同じパイプラインをサーバー部分（FastAPI+Uvicorn）とともに実装したリポジトリがあります。さらにカスタマイズしたい場合や、別の場所にホストしたい場合に役立つかもしれません。

Read the original source Next news