Hugging Faceトランスフォーマーのパイプラインを使ったNLPを始めよう

Original Blog : Getting started with NLP using Hugging Face transformers pipelines

翻訳： junichi.maruyama 

自然言語処理（NLP）の進歩は、企業がテキストデータから価値を引き出すための前例のない機会を解き放ちました。自然言語処理は、テキストの要約、人や場所などの固有名詞の認識、感情分類、テキスト分類、翻訳、質問応答など、幅広い用途に使用できます。多くの場合、大規模なテキストデータセットで事前に訓練された機械学習モデルから、高品質の結果を得ることができます。これらの事前学習済みモデルの多くは、オープンソースで公開されており、無料で使用することができます。Hugging Faceは、これらのモデルの素晴らしいソースの一つであり、彼らのTransformersライブラリは、モデルを適用し、また自分のデータにも適応させるための使いやすいツールです。また、これらのモデルを自分のデータに合わせて微調整をすることも可能です。

例えば、サポートチームを持つ企業では、訓練済みのモデルを使用して、人間が読めるテキストの要約を提供し、従業員がサポートケースの重要な問題を迅速に評価できるようにすることができます。また、この企業は、すぐに利用できる基礎モデルに基づいてワールドクラスの分類アルゴリズムを簡単にトレーニングし、サポートデータを社内のタクソノミに自動的に分類することができます。

Databricksは、Hugging Face Transformersを実行するための優れたプラットフォームです。これまでのDatabricksの記事では、事前学習済みモデルの推論とfine-tuningのためのトランスフォーマーの使用について説明してきましたが、この記事では、Lakehouse上でトランスフォーマーを使用する際のパフォーマンスと使いやすさを最適化するために、これらのベストプラクティスを統合します。また、Databricks では、事前学習済みモデルの推論とfine-tuning のための完全なノートブックの例も提供しています。

事前学習済みモデルの使用

センチメント分析やテキスト要約など、多くのアプリケーションでは、事前にトレーニングされたモデルは、追加のモデルトレーニングなしでうまく機能します。🤗 トランスフォーマーパイプラインは、テキストに対する推論に必要な様々なコンポーネントをシンプルなインターフェースで包み込みます。多くのNLPタスクでは、これらのコンポーネントはトークナイザーとモデルで構成されています。パイプラインはベストプラクティスをコード化することで、簡単に始めることができます。例えば、パイプラインは、GPUが利用可能な場合、GPUに送信されるアイテムのバッチ処理を可能にし、より良いスループットを実現します。

from transformers import pipeline
import torch
# use the GPU if available
device = 0 if torch.cuda.is_available() else -1
summarizer = pipeline("summarization", device=device) 

Spark上で推論を分散させるために、Databricksはパイプラインをpandas UDFでカプセル化することを推奨しています。Sparkはブロードキャストを使用して、pandas UDFが必要とするオブジェクトをワーカーノードに効率的に転送します。また、SparkはGPUをワーカーに自動的に再割り当てするので、マルチGPUマルチマシンクラスターをシームレスに使用することができます。

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

@pandas_udf('string')
def summarize_batch_udf(texts: pd.Series) -> pd.Series:
  pipe = summarizer(texts.to_list(), truncation=True, batch_size=1)
  summaries = [summary['summary_text'] for summary in pipe]
  return pd.Series(summaries)

summaries = df.select(df.text, summarize_batch_udf(df.text).alias("summary"))

以下は、バンドEnergy Orchardに関するWikipediaの記事のスナップショットに対する要約のサンプルである。サマライザーに渡す前に、Wikipediaのマークアップをクリーンアップしていないことに注意してください。

このセクションでは、Hugging Face Transformersを使用してDatabricks上で大規模なテキスト処理を開始することがいかに簡単であるかを示しました。次のセクションでは、これらのモデルのパフォーマンスをさらに調整する方法について説明します。

パフォーマンスチューニング

UDFのパフォーマンスチューニングには、2つのポイントがあります。1つ目は、各GPUを効率よく使いたいということで、これはTransformersパイプラインによってGPUに送られるアイテムのバッチサイズを変更することで調整することができます。もう1つは、データフレームをうまく分割してクラスタ全体を活用することです。

