Transformers 日本語BERT クラス分類モデルのファインチューニング

$ tar xvf ldcc-20140209.tar.gz

import os
from glob import glob
import pandas as pd
import linecache
from sklearn.metrics import classification_report


# カテゴリを配列で取得
categories = [name for name in os.listdir("text") if os.path.isdir("text/" + name)]
print(categories)

datasets = pd.DataFrame(columns=["text", "category"])
for cat in categories:
    path = "text/" + cat + "/*.txt"
    files = glob(path)
    for text_name in files:
        body = ''
        for i in range(3,1000):
            text=linecache.getline(text_name, i)
            if text=='':
                break
            body+=text

        datasets.loc[len(datasets)]=[body,cat]

# データフレームシャッフル、先頭の200データだけ抽出
datasets = datasets.sample(frac=1,random_state=0).reset_index(drop=True)
datasets = datasets.iloc[:200]
print(datasets)

# 正解ラベル（カテゴリー）をデータセットから取得
categories = list(set(datasets['category']))
print(categories)

# カテゴリーのID辞書を作成
id2cat = dict(zip(list(range(len(categories))), categories))
cat2id = dict(zip(categories, list(range(len(categories)))))
print(id2cat)
print(cat2id)

# DataFrameにカテゴリーID列を追加
datasets['label'] = datasets['category'].map(cat2id)

# データセットを本文とカテゴリーID列だけにする
datasets = datasets[['text', 'label']]
display(datasets.head())

from datasets import Dataset

dataset_packed = Dataset.from_pandas(datasets)
dataset_split = dataset_packed.train_test_split(test_size=0.2, seed=0)
print(dataset_split)

from transformers import AutoTokenizer

tokenizer= AutoTokenizer.from_pretrained('cl-tohoku/bert-base-japanese-v2')

def preprocess_function(examples):
    MAX_LENGTH = 512
    return tokenizer(examples["text"], max_length=MAX_LENGTH, truncation=True)

tokenized_dataset = dataset_split.map(preprocess_function, batched=True)

from transformers import DataCollatorWithPadding
from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
model = AutoModelForSequenceClassification.from_pretrained("cl-tohoku/bert-base-japanese-v2", num_labels=9)

from sklearn.metrics import accuracy_score, f1_score

def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    f1 = f1_score(labels, preds, average='weighted')
    acc = accuracy_score(labels, preds)
    return {'accuracy':acc, 'f1':f1}

training_args = TrainingArguments(
    output_dir="./results",
    evaluation_strategy='epoch',
    logging_strategy='epoch',
    save_strategy='epoch',
    save_total_limit=1,
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=5,
    weight_decay=0.01,
    no_cuda=True, # GPUを使用する場合はFalse
)

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

trainer.train()

trainer.save_state()
trainer.save_model()

pred_result = trainer.predict(tokenized_dataset['test'], ignore_keys=['loss', 'last_hidden_state', 'hidden_states', 'attentions'])
pred_label= pred_result.predictions.argmax(axis=1).tolist()
print(pred_label)

from sklearn.metrics import classification_report
print(classification_report(tokenized_dataset['test']['label'], pred_label, target_names=categories))