Inferência de modelos usando Hugging Face Transformers para NLP

Este artigo mostra como usar Hugging Face Transformers para inferência de modelo de processamento de linguagem natural (NLP).

Os transformadores Hugging Face fornecem a classe de pipelines para usar o modelo pré-treinado para inferência. 🤗 Os pipelines Transformers oferecem suporte a uma ampla variedade de tarefas NLP que você pode usar facilmente no Databricks.

Requisitos

  • MLflow 2.3

  • Quaisquer clusters com a biblioteca Hugging Face transformers instalada podem ser usados para inferência de lotes. A biblioteca transformers vem pré-instalada no Databricks Runtime 10.4 LTS MLe acima. Muitos dos modelos populares de NLP funcionam melhor em hardware de GPU, portanto, você pode obter o melhor desempenho usando hardware de GPU recente, a menos que use um modelo especificamente otimizado para uso em CPUs.

Use Pandas UDFs para distribuir computação de modelo em clustersSpark

Ao experimentar modelos pré-treinados, você pode usar Pandas UDFs para agrupar o modelo e realizar cálculos em CPUs ou GPUs worker . Pandas UDFs distribuem o modelo para cada worker.

Você também pode criar um pipeline Hugging Face Transformers para tradução automática e usar um Pandas UDF para executar o pipeline no worker de clusters 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)

Definir o device dessa maneira garante que as GPUs sejam usadas se estiverem disponíveis nos clusters.

Os pipelines Hugging Face para tradução retornam uma lista de objetos Python dict, cada um com uma única key translation_text e um valor contendo o texto traduzido. Este UDF extrai a tradução dos resultados para retornar uma série Pandas com apenas o texto traduzido. Se seu pipeline foi construído para usar GPUs definindo device=0, o Spark reatribui automaticamente GPUs nos nós worker se seus clusters tiverem instâncias com várias GPUs.

Para usar o UDF para traduzir uma coluna de texto, você pode chamar o UDF em uma instrução select :

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')))

Retornar tipos de resultados complexos

Usando Pandas UDFs, você também pode retornar uma saída mais estruturada. Por exemplo, no reconhecimento de entidade nomeada, os pipelines retornam uma lista de objetos dict contendo a entidade, sua extensão, tipo e uma pontuação associada. Embora semelhante ao exemplo de tradução, o tipo de retorno para a anotação @pandas_udf é mais complexo no caso de reconhecimento de entidade nomeada.

Você pode ter uma noção dos tipos de retorno a serem usados por meio da inspeção dos resultados do pipeline, por exemplo, executando o pipeline no driver.

Neste exemplo, use o seguinte código:

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)

Para render a anotação:

[[{'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}]]

Para representar isso como um tipo de retorno, você pode usar um array de struct campos, listando as entradas dict como os campos do 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')))

Ajustar o desempenho

Há vários aspectos key para ajustar o desempenho do UDF. A primeira é usar cada GPU de forma eficaz, o que você pode ajustar alterando o tamanho dos lotes enviados à GPU pelo pipeline do Transformers. A segunda é garantir que o DataFrame seja bem particionado para utilizar todos os clusters.

Finalmente, você pode querer armazenar em cache o modelo Hugging Face para economizar tempo de carregamento do modelo ou custos de entrada.

Escolha um tamanho de lote

Embora as UDFs descritas acima devam funcionar prontas para uso com um batch_size de 1, isso pode não usar os recursos disponíveis para o worker de forma eficiente. Para melhorar o desempenho, ajuste o tamanho dos lotes ao modelo e hardware nos clusters. Databricks recomenda tentar vários tamanhos de lote para o pipeline em seus clusters para encontrar o melhor desempenho. Leia mais sobre lotes de pipeline e outras opções de desempenho na documentação do Hugging Face.

Tente encontrar um tamanho de lote grande o suficiente para conduzir a utilização total da GPU, mas não resultar em erros CUDA out of memory . Ao receber erros CUDA out of memory durante o ajuste, você precisa desconectar e reconectar o Notebook para liberar a memória usada pelo modelo e os dados na GPU.

Ajuste o paralelismo com programar em nível de palco

Por default, o Spark programa uma tarefa por GPU em cada máquina. Para aumentar o paralelismo, você pode usar o programar em nível de estágio para informar ao Spark quantas tarefas devem ser executadas por GPU. Por exemplo, se você quiser que o Spark execute duas tarefas por GPU, especifique isso da seguinte maneira:

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)

Dados de repartição para usar todo o hardware disponível

A segunda consideração para o desempenho é fazer uso total do hardware em seus clusters. Geralmente, um pequeno múltiplo do número de GPUs em seu worker (para clusters de GPU) ou número de núcleos no worker em seu cluster (para clusters de CPU) funciona bem. Seu DataFrame de entrada pode já ter partições suficientes para aproveitar o paralelismo dos clusters . Para ver quantas partições o DataFrame contém, use df.rdd.getNumPartitions(). Você pode reparticionar um DataFrame usando repartitioned_df = df.repartition(desired_partition_count).

Cache o modelo em DBFS ou em pontos de montagem

Se você estiver carregando frequentemente um modelo de clusters diferentes ou reiniciados, você também pode querer armazenar em cache o modelo Hugging Face no volumeDBFS root ou em um ponto de montagem. Isso pode diminuir os custos de entrada e reduzir o tempo de carregamento do modelo em clusters novos ou reiniciados. Para fazer isso, defina a variável de ambiente TRANSFORMERS_CACHE em seu código antes de carregar o pipeline.

Por exemplo:

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

Como alternativa, você pode obter resultados semelhantes registrando o modelo no MLflow com o sabor `transformers` do MLflow.

Notebook: Hugging Face Transformers e o log do MLflow

Para começar rapidamente com o código de exemplo, este Notebook é um exemplo de ponta a ponta para resumo de texto usando a inferência de pipelines do Hugging Face Transformers e o log de MLflow.

Notebook de inferência de pipelines Hugging FaceTransformers

Abra o bloco de anotações em outra guia

Recursos adicionais

Você pode ajustar seu modelo Hugging Face com os seguintes guias:

Saiba mais sobre o que são transformadores Hugging Face ?