しかし、これではワーカーが利用できるリソースを効率的に利用できていない可能性があります。パフォーマンスを向上させるためには、使用しているモデルやハードウェアに合わせてバッチサイズを調整することができます。Databricksでは、クラスタ上のパイプラインのバッチサイズをいろいろ試してみて、最高のパフォーマンスを見つけることを推奨しています。パイプラインのバッチ処理とその他のパフォーマンスオプションの詳細については、Hugging Faceのドキュメントをご覧ください。クラスタのライブganglia metricsを表示し、GPUプロセッサ利用率の「gpu0-util」やGPUメモリ利用率の「gpu0_mem_util」など、メトリックを選択することでGPUパフォーマンスを監視することができます。

バッチサイズを調整する目的は、GPUをフルに活用しつつ、「CUDA out of memory」エラーが発生しない程度の大きさに設定することです。

クラスタのハードウェアをうまく利用するために、Spark DataFrameを再分割する必要がある場合があります。一般的に、ワーカーのGPU数（GPUクラスタの場合）またはクラスタ内のワーカー全体のコア数（CPUクラスタの場合）の倍数が、実際にうまく機能します。このUDFの使い方は、Sparkの他のUDFの使い方と同じです。例えば、select文で使用して、モデル推論の結果を持つ列を作成することができます。

sample = sample.repartition(32)
summaries = sample.select(sample.text, summarize_batch_udf(sample.text).alias("summary"))

構築済みのモデルをMLflowのモデルとしてラッピングする

事前にトレーニングされたモデルをMLflowモデルとして保存することで、バッチまたはリアルタイム推論のためのモデルのデプロイがさらに容易になります。また、モデルレジストリによるモデルのバージョン管理が可能になり、推論ワークロードのモデル読み込みコードを簡素化することができます。

パイプラインのカスタムモデルを作成し、モデルの読み込み、GPU使用量の初期化、推論機能などをカプセル化するのです。

import mlflow

class SummarizationPipelineModel(mlflow.pyfunc.PythonModel):
  def load_context(self, context):
    device = 0 if torch.cuda.is_available() else -1
    self.pipeline = pipeline("summarization", context.artifacts["pipeline"], device=device)
    
  def predict(self, context, model_input): 
    texts = model_input.iloc[:,0].to_list() # get the first column
    pipe = self.pipeline(texts, truncation=True, batch_size=8)
    summaries = [summary['summary_text'] for summary in pipe]
    return pd.Series(summaries)

このコードは、上記で説明したpandas_udfを作成し使用するためのコードと密接に類似しています。1つの違いは、パイプラインがMLflowモデルのコンテキストで利用可能になったファイルからロードされることです。これは、モデルをログに記録するときに MLflow に提供されます。Hugging Face transformers パイプラインは、モデルをドライバ上のローカルファイルに保存することを容易にし、それを MLflow pyfunc インターフェイスの log_model 関数に渡します。

summarizer.save_pretrained("./pipeline")
with mlflow.start_run() as run:        
  mlflow.pyfunc.log_model(artifacts={'pipeline': "./pipeline"}, artifact_path="summarization_model", python_model=SummarizationPipelineModel())

MLflowモデルを使ったバッチ推論

MLflowは、ログに記録されたモデルや登録されたモデルをspark UDFに読み込むための簡単なインターフェイスを提供します。Model Registryやlogged experiment run UIからモデルURIを検索することができます。

logged_model_uri = f"runs:/{run.info.run_id}/summarization_model"
loaded_model = mlflow.pyfunc.spark_udf(spark, model_uri=logged_model_uri, result_type='string')
summaries = df.select(df.title, df.text, loaded_model(df.text).alias("summary"))

Transformers Trainer🤗を使って、1台のマシンでモデルのfine-tuningを行う

訓練済みのモデルがそのままではニーズに応えられないこともあり、独自のデータでモデルのfine-tuningを行う必要があります。例えば、サポートチケットをサポートチームのオントロジーに分類する基礎モデルに基づいてテキスト分類器を作成したい場合や、自分のデータでカスタムスパム分類器を作成したい場合などがあります。

モデルのfine-tuningをするためにDatabricksを離れる必要はありません。適度な大きさのデータセットであれば、GPUをサポートした1台のマシンでこれを行うことができます。Hugging Face transformers Trainerユーティリティを使用すると、モデルトレーニングのセットアップと実行が非常に簡単になります。より大きなデータセットの場合、Databricksは分散型マルチマシン・マルチGPUディープラーニングもサポートしています。

手順としては、GPUをサポートしたシングルマシンのクラスタを作成し、データセットを準備してドライバーにダウンロードし、Trainerを使用してモデルトレーニングを実行し、結果のモデルをMLflowにログ出力します。

データの準備とダウンロード

