使用 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 模型: