PyTorchのBERTで日本語ニュース分類

379
PyTorchのBERTで日本語ニュース分類
目次

はじめに

Hugging FaceがTransformersというライブラリを公開しており、その中にはBERTの実装、さらには日本語対応した事前学習モデルが用意されています。

このライブラリを使えば、BERTを使って簡単にさまざまなfine-tuningができるようなので、試してみました。

以下では、BERTを使ったカテゴリ分類の手順について説明をしていきます。

BERTによる日本語ニュース分類

今回は、日本語のニュース（ライブドアニュース）をカテゴリ分類する簡単なfine-tuningを行います。

BERTを使った学習は、2段階になっています。
1段階目は事前学習で、ここで汎用的な言語モデルの作成が行われます。
2段階目はfine-tuningで、事前学習で作成されたモデルに対して、課題に応じた学習を行います。

今回は、事前学習済みモデルを使って、fine-tuningを行います。

環境

macOS 12.5
Python 3.10.3

ライブラリのインストール

まず、今回のプログラムを実行するのに必要なライブラリをインストールします。

pip install scikit-learn
pip install transformers
pip install fugashi
pip install ipadic
pip install datasets

transformersが今回使用するメインのライブラリ、fugashi、ipadicは日本語のトークナイザーで必要となります。

データの準備

ダウンロードしたlivedoorニュースを加工し、タグ（カテゴリ）とセットにしてCSVとして保存するプログラムdataset.pyを作成します。

データセットの準備

RONDHUIT社のダウンロードページから、ldcc-20140209.tar.gzをダウンロードし、任意のディレクトリに解凍します。
本稿では、作成するプログラムと同じディレクトリに解凍し、ディレクトリ名をlivedoor_newsとしています。

プログラムを実行する前に、あらかじめdata、model、results、logsディレクトリを作成しておいてください（classification.py、dataset.py、train.pyは）。

.
├── data
├── livedoor_news
│   ├── dokujo-tsushin
│   ├── it-life-hack
│   ├── kaden-channel
│   ├── livedoor-homme
│   ├── movie-enter
│   ├── peachy
│   ├── smax
│   ├── sports-watch
│   ├── topic-news
│   ├── CHANGES.txt
│   └── README.txt
├── model
├── results
├── logs
├── classification.py
├── dataset.py
└── train.py

下記のプログラムでは、ニュースの各ファイルの本文を抽出し、改行、全角スペース、タブを除去して、1行の文に変換したのちに、カテゴリの番号（0-9）とのセットにします。


# dataset.py

import glob
import os

raw_data_path = "./livedoor_news"  # ライブドアニュースを格納したディレクトリ

dir_files = os.listdir(path=raw_data_path)
dirs = [f for f in dir_files if os.path.isdir(os.path.join(raw_data_path, f))]
text_label_data = []  # 文章とラベル（カテゴリ）のセット

for i in range(len(dirs)):
    dir = dirs[i]
    files = glob.glob(os.path.join(raw_data_path, dir, "*.txt"))

    for file in files:
        if os.path.basename(file) == "LICENSE.txt": # 各ディレクトリにあるLICENSE.txtを除外する
            continue

        with open(file, "r") as f:
            text = f.readlines()[3:]
            text = "".join(text)
            text = text.translate(str.maketrans({"\n":"", "\t":"", "\r":"", "\u3000":""})) 
            text_label_data.append([text, i])

学習用、テスト用データの作成、保存

先ほど作成した、本文、ラベル（カテゴリ）のセットを、学習用、評価用のデータに分割し、CSVファイル（news_train.csv、news_test.csv）として保存します。


# dataset.py

import csv
from sklearn.model_selection import train_test_split

news_train, news_test =  train_test_split(text_label_data, shuffle=True)  # データを学習用とテスト用に分割
data_path = "./data"

with open(os.path.join(data_path, "news_train.csv"), "w") as f:
    writer = csv.writer(f)
    writer.writerows(news_train) # 

with open(os.path.join(data_path, "news_test.csv"), "w") as f:
    writer = csv.writer(f)
    writer.writerows(news_test)

学習

先ほど作成したデータを入力とし、BERTを使ってニュースのカテゴリ分類を学習（fine-tuning）させるプログラムtrain.pyを作成します。

モデル、トークナイザーの読み込み

transformersに含まれている文章を分類するためのモデルBertForSequenceClassification、日本語を形態素解析するためのトークナイザーBertJapaneseTokenizerを読み込みます。

cl-tohoku/bert-base-japanese-whole-word-maskingは事前学習済みの日本語BERTモデルです。
このモデルは、東北大学の乾研究室によって作成されたもので、こちらのページで公開されています。

BertForSequenceClassification、BertJapaneseTokenizerを読み込んだ際に、自動的にダウンロードされるため、あらかじめダウンロードをする必要はありません。


# train.py

from transformers import BertForSequenceClassification, BertJapaneseTokenizer

# モデル
model = BertForSequenceClassification.from_pretrained("cl-tohoku/bert-base-japanese-whole-word-masking", num_labels=9)
# model.cuda() # cudaを使う場合は、この行を有効にする
# トークナイザー
tokenizer = BertJapaneseTokenizer.from_pretrained("cl-tohoku/bert-base-japanese-whole-word-masking")