まず、トレーニングデータをトレーナーの期待に沿うような表にフォーマットすることから始めます。テキスト分類の場合、これは2つの列を持つテーブルです：テキスト列とラベルの列。このexample notebook では、テキストメッセージのスパムデータを読み込んでいます：

Hugging Face transformersでは、テキスト分類のモデルローダーとしてAutoModelForSequenceClassificationを提供しており、カテゴリーラベルとして整数IDを想定しています。ただし、整数ラベルから文字列ラベルへのマッピングを指定する必要があります。文字列ラベルを持つDataFrameがあれば、次のように情報を収集することができます：

labels = sms.select(sms.label).groupBy(sms.label).count().collect()
id2label = {index: row.label for (index, row) in enumerate(labels)} 
label2id = {row.label: index for (index, row) in enumerate(labels)}

そして、整数のidをpandas_udfでラベルカラムとして作成します：

from pyspark.sql.functions import pandas_udf
import pandas as pd
@pandas_udf('integer')
def replace_labels_with_ids(labels: pd.Series) -> pd.Series:
  return labels.apply(lambda x: label2id[x])

sms_id_labels = sms.select(replace_labels_with_ids(sms.label).alias('label'), sms.text)

データをトレーニング/テストに分割し、ドライバのファイルシステムで利用できるようにする。その方法として、DBFSのルートボリュームやマウントポイントを利用する方法がある。

(train, test) = sms_id_labels.persist().randomSplit([0.8, 0.2])
# Write the tables to DBFS. 
train_dbfs_path = f"{tutorial_path}/sms_train"
test_dbfs_path = f"{tutorial_path}/sms_test"
train_df = train.write.parquet(train_dbfs_path, mode="overwrite")
test_df = test.write.parquet(test_dbfs_path, mode="overwrite")

これらのパーケットファイルは、マウントされた/dbfsパスを使用してドライバ上のファイルシステムで利用できるようになりました。トレーニングおよび評価データセットを作成するために、🤗 データセットユーティリティを使用して、パーケットファイルへのパスを指定することができます。

from datasets import load_dataset
train_test = load_dataset("parquet", data_files={"train":f"/dbfs{train_dbfs_path}/*.parquet", "test":f"/dbfs{test_dbfs_path}/*.parquet"})

このモデルは、ダウンロードしたデータのテキストではなく、トークン化された入力を想定しています。ベースモデルとの互換性を確保するために、ベースモデルからロードされたAutoTokenizerを使用します。HuggingFaceデータセットでは、トレーニングデータとテストデータの両方に一貫してトークナイザーを直接適用することができます。

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(base_model)
def tokenize_function(examples):
    return tokenizer(examples["text"], padding=False, truncation=True)

train_test_tokenized = train_test.map(tokenize_function, batched=True)

トレーナーの構築

Trainer のクラスでは、メトリクス、ベースモデル、トレーニングの設定をユーザが提供する必要があります。デフォルトでは、Trainerはメトリックとして損失を計算し使用しますが、これはあまり解釈しやすいものではありません。以下は、モデルのトレーニング中に精度を追加で計算するメトリクス関数を作成する例です。

import numpy as np
import evaluate
metric = evaluate.load("accuracy")
def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

テキスト分類の場合は、AutoModelForSequenceClassificationを使用して、テキスト分類用のベースモデルを読み込みます。ここでは、クラスの数とラベルのマッピングを指定します。

from transformers import AutoModelForSequenceClassification
model = AutoModelForSequenceClassification.from_pretrained(base_model, num_labels=2, label2id=label2id, id2label=id2label)

最後に、トレーニングの設定を作成する必要があります。TrainingArgumentsクラスでは、出力ディレクトリ、評価戦略、学習率、その他のparametersを指定することができます。

from transformers import TrainingArguments, Trainer
training_output_dir = "sms_trainer"
training_args = TrainingArguments(output_dir=training_output_dir, evaluation_strategy="epoch")

data collatorを使用すると、訓練と評価のデータセットの入力を一括して処理します。DataCollatorWithPaddingをデフォルトで使用すると、テキスト分類のベースライン性能が良好になります。

from transformers import DataCollatorWithPadding
data_collator = DataCollatorWithPadding(tokenizer)

これらのパラメータをすべて構築した上で、Trainerを作成することができます。

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_test_dataset["train"],
    eval_dataset=train_test_dataset["test"],
    compute_metrics=compute_metrics,
    data_collator=data_collator,
)

トレーニングの実行とモデルのロギング

