機械翻訳モデルのファインチューニング¶

from datasets import load_dataset

# 英日翻訳用データセットを読み込み
split_datasets = load_dataset("Helsinki-NLP/opus-100", "en-ja")
split_datasets

# データセットのサンプルを確認
split_datasets["train"][1]["translation"]

from transformers import pipeline

# 英日翻訳用の事前訓練済みモデルを読み込み
model_checkpoint = "Helsinki-NLP/opus-mt-en-jap"
translator = pipeline("translation", model=model_checkpoint)
translator("I'm being held in a basement. I've been abducted with two other girls.")

from transformers import AutoTokenizer

# 英日翻訳モデルのTokenizerを初期化
model_checkpoint = "Helsinki-NLP/opus-mt-en-jap"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, return_tensors = "pt")

# 英語と日本語の文章を取得
en_sentence = split_datasets["train"][1]["translation"]["en"]
ja_sentence = split_datasets["train"][1]["translation"]["ja"]

# 入力文章とターゲット文章をトークン化
inputs = tokenizer(en_sentence, text_target=ja_sentence)
inputs

# 間違った方法：日本語文章を英語tokenizerで処理
wrong_targets = tokenizer(ja_sentence)
print(tokenizer.convert_ids_to_tokens(wrong_targets["input_ids"]))
print(tokenizer.convert_ids_to_tokens(inputs["labels"]))

# 最大長を設定
max_length = 128

def preprocess_function(examples):
    # 英語と日本語の文章を分離
    inputs = [ex["en"] for ex in examples["translation"]]
    targets = [ex["ja"] for ex in examples["translation"]]

    # 入力とターゲットを同時にトークン化
    model_inputs = tokenizer(
        inputs, text_target=targets, max_length=max_length, truncation=True
    )
    return model_inputs

# データセット全体に前処理を適用
tokenized_datasets = split_datasets.map(
    preprocess_function,
    batched=True,
    remove_columns=split_datasets["train"].column_names
)

from transformers import AutoModelForSeq2SeqLM
import torch

# Sequence-to-Sequence 学習用モデルを読み込み
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)

# デバイスの設定（MPS, GPU, CPUの順で利用可能なものを使用）
if torch.backends.mps.is_available():
    device = torch.device("mps")
else:
    device = torch.device("cpu")

model = model.to(device)

from transformers import DataCollatorForSeq2Seq

# Seq2Seq用のデータコレクターを初期化
data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)

# データコレクターのテスト実行
batch = data_collator([tokenized_datasets["train"][i] for i in range(1, 3)])

# パディングされたラベルを確認
batch["labels"]

# デコーダー入力IDを確認
batch["decoder_input_ids"]

# 元のラベルを確認
for i in range(1, 3):
    print(tokenized_datasets["train"][i]["labels"])

# !pip install sacrebleu

import evaluate

# SacreBLEUメトリックを読み込み
metric = evaluate.load("sacrebleu")

# BLEUスコアの計算例
predictions = [
    "This plugin lets you translate web pages between several languages automatically."
]
references = [
    [
        "This plugin allows you to automatically translate web pages between several languages."
    ]
]
metric.compute(predictions=predictions, references=references)

# 悪い予測例1：繰り返しが多い場合
predictions = ["This This This This"]
references = [
    [
        "This plugin allows you to automatically translate web pages between several languages."
    ]
]
metric.compute(predictions=predictions, references=references)

# 悪い予測例2：短すぎる場合
predictions = ["This plugin"]
references = [
    [
        "This plugin allows you to automatically translate web pages between several languages."
    ]
]
metric.compute(predictions=predictions, references=references)

import numpy as np

def compute_metrics(eval_preds):
    preds, labels = eval_preds
    # モデルが予測logit以上を返す場合
    if isinstance(preds, tuple):
        preds = preds[0]

    # 予測をデコードしてテキストに変換
    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)

    # ラベルの-100をパディングトークンIDに置き換えてデコード
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

    # 簡単な後処理：空白の削除
    decoded_preds = [pred.strip() for pred in decoded_preds]
    decoded_labels = [[label.strip()] for label in decoded_labels]

    # BLEUスコアを計算
    result = metric.compute(predictions=decoded_preds, references=decoded_labels)
    return {"bleu": result["score"]}

