使用 Hugging Face Transformers 進行 NLP 的模型推斷

文章
11/08/2024

重要

此文件已淘汰，且可能未更新。不再支援此內容中提及的產品、服務或技術。
Databricks 建議 ai_query 改用批次推斷。請參閱使用ai_query執行批次推斷。

本文說明如何使用 Hugging Face Transformers 進行自然語言處理 (NLP) 模型推斷。

Hugging Face Transformers 提供管線類別來使用預先訓練的模型進行推斷。 🤗 Transformers 管線支援廣泛的 NLP 工作，您可以在 Azure Databricks 上輕鬆使用這些工作。

需求

MLflow 2.3
已安裝 Hugging Face transformers 程式庫的任何叢集都可用於批次推斷。 transformers 程式庫已預先安裝於 Databricks Runtime 10.4 LTS ML 和更新版本。許多熱門的 NLP 模型最適合 GPU 硬體，因此，除非您使用特別針對 CPU 使用最佳化的模型，否則您可以使用最新的 GPU 硬體取得最佳效能。

使用 Pandas UDF 在 Spark 叢集上散發模型計算

在體驗預先訓練的模型時，您可以使用 Pandas UDF 來包裝模型並在背景工作角色 CPU 或 GPU 上執行計算。 Pandas UDF 會將模型散發給每個背景工作角色。

您也可以建立用於機器翻譯的 Hugging Face Transformers 管線，並使用 Pandas UDF 在 Spark 叢集的背景工作角色上執行此管線：

import pandas as pd
from transformers import pipeline
import torch
from pyspark.sql.functions import pandas_udf

device = 0 if torch.cuda.is_available() else -1
translation_pipeline = pipeline(task="translation_en_to_fr", model="t5-base", device=device)

@pandas_udf('string')
def translation_udf(texts: pd.Series) -> pd.Series:
  translations = [result['translation_text'] for result in translation_pipeline(texts.to_list(), batch_size=1)]
  return pd.Series(translations)

以此方式設定 device 可確保使用叢集上可用的 GPU。

用於翻譯的 Hugging Face 管線可傳回 Python dict 物件的清單，其中每個物件都有一個 translation_text 索引鍵和一個包含翻譯文字的值。此 UDF 從結果中擷取翻譯，以傳回僅包含翻譯文字的 Pandas 系列。如果透過設定 device=0 將管線建構為使用 GPU，則在叢集中的執行個體包含多個 GPU 的情況下，Spark 會自動在背景工作角色節點上重新指派 GPU。

若要使用 UDF 翻譯文字資料行，您可以在 select 陳述式中呼叫 UDF：

texts = ["Hugging Face is a French company based in New York City.", "Databricks is based in San Francisco."]
df = spark.createDataFrame(pd.DataFrame(texts, columns=["texts"]))
display(df.select(df.texts, translation_udf(df.texts).alias('translation')))

傳回複雜結果類型

使用 Pandas UDF，您也可以傳回更多結構化輸出。例如，在具名實體辨識中，管線會傳回 dict 物件清單，其中包含實體、其範圍、類型和關聯的分數。雖然與翻譯範例類似，但在具名實體辨識案例中，@pandas_udf 註釋的傳回類型更複雜。

您可以透過檢查管線結果來了解要使用的傳回類型，例如透過在驅動程式上執行管線。

在此範例中，使用下列程式碼：

from transformers import pipeline
import torch
device = 0 if torch.cuda.is_available() else -1
ner_pipeline = pipeline(task="ner", model="Davlan/bert-base-multilingual-cased-ner-hrl", aggregation_strategy="simple", device=device)
ner_pipeline(texts)

若要產生註釋，請執行以下操作：

[[{'entity_group': 'ORG',
   'score': 0.99933606,
   'word': 'Hugging Face',
   'start': 0,
   'end': 12},
  {'entity_group': 'LOC',
   'score': 0.99967843,
   'word': 'New York City',
   'start': 42,
   'end': 55}],
 [{'entity_group': 'ORG',
   'score': 0.9996372,
   'word': 'Databricks',
   'start': 0,
   'end': 10},
  {'entity_group': 'LOC',
   'score': 0.999588,
   'word': 'San Francisco',
   'start': 23,
   'end': 36}]]

