Transformersを使ったトークン分類：固有表現認識（NER）の実践ガイド¶

from datasets import load_dataset

# CoNLL-2003データセットを読み込み
raw_datasets = load_dataset("eriktks/conll2003", revision="convert/parquet")

# 最初のサンプルのトークンを表示
raw_datasets["train"][0]["tokens"]

# 対応するNERタグを表示
raw_datasets["train"][0]["ner_tags"]

# NERタグの特徴量情報を取得
ner_feature = raw_datasets["train"].features["ner_tags"]
print(ner_feature)

# ラベル名のリストを取得
label_names = ner_feature.feature.names
print(label_names)

# ラベルとトークンを整列表示する関数
words = raw_datasets["train"][0]["tokens"]
labels = raw_datasets["train"][0]["ner_tags"]
line1 = ""
line2 = ""
for word, label in zip(words, labels):
    full_label = label_names[label]
    max_length = max(len(word), len(full_label))
    line1 += word + " " * (max_length - len(word) + 1)
    line2 += full_label + " " * (max_length - len(full_label) + 1)

print(line1)
print(line2)

# インデックス4の例を表示
words = raw_datasets["train"][4]["tokens"]
labels = raw_datasets["train"][4]["ner_tags"]
line1 = ""
line2 = ""
for word, label in zip(words, labels):
    full_label = label_names[label]
    max_length = max(len(word), len(full_label))
    line1 += word + " " * (max_length - len(word) + 1)
    line2 += full_label + " " * (max_length - len(full_label) + 1)

print(line1)
print(line2)

# POSタグのラベル名を確認
pos_label_names = raw_datasets["train"].features["pos_tags"].feature.names
print(pos_label_names)

# 最初のサンプルのPOSタグを表示
words = raw_datasets["train"][0]["tokens"]
pos_labels = raw_datasets["train"][0]["pos_tags"]
line1 = ""
line2 = ""
for word, label in zip(words, pos_labels):
    full_label = pos_label_names[label]
    max_length = max(len(word), len(full_label))
    line1 += word + " " * (max_length - len(word) + 1)
    line2 += full_label + " " * (max_length - len(full_label) + 1)

print(line1)
print(line2)

# チャンクタグのラベル名を確認
chunk_label_names = raw_datasets["train"].features["chunk_tags"].feature.names
print(chunk_label_names)

# 最初のサンプルのチャンクタグを表示
words = raw_datasets["train"][0]["tokens"]
chunk_labels = raw_datasets["train"][0]["chunk_tags"]
line1 = ""
line2 = ""
for word, label in zip(words, chunk_labels):
    full_label = chunk_label_names[label]
    max_length = max(len(word), len(full_label))
    line1 += word + " " * (max_length - len(word) + 1)
    line2 += full_label + " " * (max_length - len(full_label) + 1)

print(line1)
print(line2)

from transformers import AutoTokenizer

# 使用するモデルのチェックポイントを指定
model_checkpoint = "bert-base-cased"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

# Fasttokenizerかどうかを確認
print(tokenizer.is_fast)

# 事前トークン化された入力をトークン化
inputs = tokenizer(raw_datasets["train"][0]["tokens"], is_split_into_words=True)
print(inputs.tokens())

# 各トークンの対応する単語IDを確認
print(inputs.word_ids())

def align_labels_with_tokens(labels, word_ids):
    """ラベルをトークンに整列させる関数"""
    new_labels = []
    current_word = None
    for word_id in word_ids:
        if word_id != current_word:
            # 新しい単語の開始
            current_word = word_id
            label = -100 if word_id is None else labels[word_id]
            new_labels.append(label)
        elif word_id is None:
            # 特殊トークン
            new_labels.append(-100)
        else:
            # 前のトークンと同じ単語
            label = labels[word_id]
            # ラベルがB-XXXの場合、I-XXXに変更
            if label % 2 == 1:
                label += 1
            new_labels.append(label)

    return new_labels

# ラベル整列のテスト
labels = raw_datasets["train"][0]["ner_tags"]
word_ids = inputs.word_ids()
print("元のラベル:", labels)
print("整列後ラベル:", align_labels_with_tokens(labels, word_ids))

def align_labels_with_tokens_v2(labels, word_ids):
    """単語ごとに1つのラベルのみを割り当てる代替バージョン"""
    new_labels = []
    current_word = None
    for word_id in word_ids:
        if word_id != current_word:
            # 新しい単語の開始
            current_word = word_id
            label = -100 if word_id is None else labels[word_id]
            new_labels.append(label)
        elif word_id is None:
            # 特殊トークン
            new_labels.append(-100)
        else:
            # 前のトークンと同じ単語（サブトークン）
            new_labels.append(-100)

    return new_labels

# 代替バージョンのテスト
labels = raw_datasets["train"][0]["ner_tags"]
word_ids = inputs.word_ids()
print("元のラベル:", labels)
print("代替整列:", align_labels_with_tokens_v2(labels, word_ids))

def tokenize_and_align_labels(examples):
    """データセット全体をトークン化してラベルを整列させる関数"""
    tokenized_inputs = tokenizer(
        examples["tokens"], truncation=True, is_split_into_words=True
    )
    all_labels = examples["ner_tags"]
    new_labels = []
    for i, labels in enumerate(all_labels):
        word_ids = tokenized_inputs.word_ids(i)
        new_labels.append(align_labels_with_tokens(labels, word_ids))

    tokenized_inputs["labels"] = new_labels
    return tokenized_inputs

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

