使用 Hugging Face Transformers 进行 NLP 的模型推理

本文介绍如何使用 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 硬件上可以保持最佳工作状态，因此你可以使用最新的 GPU 硬件来获得最佳性能，除非使用专门针对 CPU 优化的模型。

使用 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}]]

若要将此数据表示为返回类型，可以使用 struct 字段的 array，将 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，对此可以通过更改转换器管道发送到 GPU 的批大小来进行调整。 第二个方面是确保数据帧已合理分区以利用整个群集。

最后，建议缓存 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 群集）的较小倍数就能获得不错的效果。 输入数据帧可能已经有足够的分区来利用群集的并行度。 若要查看数据帧包含多少个分区，请使用 df.rdd.getNumPartitions()。 可以使用 repartitioned_df = df.repartition(desired_partition_count) 将数据帧重新分区。

将模型缓存在 DBFS 或装载点上

如果你经常从不同群集或重启的群集加载模型，则还建议将 Hugging Face 模型缓存在 DBFS 根卷或某个装载点上。 这可以降低流入成本，并减少在新群集或重启的群集上加载模型所需的时间。 为此，请在加载管道之前，先在代码中设置 TRANSFORMERS_CACHE 环境变量。

例如：

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

也可通过使用 MLflow transformers 风格将模型记录到 MLflow 来获得类似的结果。

笔记本：Hugging Face 转换器推理和 MLflow 日志记录

为了让你快速开始使用示例代码，此笔记本提供了一个端到端示例，它使用 Hugging Face 转换器管道推理和 MLflow 日志记录进行文本汇总。

Hugging Face 转换器管道推理笔记本

获取笔记本

其他资源

可按照以下指南微调 Hugging Face 模型：

• 准备数据以微调 Hugging Face 模型
• 为单个 GPU 微调 Hugging Face 模型

详细了解什么是 Hugging Face Transformers？