【Python】BERTの実装方法｜文章分類, pytorch

BERTを使えば簡単に文章分類をすることができます。

from torch import nn
from transformers import AutoModel

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.bert = AutoModel.from_pretrained("bert-base-uncased")
        self.classifier = nn.Linear(in_features = 768, out_features = 20)
    
    def forward(self, input_ids, attention_mask, token_type_ids):
        outputs = self.bert(input_ids = input_ids, attention_mask = attention_mask, token_type_ids = token_type_ids)
        pooler_output = outputs.pooler_output
        logits = self.classifier(pooler_output).squeeze(-1)
        return logits

BERTは計算が重くGPUが必須なので、CPUで計算したいなら以下の方法を使いましょう。

GPUがないなら、GoogleColabを使うといいですよ。

データ準備

インポート

from sklearn.datasets import fetch_20newsgroups
import pandas as pd

train_data = fetch_20newsgroups(subset = 'train')
valid_data = fetch_20newsgroups(subset = 'test')

train = pd.DataFrame({"text" : train_data["data"], "target" : train_data["target"]})
valid = pd.DataFrame({"text" : valid_data["data"], "target" : valid_data["target"]})
print(train.shape, valid.shape)
train.head()

今回使う文章データを準備します。

scikit-learnから文章のデータセットをインポートしましょう。

“fetch_20newsgroups”は文章が20個のグループに分けられたデータセットです。

こんな感じで文章(text)と予測したい値(target)があります。

“target”は０～19の数値で、20個のグループを表現しています。

データを見てみます。↓

text = train["text"].values[0]
print(text)

メールの文章です。

この文章の”target”は７です。↓

target = train["target"].values[0]
print(target)

“target”の種類は”target_names”に入っています。

print(train_data.target_names)
print(train_data.target_names[7])

これの７個目は”rec.autos”で、つまり車関係の文章という意味です。

こんな感じで文章とそのグループが与えられています。

Text Cleaning

文章の改行やスペースが多いのできれいにしましょう。

print(text)

現状はこの通りです。

“re”という正規表現ライブラリを使えばうまく処理できます。

import re
def cleaning(text):
    text = re.sub("\n", " ", text) # 改行削除
    text = re.sub("[^A-Za-z0-9]", " ", text) # 記号削除
    text = re.sub("[' ']+", " ", text) # スペース統一
    return text.lower() # 小文字で出力

この処理方法が必ず正しいわけではありませんので、参考程度にしてください。

train["cleaned_text"] = train["text"].map(cleaning)
valid["cleaned_text"] = valid["text"].map(cleaning)
train.head()

pandasのmapを使って各データに関数を実行しました。

こんな感じで改行がない小文字になっていればOKです。

Tokenizer

文章をそのままモデルに入れて計算することはできません。

なので、文章を数値(id)に変換するTokenizerを用意します。

Tokenizerの作成

!pip install transformers -q
from transformers import AutoTokenizer

MODEL_NAME = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

“bert-base-uncased”という形式のTokenizerを作ったと理解しておけば大丈夫です。

bertの他にもrobertaとかdebertaとか色々あり、性能も違います。

text = train["cleaned_text"].values[0]
print(text)

print("\n")

encoded = tokenizer(text)
print(encoded.keys())

辞書型データになっており、キーは以下の通りです。

input_ids

input_ids = encoded["input_ids"]
print(len(input_ids))
print(input_ids)

input_idsは文章をトークンと呼ばれるidに変換したものです。

もちろんidの数はもともとの文章の長さによって変わります。

正しく変換されているかは、逆変換であるdecodeを使えば確かめられます。↓

print(text[:100])
print(tokenizer.decode(input_ids)[:100])

こんな感じで、Tokenizerの中では文章を数値で認識できています。

attention_mask

attention_mask = encoded["attention_mask"]
print(len(attention_mask))
print(attention_mask)

attention_maskは学習時に使用するトークンと無視するトークンを識別するものです。

１なら使用、０なら無視されます。今はすべて１のはずです。

例えば、後で解説するパディングで追加されたダミーのidは不要になります。

