マスク言語モデルのファインチューニング¶

from transformers import AutoModelForMaskedLM

# DistilBERTの事前学習済みモデルを読み込み
model_checkpoint = "distilbert-base-uncased"
model = AutoModelForMaskedLM.from_pretrained(model_checkpoint)

# モデルのパラメータ数を計算（百万単位）
distilbert_num_parameters = model.num_parameters() / 1_000_000
print(f"'>>> DistilBERT パラメータ数: {round(distilbert_num_parameters)}M'")
print(f"'>>> BERT パラメータ数: 110M'")

# テスト用のサンプルテキスト
text = "This is a great [MASK]."

from transformers import AutoTokenizer

# 対応するトークナイザーを読み込み
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

import torch

# テキストをトークン化してモデルに入力
inputs = tokenizer(text, return_tensors="pt")
token_logits = model(**inputs).logits

# [MASK]トークンの位置を特定し、そのロジットを抽出
mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]
mask_token_logits = token_logits[0, mask_token_index, :]

# 最も高いロジットを持つ上位5つの[MASK]候補を取得
top_5_tokens = torch.topk(mask_token_logits, 5, dim=1).indices[0].tolist()

# 結果を表示
for token in top_5_tokens:
    print(f"'>>> {text.replace(tokenizer.mask_token, tokenizer.decode([token]))}'")

from datasets import load_dataset

# IMDbデータセットを読み込み
imdb_dataset = load_dataset("imdb")
imdb_dataset

# ランダムに3つのサンプルを選択
sample = imdb_dataset["train"].shuffle(seed=42).select(range(3))

for row in sample:
    print(f"\n'>>> レビュー: {row['text']}'")
    print(f"'>>> ラベル: {row['label']}'")

def tokenize_function(examples):
    # テキストをトークン化（切り捨てなし）
    result = tokenizer(examples["text"])
    # 高速トークナイザーの場合、word IDsを取得
    if tokenizer.is_fast:
        result["word_ids"] = [result.word_ids(i) for i in range(len(result["input_ids"]))]
    return result

# 高速マルチスレッド処理を有効化
tokenized_datasets = imdb_dataset.map(
    tokenize_function, batched=True, remove_columns=["text", "label"]
)
tokenized_datasets

# モデルの最大コンテキストサイズを確認
tokenizer.model_max_length

# スライシングにより各特徴のリストのリストを生成
tokenized_samples = tokenized_datasets["train"][:3]

for idx, sample in enumerate(tokenized_samples["input_ids"]):
    print(f"'>>> レビュー {idx} の長さ: {len(sample)}'")

# すべての例を連結
concatenated_examples = {
    k: sum(tokenized_samples[k], []) for k in tokenized_samples.keys()
}
total_length = len(concatenated_examples["input_ids"])
print(f"'>>> 連結されたレビューの長さ: {total_length}'")

# チャンクサイズを設定
chunk_size = 128

# 連結された例の特徴ごとにチャンクを作成
chunks = {
    k: [t[i : i + chunk_size] for i in range(0, total_length, chunk_size)]
    for k, t in concatenated_examples.items()
}

# 各チャンクの長さを確認
for chunk in chunks["input_ids"]:
    print(f"'>>> チャンク長: {len(chunk)}'")

def group_texts(examples):
    # すべてのテキストを連結
    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
    # 連結されたテキストの長さを計算
    total_length = len(concatenated_examples[list(examples.keys())[0]])
    # 最後のチャンクがchunk_sizeより小さい場合は削除
    total_length = (total_length // chunk_size) * chunk_size
    # max_lenのチャンクに分割
    result = {
        k: [t[i : i + chunk_size] for i in range(0, total_length, chunk_size)]
        for k, t in concatenated_examples.items()
    }
    # 新しいlabelsカラムを作成
    result["labels"] = result["input_ids"].copy()
    return result

# グループ化されたテキストに変換
lm_datasets = tokenized_datasets.map(group_texts, batched=True)
lm_datasets

# チャンクの内容を確認
tokenizer.decode(lm_datasets["train"][1]["input_ids"])

# ラベルの内容を確認
tokenizer.decode(lm_datasets["train"][1]["labels"])

from transformers import DataCollatorForLanguageModeling

# マスク言語モデリング用のデータコレーターを作成
# 15%の確率でトークンをマスク
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15)

# サンプルデータを準備
samples = [lm_datasets["train"][i] for i in range(2)]
for sample in samples:
    _ = sample.pop("word_ids")

# データコレーターを適用してマスキング結果を確認
for chunk in data_collator(samples)["input_ids"]:
    print(f"\n'>>> {tokenizer.decode(chunk)}'")

import collections
import numpy as np
from transformers import default_data_collator

# 全単語マスキングの確率
wwm_probability = 0.2

def whole_word_masking_data_collator(features):
    for feature in features:
        word_ids = feature.pop("word_ids")

        # 単語と対応するトークンインデックス間のマップを作成
        mapping = collections.defaultdict(list)
        current_word_index = -1
        current_word = None
        for idx, word_id in enumerate(word_ids):
            if word_id is not None:
                if word_id != current_word:
                    current_word = word_id
                    current_word_index += 1
                mapping[current_word_index].append(idx)

        # 単語をランダムにマスク
        mask = np.random.binomial(1, wwm_probability, (len(mapping),))
        input_ids = feature["input_ids"]
        labels = feature["labels"]
        new_labels = [-100] * len(labels)
        for word_id in np.where(mask)[0]:
            word_id = word_id.item()
            for idx in mapping[word_id]:
                new_labels[idx] = labels[idx]
                input_ids[idx] = tokenizer.mask_token_id
        feature["labels"] = new_labels

    return default_data_collator(features)

