使用 Hugging Face 推理終端搭建強大的“語音識別 + 說話人分割 + 投機解碼”工作流

HuggingFace發表於2024-06-06
Hugging Face

Whisper 是當前最先進的開源語音識別模型之一，毫無疑問，也是應用最廣泛的模型。如果你想部署 Whisper 模型，Hugging Face 推理終端 能夠讓你開箱即用地輕鬆部署任何 Whisper 模型。但是，如果你還想疊加其它功能，如用於分辨不同說話人的說話人分割，或用於投機解碼的輔助生成，事情就有點麻煩了。因為此時你需要將 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，因此本例使用 PyTorch 2.2 以加速推理。

主要模組

下圖展示了我們設計的方案的系統框圖:

系統框圖

在實現時，ASR 和說話人分割流水線採用了模組化的方法，因此是可重用的。說話人分割流水線是基於 ASR 的輸出的，如果不需要說話人分割，則可以僅用 ASR 的部分。我們建議使用 Pyannote 模型 做說話人分割，該模型目前是開源模型中的 SOTA。

我們還使用了投機解碼以加速模型推理。投機解碼透過使用更小、更快的模型來打草稿，再由更大的模型來驗證，從而實現加速。具體請參閱 這篇精彩的博文 以詳細瞭解如何對 Whisper 模型使用投機解碼。

投機解碼有如下兩個限制:

  • 輔助模型和主模型的解碼器的架構應相同
  • 在很多實現中，batch size 須為 1

在評估是否使用投機解碼時，請務必考慮上述因素。根據實際用例不同，有可能支援較大 batch size 帶來的收益比投機解碼更大。如果你不想使用輔助模型，只需將配置中的 assistant_model 置為 None 即可。

如果你決定使用輔助模型，distil-whisper 是一個不錯的 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 呼叫 (而不是透過 GUI) 建立終端。

你還可以在程式碼中硬編碼模型名，而不將其作為環境變數傳入，但 請注意，說話人分割流水線需要顯式地傳入 HF 令牌 (hf_token )。 出於安全考量，我們不允許對令牌進行硬編碼，這意味著你必須透過 API 呼叫建立終端才能使用說話人分割模型。

提醒一下，所有與說話人分割相關的預處理和後處理工具程式都在 diarization_utils.py 中。

該方案中，唯一必選的元件是 ASR 模型。可選項是: 1) 投機解碼，你可指定一個輔助模型用於此; 2) 說話人分割模型，可用於對轉錄文字按說話人進行分割。

部署至推理終端

如果僅需 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 進行了一系列基準測試，如下:

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

# 長音訊: 60s; 短音訊: 8s
長音訊 _ 投機解碼 ..................: avg=4.15s min=3.84s med=3.95s max=6.88s p(90)=4.03s p(95)=4.89s
長音訊 _ 直接解碼 ..............: avg=3.48s min=3.42s med=3.46s max=3.71s p(90)=3.56s p(95)=3.61s
短音訊 _ 輔助解碼 .................: avg=326.96ms min=313.01ms med=319.41ms max=960.75ms p(90)=325.55ms p(95)=326.07ms
短音訊 _ 直接解碼 .............: avg=784.35ms min=736.55ms med=747.67ms max=2s p(90)=772.9ms p(95)=774.1ms

如你所見，當音訊較短 (batch size 為 1) 時，輔助生成能帶來顯著的效能提升。如果音訊很長，推理系統會自動將其切成多 batch，此時由於上文述及的限制，投機解碼可能會拖慢推理。

推理引數

所有推理引數都在 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 推理終端搭建模組化的 “ASR + 說話人分割 + 投機解碼”工作流。該方案使用了模組化的設計，使使用者可以根據需要輕鬆配置並調整流水線，並輕鬆地將其部署至推理終端！更幸運的是，我們能夠基於社群提供的優秀公開模型及工具實現我們的方案:

  • OpenAI 的一系列 Whisper 模型
  • Pyannote 的 說話人分割模型
  • Insanely Fast Whisper 程式碼庫，這是本文的主要靈感來源

本文相關的程式碼已上傳至 這個程式碼庫中，其中包含了本文論及的流水線及其服務端程式碼 (FastAPI + Uvicorn)。如果你想根據本文的方案進一步進行定製或將其託管到其他地方，這個程式碼庫可能會派上用場。

英文原文: https://hf.co/blog/asr-diarization

原文作者: Sergei Petrov，Vaibhav Srivastav，Pedro Cuenca，Philipp Schmid

譯者: Matrix Yao (姚偉峰)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。

相關文章