若要將此資料表示為傳回類型，您可以使用 array 欄位的 struct，將 dict 項目列為 struct 的欄位：

import pandas as pd
from pyspark.sql.functions import pandas_udf

@pandas_udf('array<struct<word string, entity_group string, score float, start integer, end integer>>')
def ner_udf(texts: pd.Series) -> pd.Series:
  return pd.Series(ner_pipeline(texts.to_list(), batch_size=1))

display(df.select(df.texts, ner_udf(df.texts).alias('entities')))

調整效能

可以在多個重要方面調整 UDF 的效能。第一個方面是有效使用每個 GPU，對此可以透過變更 Transformers 管線傳送至 GPU 的批次大小來進行調整。第二個方面是確保 DataFrame 已妥善分割以利用整個叢集。

最後，您可能想要快取 Hugging Face 模型，以節省模型載入時間或輸入成本。

選擇批次大小

雖然使用值為 1 的 batch_size 就能現成地使用上述 UDF，但這樣可能無法有效利用可供背景工作角色使用的資源。若要改善效能，請將批次大小調整為叢集中的模型和硬體。 Databricks 建議為叢集上的管線嘗試各種批次大小，以找出最佳效能。在 Hugging Face 文件中深入了解管線批次處理和其他效能選項。

嘗試找到一個足夠大的批次大小，以全面促進 GPU 的利用，但不會導致 CUDA out of memory 錯誤。如果在調整期間收到 CUDA out of memory 錯誤，您需要中斷連結再重新連結筆記本，以釋放模型和 GPU 中的資料使用的記憶體。

透過檢視叢集的即時叢集計量並選擇某個計量 (例如，表示 GPU 處理器利用率的 gpu0-util，或表示 GPU 記憶體利用率的 gpu0_mem_util)，來監視 GPU 效能。

使用階段層級調整平行處理原則

根據預設，Spark 在每部機器上為每個 GPU 排程一項工作。若要提高平行處理原則，您可以使用階段層級排程來告知 Spark 每個 GPU 要執行的工作數。例如，如果希望 Spark 為每個 GPU 執行兩項工作，您可以按下列方式進行指定：

from pyspark.resource import TaskResourceRequests, ResourceProfileBuilder

task_requests = TaskResourceRequests().resource("gpu", 0.5)

builder = ResourceProfileBuilder()
resource_profile = builder.require(task_requests).build

rdd = df.withColumn('predictions', loaded_model(struct(*map(col, df.columns)))).rdd.withResources(resource_profile)

重新分割資料以使用所有可用的硬體

效能的第二個考量是充分利用叢集中的硬體。一般情況下，使用背景工作角色上 GPU 數目 (對於 GPU 叢集) 或叢集中背景工作角色上的核心數目 (對於 CPU 叢集) 的較小倍數就能取得不錯的效果。輸入 DataFrame 可能已有足夠的分割區來利用叢集的平行處理原則。若要檢視 DataFrame 包含的分割區數量，請使用 df.rdd.getNumPartitions()。您可以使用 repartitioned_df = df.repartition(desired_partition_count) 重新分割 DataFrame。

將模型快取在 DBFS 或掛接點上

如果經常從不同叢集或重新啟動的叢集載入模型，您可能還希望將 Hugging Face 模型快取在 DBFS 根磁碟區中或掛接點上。這可以降低輸入成本，並減少在新叢集或重新啟動的叢集上載入模型所需的時間。為此，請在載入管線之前，先在程式碼中設定 TRANSFORMERS_CACHE 環境變數。

例如：

import os
os.environ['TRANSFORMERS_CACHE'] = '/dbfs/hugging_face_transformers_cache/'

或者，您可以透過使用 MLflow transformers 類別將模型記錄至 MLflow 來達成類似結果。

筆記本：Hugging Face Transformers 推斷和 MLflow 記錄

為了快速開始使用範例程式碼，此筆記本提供了一個端對端範例，它使用 Hugging Face Transformers 管線推斷和 MLflow 記錄進行文字摘要。

Hugging Face Transformers 管線推斷筆記本

取得筆記本

其他資源

您可以使用下列指南來微調 Hugging Face 模型：

深入了解什麼是 Hugging Face Transformers？

分享方式：