Hugging FaceはMLflowとうまく連携し、MLflowCallback.を使用してモデルトレーニング中に自動的にメトリクスをログに記録します。しかし、訓練されたモデルを自分でログに記録する必要があります。上記の例と同様に、Databricksは学習済みモデルをトランスフォーマーパイプラインでラップし、MLflowのpyfunc log_model機能を使用することを推奨します。そのためには、カスタムモデルクラスが必要です。

import mlflow
import torch

pipeline_artifact_name = "pipeline"
class TextClassificationPipelineModel(mlflow.pyfunc.PythonModel):
  
  def load_context(self, context):
    device = 0 if torch.cuda.is_available() else -1
    self.pipeline = pipeline("text-classification", context.artifacts[pipeline_artifact_name], device=device)
    
  def predict(self, context, model_input): 
    texts = model_input[model_input.columns[0]].to_list()
    pipe = self.pipeline(texts, truncation=True, batch_size=8)
    labels = [prediction['label'] for prediction in pipe]
    return pd.Series(labels)

Tokenizerと学習済みモデルから変換パイプラインを構築し、ローカルディスクに書き込む。最後に、モデルをMLflowに記録します。

from transformers import pipeline

model_output_dir = "./sms_model"
pipeline_output_dir = "./sms_pipeline"
model_artifact_path = "sms_spam_model"

with mlflow.start_run() as run:
  trainer.train()
  trainer.save_model(model_output_dir)
  pipe = pipeline("text-classification", model=AutoModelForSequenceClassification.from_pretrained(model_output_dir), batch_size=8, tokenizer=tokenizer)
  pipe.save_pretrained(pipeline_output_dir)
  mlflow.pyfunc.log_model(artifacts={pipeline_artifact_name: pipeline_output_dir}, artifact_path=model_artifact_path, python_model=TextClassificationPipelineModel())

推論用モデルのロードは、MLflowでラップされた事前学習済みモデルのロードと同じです。

logged_model = "runs:/{run_id}/{model_artifact_path}".format(run_id=run.info.run_id, model_artifact_path=model_artifact_path)

# Load model as a Spark UDF. Override result_type if the model does not return double values.
sms_spam_model_udf = mlflow.pyfunc.spark_udf(spark, model_uri=logged_model, result_type='string')

test = test.select(test.text, test.label, sms_spam_model_udf(test.text).alias("prediction"))
display(test)

このセクションでは、Hugging Face Transformer Trainer API を直接使用して、新しいテキスト分類問題に対してモデルをFine-Tuningする方法を示しました。AutoModel classes for Natural Language Processingは素晴らしい基礎を提供します。.

まとめ

このブログ記事では、いくつかのベストプラクティスを示し、DatabricksでNLPタスクのためにトランスフォーマーを使い始めるのがいかに簡単であるかを示しました。

推論のために思い出すべき重要なポイントは以下の通りです：

• 🤗トランスフォーマーパイプラインは、トランスフォーマーモデルを簡単に使用することができます、
• フルクラスタを利用するために必要であれば、データを再分割します、
• GPUを効率的に使用するために、バッチサイズを調整することができます、
• Sparkは、マルチマシンGPUクラスターでGPUを自動的に割り当てます、
• Pandas UDFは、モデルのブロードキャストとデータのバッチングを管理し
• パイプラインは、MLflowへのトランスフォーマーモデルのロギングを簡素化します。

シングルマシンモデルのトレーニングで思い出すべきポイント：

• 🤗トランスフォーマー トレーナーは、モデルのFine-Tuningをするための身近な手段です
• Spark上でデータセットを準備し、モデリングタスクに必要な場合はラベルをidにマッピングし、トークン化はTransformersにお任せします🤗
• データセットをドライバーのファイルシステムで利用できるようにする、
• AutoTokenizerを使用してデータセットをトークン化し、モデルに適したトークナイザーをロードします、
• Fine-Tuningを行う場合は、Trainerを使用します、
• トークナイザーとFine-Tuningされたモデルからパイプラインを構築し
• パイプラインをラップするカスタムモデルを使用して、MLflowにログを記録します。

Databricksは、Databricks上でモデルトレーニングや推論をスケールさせるよりシンプルな方法への投資を続けています。データロード、分散モデルトレーニング、そしてMLflowモデルとしてのTransformersパイプラインとモデルの保存に関する改良にご期待ください。

これらの例で始めるには、事前訓練されたモデルの使用とfine-tuningのためのこれらのノートブックをインポートしてください。