学習用、テスト用データの読み込み

学習の入力データとして、先ほど保存したデータを読み込みます。


# train.py

import os
from datasets import load_dataset

# トークナイズ用関数
def tokenize(batch):
    return tokenizer(batch["text"], padding=True, truncation=True, max_length=128)

data_path = "./data"

# 学習用データ
train_data = load_dataset("csv", data_files=os.path.join(data_path, "news_train.csv"), column_names=["text", "label"], split="train")
train_data = train_data.map(tokenize, batched=True, batch_size=len(train_data))
train_data.set_format("torch", columns=["input_ids", "label"])

# テスト用データ
test_data = load_dataset("csv", data_files=os.path.join(data_path, "news_test.csv"), column_names=["text", "label"], split="train")
test_data = test_data.map(tokenize, batched=True, batch_size=len(test_data))
test_data.set_format("torch", columns=["input_ids", "label"])

Trainerの初期化

Trainerに、学習対象のモデル、学習用パラメーター、評価用関数、学習用データ、評価用データを設定して初期化します。


# train.py

# 評価用関数
from sklearn.metrics import accuracy_score

def compute_metrics(result):
    labels = result.label_ids
    preds = result.predictions.argmax(-1)
    acc = accuracy_score(labels, preds)
    return {
        "accuracy": acc,
    }

# Trainerの設定
from transformers import Trainer, TrainingArguments

# 学習用パラメーター
training_args = TrainingArguments(
    output_dir = "./results",
    num_train_epochs = 2,
    per_device_train_batch_size = 8,
    per_device_eval_batch_size = 32,
    warmup_steps = 500,  # 学習係数が0からこのステップ数で上昇
    weight_decay = 0.01,  # 重みの減衰率
    # evaluate_during_training = True,  # ここの記述はバージョンによっては必要ありません
    logging_dir = "./logs",
)

# Trainerの初期化
trainer = Trainer(
    model = model, # 学習対象のモデル
    args = training_args, # 学習用パラメーター
    compute_metrics = compute_metrics, # 評価用関数
    train_dataset = train_data, # 学習用データ
    eval_dataset = test_data, # テスト用データ
)

モデルの学習

Trainerを使って、モデルの学習、評価を行います。


# train.py

trainer.train() # 学習
trainer.evaluate() # 評価

cudaが使えない手元のマシンで実行したため、3時間弱かかってしまいました。
ちなみに、Google Colaboratoryでcudaを使って実行した場合は、5分弱で終わりました。

学習済みモデルの保存

後で使うために、学習済みモデル、トークナイザーを保存します。


# train.py

model_dir = "./model"
trainer.save_model(model_dir)
tokenizer.save_pretrained(model_dir)

分類の実行

では、学習したモデルを使って、ニュースが正しく分類できるかどうかを確認します。
今回はsports-watchディレクトリにあるsports-watch-4764756.txtを入力として分類を行います。


# classification.py

import os
import torch
from transformers import BertForSequenceClassification, BertJapaneseTokenizer

# 学習済みモデルの読み込み
model_dir = "./model"
loaded_model = BertForSequenceClassification.from_pretrained(model_dir)
# loaded_model.cuda() # cudaを使う場合は、この行を有効にする
loaded_tokenizer = BertJapaneseTokenizer.from_pretrained(model_dir)

# 分類するデータの読み込み
file = "./livedoor_news/sports-watch/sports-watch-4764756.txt"  # sports-watchの適当なニュース

with open(file, "r") as f:
    sample_text = f.readlines()[3:]
    sample_text = "".join(sample_text)
    sample_text = sample_text.translate(str.maketrans({"\n":"", "\t":"", "\r":"", "\u3000":""})) 

max_length = 512
words = loaded_tokenizer.tokenize(sample_text)
word_ids = loaded_tokenizer.convert_tokens_to_ids(words)  # 単語をインデックスに変換
word_tensor = torch.tensor([word_ids[:max_length]])  # Tensorに変換

# 予測の実行
# word_tensor.cuda()  # cudaを使う場合は、この行を有効にする
y = loaded_model(word_tensor)  # 結果の予測
pred = y[0].argmax(-1)  # 最大値のインデックス（ディレクトリの番号）

# 結果の標準
path = "./livedoor_news"
dir_files = os.listdir(path=path)
dirs = [f for f in dir_files if os.path.isdir(os.path.join(path, f))]  # ディレクトリ一覧
print("結果は", dirs[pred])

実行すると

result: sports-watch

と表示され、正しく分類ができることを確認できました。

まとめ

以前は、GPUマシンで数日間かけてBERTに日本語の文章を食わせて、事前学習を行なっていましたが、便利なライブラリの出現により、BERTを使って簡単にさまざまな学習（fine tuning）を試すことができるようになりました。

また、実装が必要なのは、データの作成、読み込みが大部分で、学習自体は数行で行えるため非常に簡単です。

今後は、このライブラリを使って様々なタスクを試していこうと思います。

今回作成したプログラムはGitHubで公開しています。
https://github.com/age884/bert_japanese_news_classification