共用方式為


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

重要

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

本文說明如何使用 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?