Hugging Face Trainerで小さいモデルをファインチューニング

1. 概要

このガイドでは、transformersライブラリのTrainer APIを使用して、小さいモデル(DistilBERT)をファインチューニングし、商品レビューの感情分類を行う方法を説明します。このガイドを読み終えると、データセットの準備方法、tokenize、TrainingArgumentsの設定、モデルの訓練、評価、推論用のモデル保存方法がわかるようになります。

本ガイドを通じて使用する例:OrderReview – 顧客(Customer)による注文(Order)のレビュー。positive/negativeのラベルが付与されています。

2. 環境の準備

必要なライブラリをインストールします。強力なGPUは不要です。DistilBERTのような小さいモデルは小規模なデータセットであればCPUでも十分に動作します(ただし処理時間はかかります)。

pip install transformers datasets evaluate accelerate torch scikit-learn

注意: マシンにGPUがある場合は、対応するCUDA版のtorchをインストールすることで、訓練速度が大幅に向上します。CPUの場合は、デモンストレーションのために、データ行数とエポック数を制限して実行時間を削減することをお勧めします。

3. データセットの準備

巨大なデータセットを読み込む代わりに、注文レビューをシミュレートする小規模なセットを作成します。理解しやすくするためです。実際には、load_dataset("imdb")や自身のCSVデータに置き換えることができます。

from datasets import Dataset

reviews = [
    {"text": "注文が早く届き、梱包が丁寧で、とても満足しています", "label": 1},
    {"text": "商品が壊れていて、配送が遅く、説明と異なります", "label": 0},
    {"text": "品質が良く、価格が適切で、また購入します", "label": 1},
    {"text": "サービスが悪く、カスタマーサポートが遅い", "label": 0},
    {"text": "この注文にとても満足しています、5つ星", "label": 1},
    {"text": "商品が配送中に傷んでいました、がっかりです", "label": 0},
]

dataset = Dataset.from_list(reviews)
dataset = dataset.train_test_split(test_size=0.33, seed=42)
print(dataset)
DatasetDict({
    train: Dataset({features: ['text', 'label'], num_rows: 4})
    test: Dataset({features: ['text', 'label'], num_rows: 2})
})

Dataset.from_listは、辞書のリストをHugging FaceのDatasetに変換します。train_test_splitは訓練セットとテストセットに分割します。実データセットの場合は数千行以上あるべきですが、これは処理フローを示すための例です。

4. データのトークン化

モデルはテキストを直接理解できません。tokenizer はテキストをinput_idsとattention_maskに変換する必要があります。

from transformers import AutoTokenizer

checkpoint = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

def tokenize_fn(batch):
    return tokenizer(batch["text"], truncation=True, padding="max_length", max_length=32)

tokenized_dataset = dataset.map(tokenize_fn, batched=True)
print(tokenized_dataset["train"][0].keys())
dict_keys(['text', 'label', 'input_ids', 'attention_mask'])
  • truncation=True: max_lengthより長い文を切り詰めます。
  • padding="max_length": すべての文を同じ長さにパディングします。これによってbatchの並列処理が容易になります。
  • map(..., batched=True): tokenize関数をbatch単位で適用し、1行ずつ処理するより高速です。

注意: 実際には、英語のDistilBERTはベトナム語のレビューのファインチューニングには適していません。ここはテクニックのデモンストレーションです。ベトナム語データの場合は、vinai/phobert-baseなどのcheckpointを選択してください。

5. モデルの読み込みとラベルの定義

from transformers import AutoModelForSequenceClassification

label2id = {"negative": 0, "positive": 1}
id2label = {v: k for k, v in label2id.items()}

model = AutoModelForSequenceClassification.from_pretrained(
    checkpoint,
    num_labels=2,
    id2label=id2label,
    label2id=label2id,
)

num_labels=2は、事前訓練されたencoderの上に2つの出力を持つ新しい分類ヘッドを作成します(positive/negative)。id2label/label2idは、モデルの出力が0/1の数値ではなく、人間が読める形のラベルを返すようにします。

6. TrainingArgumentsの設定

TrainingArgumentsは、訓練ループのすべてのハイパーパラメータと設定をまとめます。

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./order-review-model",
    eval_strategy="epoch",
    save_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    num_train_epochs=3,
    weight_decay=0.01,
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    logging_steps=1,
)
  • eval_strategy="epoch": 各ステップではなく各エポック後に評価を行います。
  • load_best_model_at_end=True: 訓練後、Trainerはmetric_for_best_modelに基づいて最良のcheckpointを自動的に読み込みます。
  • per_device_train_batch_size=4: デモデータセットが小さいため、batchサイズを小さくしています。実データの場合は、VRAMに応じて16/32に増やしてください。

7. メトリクス計算関数の定義

Trainerがモデルをどのような基準で評価するかを知る必要があります。

import numpy as np
import evaluate

accuracy = evaluate.load("accuracy")
f1 = evaluate.load("f1")

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    acc = accuracy.compute(predictions=predictions, references=labels)
    f1_score = f1.compute(predictions=predictions, references=labels)
    return {"accuracy": acc["accuracy"], "f1": f1_score["f1"]}

eval_predはTrainerが各評価時に渡すタプル(logits, labels)です。np.argmaxは最も高いlogitを持つラベルを予測として選択します。