from transformers import DataCollatorForTokenClassification

# トークン分類用のデータコレーターを作成
data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)

# データコレーターをテスト
batch = data_collator([tokenized_datasets["train"][i] for i in range(2)])
print(batch["labels"])

# 元のラベルと比較
for i in range(2):
    print(tokenized_datasets["train"][i]["labels"])

# seqeval評価指標を読み込み
import evaluate

metric = evaluate.load("seqeval")

# 最初の訓練例のラベルを文字列に変換
labels = raw_datasets["train"][0]["ner_tags"]
label_names = raw_datasets["train"].features["ner_tags"].feature.names
labels = [label_names[i] for i in labels]
print(labels)

# 偽の予測を作成してメトリクスをテスト
predictions = labels.copy()
predictions[2] = "O"
result = metric.compute(predictions=[predictions], references=[labels])
print(result)

import numpy as np

def compute_metrics(eval_preds):
    """評価指標を計算する関数"""
    logits, labels = eval_preds
    predictions = np.argmax(logits, axis=-1)

    # 無視されるインデックス（特殊トークン）を削除してラベルに変換
    true_labels = [[label_names[l] for l in label if l != -100] for label in labels]
    true_predictions = [
        [label_names[p] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]

    # seqevalメトリクスを計算
    all_metrics = metric.compute(predictions=true_predictions, references=true_labels)
    return {
        "precision": all_metrics["overall_precision"],
        "recall": all_metrics["overall_recall"],
        "f1": all_metrics["overall_f1"],
        "accuracy": all_metrics["overall_accuracy"],
    }

# ラベルマッピング辞書を作成
id2label = {i: label for i, label in enumerate(label_names)}
label2id = {v: k for k, v in id2label.items()}
print(f"id2label: {id2label}")
print(f"label2id: {label2id}")

from transformers import AutoModelForTokenClassification

# トークン分類モデルを作成
model = AutoModelForTokenClassification.from_pretrained(
    model_checkpoint,
    id2label=id2label,
    label2id=label2id
)

# ラベル数を確認
print(f"ラベル数: {model.config.num_labels}")

from transformers import TrainingArguments

# 訓練引数を設定
args = TrainingArguments(
    "bert-finetuned-ner",           # 出力ディレクトリ
    eval_strategy="epoch",          # エポックごとに評価
    save_strategy="epoch",          # エポックごとに保存
    learning_rate=2e-5,            # 学習率
    num_train_epochs=3,            # 訓練エポック数
    weight_decay=0.01,             # 重み減衰
    push_to_hub=False,             # Hubにプッシュしない
)

from transformers import Trainer

# Trainerを作成
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    processing_class=tokenizer
)

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

from torch.utils.data import DataLoader

# データローダーを作成
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 = AutoModelForTokenClassification.from_pretrained(
    model_checkpoint,
    id2label=id2label,
    label2id=label2id
)

from torch.optim import AdamW

# オプティマイザーを作成
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
)

def postprocess(predictions, labels):
    """後処理関数：予測とラベルを文字列リストに変換"""
    predictions = predictions.detach().cpu().clone().numpy()
    labels = labels.detach().cpu().clone().numpy()

    # 無視されるインデックスを削除してラベルに変換
    true_labels = [[label_names[l] for l in label if l != -100] for label in labels]
    true_predictions = [
        [label_names[p] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]
    return true_labels, true_predictions

from tqdm import tqdm
import torch

# 訓練ループを実行
progress_bar = tqdm(range(num_training_steps))
output_dir = "bert-finetuned-ner-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()
    for batch in eval_dataloader:
        with torch.no_grad():
            outputs = model(**batch)

        predictions = outputs.logits.argmax(dim=-1)
        labels = batch["labels"]

        # 予測とラベルを収集のためにパディング
        predictions = accelerator.pad_across_processes(predictions, dim=1, pad_index=-100)
        labels = accelerator.pad_across_processes(labels, dim=1, pad_index=-100)

        predictions_gathered = accelerator.gather(predictions)
        labels_gathered = accelerator.gather(labels)

        true_predictions, true_labels = postprocess(predictions_gathered, labels_gathered)
        metric.add_batch(predictions=true_predictions, references=true_labels)

    # 評価指標を計算・表示
    results = metric.compute()
    print(
        f"epoch {epoch}:",
        {
            key: results[f"overall_{key}"] 
            for key in ["precision", "recall", "f1", "accuracy"]
        }
    )

    # モデルとtokenizerを保存
    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)

# 全プロセスの同期を待つ
accelerator.wait_for_everyone()
# ベースモデルを取得
unwrapped_model = accelerator.unwrap_model(model)
# モデルを保存
unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)

from transformers import pipeline

# ファインチューニングしたモデルでパイプラインを作成
# (実際のチェックポイントに置き換えてください)
model_checkpoint = "bert-finetuned-ner-accelerate"
token_classifier = pipeline(
    "token-classification", 
    model=model_checkpoint, 
    aggregation_strategy="simple"
)

# モデルをテスト
result = token_classifier("My name is Sylvain and I work at Hugging Face in Brooklyn.")
print(result)