HuggingFaceを使ってCausal Language ModelのFine-tuning

Causal Language Model（CLM）は、テキストを生成するために使われるモデルの一種で、言語の生成における因果関係（causality）を重視します。主に次に来る単語を予測するタスクに使われ、一般的に次のような用途に向いています：

テキスト生成：次に来る単語を予測し、シーケンス全体を生成。
チャットボット：ユーザーの入力に対し自然な応答を生成。
オートコンプリート：入力されたフレーズの続きを予測。

Causal Language Modelの特徴

次単語予測（Autoregressive）
CLMは前の単語やトークンの情報のみを使って次の単語を予測します。このように、モデルは各ステップで過去の情報に基づき次の単語を生成しますが、未来の情報（これから出現する単語）は参照しません。
トレーニング方法
トレーニングには「左から右」への順序でシーケンスを入力し、それぞれのトークンに次のトークンを予測させます。例えば、文「I like cats」なら、「I」の次に「like」、「like」の次に「cats」を予測するように学習します。

代表的なモデル

Hugging Face Transformersで提供されているGPT系のモデル（GPT-2、GPT-3など）はCausal Language Modelingの代表例です。

HuggingFaceを使ってCLMをFine-tuningする

ここでは、HuggingFaceのCausal Language Modelingで公開されているコードを実行しながら、FTのステップについてまとめていきます。

huggingface.co

大きくは、以下の4つのステップがあります。

データセットのロード
データの前処理
学習
推論

1. データセットのロード

ELI5というデータセットを使います。ELI5は、以下のようなデータセットです（https://facebookresearch.github.io/ELI5/より）

ELI5は長文形式の質問応答のためのデータセットです。約27万件の複雑で多様な質問が含まれており、複数の文での説明的な回答が求められます。各質問に答えるために、ウェブ検索の結果が根拠となる文書として使用されています。

27万件のデータセットは規模が大きいので、サブセットとして5000件を使います。

from datasets import load_dataset
eli5 = load_dataset("eli5_category", split="train[:5000]")

ちなみに、HuggingFaceのページでは、

eli5 = load_dataset("eli5", split="train_asks[:5000]")

となっていますが、これだとDefunctDatasetError: Dataset 'eli5' is defunct and no longer accessible due to unavailability of the source dataとなってしまいます。英語版のMainブランチでは修正されていますが、日本語版ではまだのようですので注意（関連PR）。

2. データの前処理

前処理としてまず実施するのは、トークン化です。

DistilGPT2トークナイザーをロードして、eli5データセットのanswers.textフィールドに含まれるテキストデータをトークン化します。データセット全体に前処理を適用する際には、mapを使って、複数のデータを一度に処理しています。

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2")

def preprocess_function(examples):
  return tokenizer([" ".join(x) for x in examples["answers.text"]])

tokenized_eli5 = eli5.map(
  preprocess_function,
  batched=True,
  num_proc=4,
  remove_columns=eli5["train"].column_names,
)

これを実行すると、このようなWarningが出力されます。

Token indices sequence length is longer than the specified maximum sequence length for this model (1033 > 1024). Running this sequence through the model will result in indexing errors

そのため、2つ目の前処理として、モデルの最大入力長（ここでは1024）を超えないように、短いチャンクに分割します。

block_size = 128

def group_texts(examples):
  # Concatenate all texts.
  concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
  total_length = len(concatenated_examples[list(examples.keys())[0]])
  # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
  # customize this part to your needs.
  if total_length >= block_size:
    total_length = (total_length // block_size) * block_size
  # Split by chunks of block_size.
  result = {
    k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
    for k, t in concatenated_examples.items()
  }
  result["labels"] = result["input_ids"].copy()
  return result

これをデータセット全体に適用します。

lm_dataset = tokenized_eli5.map(group_texts, batched=True, num_proc=4)

最後に、バッチ生成用のDataCollatorを準備します。DataCollatorがバッチ内のデータでPadding（Dynaminc padding）を行なってくれて、その時に使用するトークンはeos_tokenとしています。

from transformers import DataCollatorForLanguageModeling

tokenizer.pad_token = tokenizer.eos_token
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

3. 学習

モデルの学習のための準備ができたので、実際に学習を行います。ハイパーパラメータの設定、Trainerの作成を行い、学習を開始します。HuggingFace Hubへのモデルアップロードと、WandBを使用しないので、push_to_hub=Falseとreport_to=noneを追加で指定します。

training_args = TrainingArguments(
  output_dir="my_awesome_eli5_clm-model",
  eval_strategy="epoch",
  learning_rate=2e-5,
  weight_decay=0.01,
  push_to_hub=False,
  report_to="none",
)

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

trainer.train()

ColabのA100を使って、おおよそ5分くらいの学習時間でした。

Perplexityを確認すると46.02で、HuggingFaceの値（49.61）から少し乖離がある結果となりました。

import math

eval_results = trainer.evaluate()
print(f"Perplexity: {math.exp(eval_results['eval_loss']):.2f}")

Inference

FTしたモデルを使って推論を試します。GPUを使える環境なため、device=0を追加で指定しています。

prompt = "Somatic hypermutation allows the immune system to"

from transformers import pipeline

generator = pipeline("text-generation", model="my_model", device=0)
generator(prompt)

生成されたテキストは、

[{'generated_text': "Somatic hypermutation allows the immune system to work with a different gene and cause the cell to fail, making it more susceptible to damage to other cells and even becoming a parasympathetic to other cells. When we don't know if we"}]