from transformers import Seq2SeqTrainingArguments

# 訓練パラメータの設定
args = Seq2SeqTrainingArguments(
    f"marian-finetuned-opus-100-en-to-ja",  # 出力ディレクトリ
    eval_strategy="no",  # 評価戦略
    save_strategy="epoch",  # 保存戦略
    learning_rate=2e-5,  # 学習率
    per_device_train_batch_size=16,  # 訓練時のバッチサイズ
    per_device_eval_batch_size=32,   # 評価時のバッチサイズ
    weight_decay=0.01,  # 重み減衰
    save_total_limit=3,  # 保存するチェックポイントの最大数
    num_train_epochs=3,  # エポック数
    predict_with_generate=True,  # 評価時にgenerate()メソッドを使用
    push_to_hub=False,  # Hubにアップロードしない
)

from transformers import Seq2SeqTrainer

# Seq2SeqTrainerの初期化
trainer = Seq2SeqTrainer(
    model,
    args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    processing_class=tokenizer,
    compute_metrics=compute_metrics,
)

# 実際の訓練実行（コメントアウト）
# trainer.train()

from torch.utils.data import DataLoader

# データセットをPyTorch形式に設定
tokenized_datasets.set_format("torch")

# 訓練用データローダーの作成
train_dataloader = DataLoader(
    tokenized_datasets["train"],
    shuffle=True,  # データをシャッフル
    collate_fn=data_collator,
    batch_size=8,
)

# 評価用データローダーの作成
eval_dataloader = DataLoader(
    tokenized_datasets["validation"], 
    collate_fn=data_collator, 
    batch_size=8
)

# モデルを再度読み込み（前回の訓練状態をリセット）
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)

from torch.optim import AdamW

# Adam optimizerの初期化
optimizer = AdamW(model.parameters(), lr=2e-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,
)

# 出力ディレクトリの設定
output_dir = "marian-finetuned-kde4-en-to-fr-accelerate"

def postprocess(predictions, labels):
    # CPUにデータを移動してnumpy配列に変換
    predictions = predictions.cpu().numpy()
    labels = labels.cpu().numpy()

    # 予測をテキストにデコード
    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)

    # ラベルの-100をパディングトークンIDに置き換えてデコード
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

    # 簡単な後処理：空白の削除とリスト形式への変換
    decoded_preds = [pred.strip() for pred in decoded_preds]
    decoded_labels = [[label.strip()] for label in decoded_labels]
    return decoded_preds, decoded_labels

from tqdm.auto import tqdm
import torch

# 進捗バーの設定
progress_bar = tqdm(range(num_training_steps))

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()
    for batch in tqdm(eval_dataloader):
        with torch.no_grad():
            # generate()メソッドを使用してテキスト生成
            generated_tokens = accelerator.unwrap_model(model).generate(
                batch["input_ids"],
                attention_mask=batch["attention_mask"],
                max_length=128,
            )
        labels = batch["labels"]

        # 予測とラベルを各プロセス間で同じ形状にパディング
        generated_tokens = accelerator.pad_across_processes(
            generated_tokens, dim=1, pad_index=tokenizer.pad_token_id
        )
        labels = accelerator.pad_across_processes(labels, dim=1, pad_index=-100)

        # 全プロセスから予測とラベルを収集
        predictions_gathered = accelerator.gather(generated_tokens)
        labels_gathered = accelerator.gather(labels)

        # テキストに変換してメトリックに追加
        decoded_preds, decoded_labels = postprocess(predictions_gathered, labels_gathered)
        metric.add_batch(predictions=decoded_preds, references=decoded_labels)

    # BLEUスコアを計算して表示
    results = metric.compute()
    print(f"epoch {epoch}, BLEU score: {results['score']:.2f}")

    # モデルの保存とアップロード
    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)