なのでattention_mask = 0にして不要なデータは無視できるようにします。

token_type_ids

token_type_ids = encoded["token_type_ids"]
print(len(token_type_ids))
print(token_type_ids)

文章を２つ入力するときに使用します。

１つ目の文章にあるトークンは０、２つ目の文章にあれば１になります。

２つの文章をインプットしたい場合に活用されますが、今回は１つだけなので意味ありません。

パディング

text1 = train["cleaned_text"].values[0]
text2 = train["cleaned_text"].values[1]

encoded = tokenizer(text1)
print(len(encoded["input_ids"]))

encoded = tokenizer(text2)
print(len(encoded["input_ids"]))

例えば、１つ目のデータの文章をTokenizerで変換すると154個のデータです。

しかし２つ目の文章では171個になっており、データの数が違います。

まずは、各文章をidに変換したときの長さを調査します。↓

import matplotlib.pyplot as plt

length = []
for text in train["cleaned_text"].values:
    encoded = tokenizer.encode_plus(text.lower())
    length.append(len(encoded["input_ids"]))
plt.hist(length)
plt.show()

import numpy as np
length = np.array(length)
print(np.quantile(length, q = 0.5))
print(np.quantile(length, q = 0.75))

idの長さは、中央値が233、75%が368くらいのはずです。

今回は計算を軽くしたいので256に統一しましょう。

MAX_LEN = 256
text = train["cleaned_text"].values[0]
encoded = tokenizer(text, padding = "max_length", max_length = MAX_LEN, truncation = True)
print(len(encoded["input_ids"]))
print(encoded["input_ids"][-10:])
print(encoded["attention_mask"][:10])
print(encoded["attention_mask"][-10:])

つまり、MAX_LEN = 256に足りない分は埋めて、超えた分はカットしています。

ここでattention_maskを確認すると、最初は１が多く最後の方は０だけになります。

これは、パディングで埋められた分のデータは学習に使わないので、０１で区別するためです。

print(tokenizer.decode(encoded["input_ids"]))

こんな感じで、埋められた位置は[PAD]として表現されています。

Dataset, DataLoader作成

pytorchのモデル作成で必要なことは以下の記事で解説しています。

import torch
from torch.utils.data import Dataset, DataLoader

class Data(Dataset):
    def __init__(self, data, tokenizer, max_length):
        self.data = data
        self.tokenizer = tokenizer
        self.max_length = max_length
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        text = self.data["cleaned_text"].values[idx]
        encoded = tokenizer(
            text,
            padding = "max_length",
            max_length = self.max_length,
            truncation = True
        )
        input_ids = torch.tensor(encoded["input_ids"], dtype = torch.int32)
        attention_mask = torch.tensor(encoded["attention_mask"], dtype = torch.int32)
        token_type_ids = torch.tensor(encoded["token_type_ids"], dtype = torch.int32)
        label = self.data["target"].values[idx]
        label = torch.tensor(label, dtype = torch.int32)
        return input_ids, attention_mask, token_type_ids, label

input_ids, attention_mask, token_type_idsを作り、出力します。

pytorchで使うtensor型にしておきましょう。

以下のコードでDataset, DataLoaderを作ります。↓

BATCH_SIZE = 8
train_ds = Data(train, tokenizer, MAX_LEN)
train_dl = DataLoader(train_ds, batch_size = BATCH_SIZE, shuffle = True, drop_last = True)
valid_ds = Data(valid, tokenizer, MAX_LEN)
valid_dl = DataLoader(valid_ds, batch_size = BATCH_SIZE * 2, shuffle = False, drop_last = False)

計算が重いのでメモリエラーになった場合はバッチサイズを下げましょう。

モデル作成

from torch import nn
from transformers import AutoModel

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.bert = AutoModel.from_pretrained(MODEL_NAME) # <-"bert-base-uncased"
        self.classifier = nn.Linear(in_features = 768, out_features = 20)
    
    def forward(self, input_ids, attention_mask, token_type_ids):
        outputs = self.bert(
            input_ids = input_ids, attention_mask = attention_mask, token_type_ids = token_type_ids
        )
        pooler_output = outputs.pooler_output
        logits = self.classifier(pooler_output).squeeze(-1)
        return logits

