藉助 Hugging Face 推理端點實現強大的 ASR + 對話分離 + 推測解碼

釋出於 2024 年 5 月 1 日

在 GitHub 上更新

贊

Whisper 是最優秀的開源語音識別模型之一，也無疑是使用最廣泛的模型。Hugging Face 推理端點讓部署任何 Whisper 模型都變得非常簡單。然而，如果你想引入額外功能，比如用於識別說話者的對話分離（diarization）流水線，或者用於推測解碼的輔助生成，事情就會變得複雜。原因是，你需要將 Whisper 與其他模型結合起來，同時仍然只暴露一個 API 端點。

我們將使用自定義推理處理程式來解決這個挑戰，它將在推理端點上實現自動語音識別（ASR）和對話分離流水線，並支援推測解碼。對話分離流水線的實現受到了著名的 Insanely Fast Whisper 的啟發，並使用 Pyannote 模型進行對話分離。

這也將展示推理端點的靈活性，以及你幾乎可以在上面託管任何東西。這裡是可供參考的程式碼。請注意，在端點初始化期間，整個程式碼倉庫都會被掛載，所以如果不想把所有邏輯都放在一個檔案中，你的 `handler.py` 可以引用倉庫中的其他檔案。在這種情況下，我們決定將內容分成幾個檔案以保持整潔。

handler.py 包含初始化和推理程式碼
diarization_utils.py 包含所有與對話分離相關的預處理和後處理
config.py 包含 ModelSettings 和 InferenceConfig。ModelSettings 定義了流水線中將使用哪些模型（不必全部使用），而 InferenceConfig 定義了預設的推理引數

從 Pytorch 2.2 開始，SDPA 原生支援 Flash Attention 2，所以我們將使用該版本以加快推理速度。

主要模組

這是端點內部結構的高層示意圖

ASR 和對話分離流水線的實現是模組化的，以適應更廣泛的用例——對話分離流水線在 ASR 輸出的基礎上執行，如果不需要對話分離，你可以只使用 ASR 部分。對於對話分離，我們建議使用 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 模型。可選地，可以指定一個輔助模型用於推測解碼，以及一個對話分離模型用於按說話者劃分轉錄文字。

在推理端點上部署

如果你只需要 ASR 部分，你可以在 `config.py` 中指定 `asr_model`/`assistant_model`，然後一鍵部署

要將環境變數傳遞給託管在推理端點上的容器，你需要使用提供的 API 以程式設計方式建立一個端點。以下是一個示例呼叫

body = {
    "compute": {
        "accelerator": "gpu",
        "instanceSize": "medium",
        "instanceType": "g5.2xlarge",
        "scaling": {
            "maxReplica": 1,
            "minReplica": 0
        }
    },
    "model": {
        "framework": "pytorch",
        "image": {
            # a default container
            "huggingface": {
                "env": {
            # this is where a Hub model gets mounted
                    "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"
                }
            }
        },
        # a model repository on the Hub
        "repository": "sergeipetrov/asrdiarization-handler",
        "task": "custom"
    },
    # the endpoint name
    "name": "asr-diarization-1",
    "provider": {
        "region": "us-east-1",
        "vendor": "aws"
    },
    "type": "private"
}

何時使用輔助模型

為了更好地說明何時使用輔助模型是有益的，這裡有一個使用 k6 進行的基準測試

# Setup:
# GPU: A10
ASR_MODEL=openai/whisper-large-v3
ASSISTANT_MODEL=distil-whisper/distil-large-v3

# long: 60s audio; short: 8s audio
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)