# 全単語マスキングの動作確認
samples = [lm_datasets["train"][i] for i in range(2)]
batch = whole_word_masking_data_collator(samples)

for chunk in batch["input_ids"]:
    print(f"\n'>>> {tokenizer.decode(chunk)}'")

# データセットサイズを設定
train_size = 10_000
test_size = int(0.1 * train_size)

# ダウンサンプリングを実行
downsampled_dataset = lm_datasets["train"].train_test_split(
    train_size=train_size, test_size=test_size, seed=42
)
downsampled_dataset

from transformers import TrainingArguments

batch_size = 64
# エポックごとに訓練損失を表示
logging_steps = len(downsampled_dataset["train"]) // batch_size
model_name = model_checkpoint.split("/")[-1]

training_args = TrainingArguments(
    output_dir=f"{model_name}-finetuned-imdb",
    overwrite_output_dir=True,
    eval_strategy="epoch",
    learning_rate=2e-5,
    weight_decay=0.01,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    push_to_hub=False,
    logging_steps=logging_steps,
)

from transformers import Trainer

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

import math

# 事前学習済みモデルの困惑度を評価
# eval_results = trainer.evaluate()
# print(f">>> 困惑度: {math.exp(eval_results['eval_loss']):.2f}")
# >>> 困惑度: 21.75

# 訓練の実行
# trainer.train()

# ファインチューニング後の困惑度を評価
# eval_results = trainer.evaluate()
# print(f">>> 困惑度: {math.exp(eval_results['eval_loss']):.2f}")
# >>> 困惑度: 11.32

def insert_random_mask(batch):
    features = [dict(zip(batch, t)) for t in zip(*batch.values())]
    masked_inputs = data_collator(features)
    # データセットの各カラムに新しい「masked」カラムを作成
    return {"masked_" + k: v.numpy() for k, v in masked_inputs.items()}

# データセットからword_idsカラムを削除
downsampled_dataset = downsampled_dataset.remove_columns(["word_ids"])

# テストセットに固定マスキングを適用
eval_dataset = downsampled_dataset["test"].map(
    insert_random_mask,
    batched=True,
    remove_columns=downsampled_dataset["test"].column_names,
)

# カラム名を変更してマスクされたバージョンを標準名に
eval_dataset = eval_dataset.rename_columns(
    {
        "masked_input_ids": "input_ids",
        "masked_attention_mask": "attention_mask",
        "masked_labels": "labels",
    }
)

from torch.utils.data import DataLoader
from transformers import default_data_collator

batch_size = 64

# 訓練用データローダー（ランダムマスキング付き）
train_dataloader = DataLoader(
    downsampled_dataset["train"],
    shuffle=True,
    batch_size=batch_size,
    collate_fn=data_collator,
)

# 評価用データローダー（固定マスキング）
eval_dataloader = DataLoader(
    eval_dataset, batch_size=batch_size, collate_fn=default_data_collator
)

# 新しいモデルインスタンスを読み込み
model = AutoModelForMaskedLM.from_pretrained(model_checkpoint)

from torch.optim import AdamW

# オプティマイザーを設定
optimizer = AdamW(model.parameters(), lr=5e-5)

from accelerate import Accelerator

# Acceleratorを初期化し、すべてのコンポーネントを準備
accelerator = Accelerator()
model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader, eval_dataloader
)

from transformers import get_scheduler

# 訓練パラメータの設定
num_train_epochs = 3
num_update_steps_per_epoch = len(train_dataloader)
num_training_steps = num_train_epochs * num_update_steps_per_epoch

# 線形学習率スケジューラーを設定
lr_scheduler = get_scheduler(
    "linear",
    optimizer=optimizer,
    num_warmup_steps=0,
    num_training_steps=num_training_steps,
)

from tqdm.auto import tqdm
import torch
import math

# プログレスバーと出力ディレクトリを設定
progress_bar = tqdm(range(num_training_steps))
output_dir = "distilbert-base-uncased-finetuned-imdb-accelerate"

for epoch in range(num_train_epochs):
    # 訓練フェーズ
    model.train()
    for batch in train_dataloader:
        # 順伝播
        outputs = model(**batch)
        loss = outputs.loss
        # 逆伝播
        accelerator.backward(loss)

        # パラメータ更新
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        progress_bar.update(1)

    # 評価フェーズ
    model.eval()
    losses = []
    for step, batch in enumerate(eval_dataloader):
        with torch.no_grad():
            outputs = model(**batch)

        loss = outputs.loss
        losses.append(accelerator.gather(loss.repeat(batch_size)))

    # 困惑度の計算
    losses = torch.cat(losses)
    losses = losses[: len(eval_dataset)]
    try:
        perplexity = math.exp(torch.mean(losses))
    except OverflowError:
        perplexity = float("inf")

    print(f">>> エポック {epoch}: 困惑度: {perplexity}")

    # モデルの保存とアップロード
    accelerator.wait_for_everyone()
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)
    if accelerator.is_main_process:
        tokenizer.save_pretrained(output_dir)

from transformers import pipeline

# ファインチューニング済みモデルでパイプラインを作成
mask_filler = pipeline(
    "fill-mask", model="distilbert-base-uncased-finetuned-imdb-accelerate"
)

# マスク予測を実行
preds = mask_filler(text)

for pred in preds:
    print(f">>> {pred['sequence']}")