model = Model().cuda()

TokenizerとModelとで同じBERTの種類にしておきましょう。

BERTの出力には複数の形式があります。今回は”pooler_output”にしました。

この出力では768個のデータがあるので、Linear層で受け取ってラベル数20にしましょう。

cudaはGPUをONにしている場合に付け足してください。

CPUで実行するならcudaは不要ですが、計算が重くて終わらないのでおススメしません。

学習

最適化手法と損失関数を定義します。

optimizer = torch.optim.Adam(model.parameters(), lr = 1e-5)
criterion = nn.CrossEntropyLoss()

学習率が高いとうまくlossが下がらないので、1e-5にしました。

以下のコードで学習させます。↓

import numpy as np
best_loss = np.inf

from tqdm import tqdm

for epoch in range(5):
    model.train() # 学習モード
    train_loss = 0
    for batch in tqdm(train_dl):
        optimizer.zero_grad() # 勾配リセット
        input_ids = batch[0].cuda()
        attention_mask = batch[1].cuda()
        token_type_ids = batch[2].cuda()
        label = batch[3].long().cuda()
        preds = model(input_ids, attention_mask, token_type_ids) # 予測計算
        loss = criterion(preds, label) # 誤差計算
        loss.backward() # 誤差伝播
        optimizer.step() # パラメータ更新
        train_loss += loss.item()
    train_loss /= len(train_dl)

    model.eval()
    valid_loss = 0
    with torch.no_grad(): # 必須
        for batch in tqdm(valid_dl):
            input_ids = batch[0].cuda()
            attention_mask = batch[1].cuda()
            token_type_ids = batch[2].cuda()
            label = batch[3].long().cuda()
            preds = model(input_ids, attention_mask, token_type_ids)
            loss = criterion(preds, label)
            valid_loss += loss.item()
        valid_loss /= len(valid_dl)

    print(f"EPOCH[{epoch}]")
    print(f"TRAIN: {train_loss}")
    print(f"VALID: {valid_loss}")
    if valid_loss < best_loss: # best_lossが最小の時保存する
        best_loss = valid_loss
        torch.save(model.state_dict(), "best_model.pth")

pytorchの使い方は以下の記事を参考にしてください。

学習時間はGPUでも長いです。CPUだともはや終わりません。

学習と検証を５回繰り返し、最も誤差の小さいモデルを保存するようにしました。

import matplotlib.pyplot as plt
plt.plot(history["train"], label = "train")
plt.plot(history["valid"], label = "valid")
plt.legend()
plt.show()

こんな感じで徐々にlossが下がっていますが、検証データでは上がったり下がったりしています。

予測結果

pred = []

model.load_state_dict(torch.load("best_model.pth", map_location = "cuda"))
model.eval()
with torch.no_grad():
    for batch in tqdm(valid_dl):
        input_ids = batch[0].cuda()
        attention_mask = batch[1].cuda()
        token_type_ids = batch[2].cuda()
        logits = model(input_ids, attention_mask, token_type_ids).cpu().numpy()
        pred.extend(logits.argmax(axis = 1))

lossが最小のときのモデルを読み込んで予測します。

20列の中で最も値の大きい列番号が予測したいラベルです。

truth = valid["target"].values
print(pred[:10])
print(truth[:10])

こんな感じであっていたり間違っていたりしていますね。

正解率、適合率、再現率を見てみましょう。↓

from sklearn.metrics import accuracy_score, precision_score, recall_score
truth = valid["target"].values
accuracy = accuracy_score(truth, pred)
precision = precision_score(truth, pred, average = "macro")
recall = recall_score(truth, pred, average = "macro")
print("ACC:", accuracy.round(3), "PRE:", precision.round(3), "REC:", recall.round(3))

うまくいっていれば80%を超えているかと思います。

まとめ

今回はBERTの使い方を解説しました。

もし計算が重すぎて手元でできない場合は、CPUでLightGBMを使うといいですよ。