8. Trainerの初期化と訓練

from transformers import Trainer, DataCollatorWithPadding

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

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

trainer.train()
Epoch   Training Loss   Validation Loss   Accuracy   F1
1       0.693100        0.671204          0.500000   0.500000
2       0.654200        0.598311          1.000000   1.000000
3       0.581900        0.520447          1.000000   1.000000

DataCollatorWithPaddingはbatch単位で動的にパディングを処理します(tokenize時の固定max_lengthパディングより効率的ですが、ここはすでに事前パディングされているため、主にbatchテンソル作成を行います)。trainer.train()は forward → backward → optimizer step → ロス記録のループを実行し、設定されたeval_strategyに従って自動的に評価を行います。

8.1. 訓練進度の監視

TensorBoardへのより詳細なロギングが必要な場合は、TrainingArgumentsにreport_to="tensorboard"を追加して、tensorboard --logdir ./order-review-model/runsを実行します。デモのような小規模データセットでは、コンソールログでロスが段階的に減少しているか確認するだけで十分です。

9. 評価とモデルの保存

metrics = trainer.evaluate()
print(metrics)

trainer.save_model("./order-review-model/final")
tokenizer.save_pretrained("./order-review-model/final")
{'eval_loss': 0.5204, 'eval_accuracy': 1.0, 'eval_f1': 1.0, 'eval_runtime': 0.12, ...}

trainer.evaluate()はeval_dataset上で再度実行し、メトリクスの辞書を返します。save_modelはconfigと重みの両方を保存します。後で読み込む際に語彙がズレないようにtokenizer もsave_pretrainedで一緒に保存することを忘れずに。

10. ファインチューニング済みモデルでのInference

pipelineを使用して、訓練したモデルをすばやくテストします。

from transformers import pipeline

classifier = pipeline(
    "text-classification",
    model="./order-review-model/final",
    tokenizer="./order-review-model/final",
)

result = classifier("配送が予定通りで、商品は期待以上に良かった")
print(result)
[{'label': 'positive', 'score': 0.9123}]

pipelineはtokenize、forward、softmaxを自動的に処理し、labelと信頼度スコアを返します。これは小規模なサービスにモデルを統合する最速の方法です。例えば、PaymentServiceが顧客フィードバックを自動分類する必要がある場合などです。

11. モデルの動作をチェックするテストの作成

デモデータセットでは、pytestを使用して簡単なテストを作成し、pipelineが正しく動作し、checkpointを変更する際に回帰が起きないことを確認できます。

import pytest
from transformers import pipeline

@pytest.fixture(scope="module")
def classifier():
    return pipeline(
        "text-classification",
        model="./order-review-model/final",
        tokenizer="./order-review-model/final",
    )

def test_positive_review_is_classified_correctly(classifier):
    result = classifier("素晴らしい商品、配送が早い")[0]

    assert result["label"] == "positive"
    assert result["score"] > 0.5

def test_negative_review_is_classified_correctly(classifier):
    result = classifier("商品が壊れていて、配送が遅く、とても失望した")[0]

    assert result["label"] == "negative"
    assert result["score"] > 0.5

scope="module"により、pipelineはファイル内のすべてのテストに対して1回だけロードされ、毎回ロードするのに費やす時間を節約できます。このテストは高レベルの動作(ラベルの正確性)のみをチェックします。モデル品質の評価は、上記の評価ステップに従うべきです。

12. よくあるエラー

  • padding/truncationを忘れる: max_lengthより長い文はbatch内で異なるテンソルサイズになるため、実行時エラーが発生します。
  • num_labelsを間違える: データセットに3つのクラスがあるのにnum_labels=2と宣言すると、Trainerは実行されますが結果が完全に間違います。分類ヘッドの出力が不十分になるからです。
  • 非常に小さいデータセットで高いF1を期待する: このガイドの例はコードフローを示すためだけのものです。訓練データが4行では、モデルが汎化することはありません。ここのaccuracy=1.0を基準にしないでください。
  • seedを設定しない: 訓練実行ごとに異なる結果が得られ、ハイパーパラメータ調整時の比較が難しくなります。訓練前にtransformersのset_seed(42)を使用してください。
  • 弱いマシンで大規模モデル(base/large)を使用する: distilbert/tinyバージョンから始めて、開発ループを高速化してから、パイプラインを確定した後に大規模モデルに切り替えてください。

Trainerを使わない場合: 訓練ループを完全にカスタマイズする必要がある場合(例:複雑なマルチタスク損失、Trainerで表現困難なカスタムscheduler)、PyTorchの純粋なループまたはaccelerateを直接使用する方が柔軟です。

13. 結論

Trainer APIを使用すると、訓練/評価ループを自分で書く必要なく、小さいモデルをファインチューニングできます。データセット、tokenizer、モデル、TrainingArgumentsを準備するだけです。注文レビュー分類のような単純な分類タスクの場合、標準フローは:トークン化 → TrainingArguments設定 → compute_metrics定義 → Trainer初期化 → train() → evaluate() → save_model()です。実データがより複雑な場合でも、データセットとcheckpointを置き換えるだけで、パイプラインの他の部分はほぼ変わりません。

huggingfacetransformersfine-tuningnlppython