BERT + LightGBM + optuna で手軽に自然言語処理モデルを構築する

BERTとLightGBM, optunaで自然言語の分類モデルをサクッと作ってみようという記事です。
データはlivedoorニュースコーパスを使用しています。

また、本記事で使用しているコードはすべて以下にあります。
https://github.com/kazuki-hayakawa/bert_lightgbm_model

実際に実行するときはgit cloneしてからお手元で試してみてください。

全体の流れ

1. livedoorニュースコーパスのデータのダウンロード
2. 日本語学習済みBERTモデルのダウンロード
3. 実験用コンテナの起動
4. 特徴量の生成
5. モデルのトレーニング
6. テストデータによるモデルの評価
7. コンテナの終了

livedoorニュースコーパスのデータのダウンロード

livedoorニュースコーパスのデータを data/raw ディレクトリにダウンロードします。ダウンロード用のスクリプトは src/data/download_livedoor_news.sh にまとめてあります。

その後、 src/data/preprocess.py を実行して前処理をし、トレーニング用とテスト用にデータを分割して保存します。

src/data/preprocess.py

import os
import glob
from tqdm import tqdm
import pandas as pd
from sklearn.model_selection import train_test_split


def read_text(text_filepath):
    """ livedoor ニュースの形式に合わせて、4行目以降の本文のみを読み取る """
    with open(text_filepath, 'r') as f:
        lines = f.readlines()
        lines = lines[3:]

    text = ' '.join(lines)
    # 全角スペース、改行コードの削除
    text = text.replace('\u3000', '').replace('\n', '')
    return text


def main():
    # 事前に download_livedoor_news.sh を実行してデータを取得しておく
    exclude_files = ['CHANGES.txt', 'README.txt', 'LICENSE.txt']
    all_file_paths = glob.glob('../../data/raw/text/**/*.txt', recursive=True)
    all_file_paths = [p for p in all_file_paths
                      if os.path.basename(p) not in exclude_files]

    df_processed = pd.DataFrame(columns=['id', 'media', 'text'])
    for idx, filepath in enumerate(tqdm(all_file_paths)):
        media = os.path.dirname(filepath).replace('../../data/raw/text/', '')
        text = read_text(filepath)
        row = pd.Series([idx + 1, media, text], index=df_processed.columns)
        df_processed = df_processed.append(row, ignore_index=True)

    df_train, df_test, _, _ = train_test_split(
        df_processed, df_processed['media'], test_size=0.1, random_state=0,
        stratify=df_processed['media']
    )
    df_train.to_csv('../../data/processed/train_dataset.csv', index=False)
    df_test.to_csv('../../data/processed/test_dataset.csv', index=False)


if __name__ == '__main__':
    main()

日本語学習済みBERTモデルのダウンロード

手順はbert-as-serviceを使って日本語BERTの文エンベディング計算サーバーを作るを参照しています。

日本語学習済みBERTモデルを models/bert_jp ディレクトリを作成し、ダウンロードしておきます。

bert-as-service でロードできるようファイル名の変更

mv model.ckpt-1400000.index bert_model.ckpt.index
mv model.ckpt-1400000.meta bert_model.ckpt.meta 
mv model.ckpt-1400000.data-00000-of-00001 bert_model.ckpt.data-00000-of-00001

語彙ファイルの作成

cut -f1 wiki-ja.vocab | sed -e "1 s/<unk>/[UNK]/g" > vocab.txt

BERT設定ファイルの作成

bert_jp/bert_config.json

{
    "attention_probs_dropout_prob" : 0.1,
    "hidden_act" : "gelu",
    "hidden_dropout_prob" : 0.1,
    "hidden_size" : 768,
    "initializer_range" : 0.02,
    "intermediate_size" : 3072,
    "max_position_embeddings" : 512,
    "num_attention_heads" : 12,
    "num_hidden_layers" : 12,
    "type_vocab_size" : 2,
    "vocab_size" : 32000
}

実験用コンテナの起動

docker-compose up -d を実行してコンテナを起動します。（ Dockerfile, docker-compose.yml はGitHubのリポジトリを参照ください）
その後、 docker-compose exec analytics /bin/bash を実行しコンテナに入ります。

特徴量の生成

BERTを操作する Bert クラスは以下のように実装しています。

src/features/bert.py

import sentencepiece as spm
from bert_serving.client import BertClient


class Bert():
    """ Bert model client
        Before usage, you need to run bert server.
    """

    def __init__(self, bert_model_path, client_ip='0.0.0.0'):
        self.bert_client = BertClient(ip=client_ip)
        self.spm_model = spm.SentencePieceProcessor()
        self.spm_model.load(bert_model_path + 'wiki-ja.model')

    def _parse(self, text):
        text = str(text).lower()
        encoded_texts = self.spm_model.EncodeAsPieces(text)
        encoded_texts = [t for t in encoded_texts if t.strip()]
        return encoded_texts

    def text2vec(self, texts):
        """

        Args:
            texts (list): 日本語文字列のリスト

        Returns:
            numpy array: テキストの分散表現テンソル
        """

        parsed_texts = list(map(self._parse, texts))
        tensor = self.bert_client.encode(parsed_texts, is_tokenized=True)
        return tensor

これを利用して自然言語をベクトルに変換します。同時に、目的変数となるメディア名も整数型のラベルに変換します。

src/features/build_features.py を実行します。

src/features/build_features.py

import subprocess
import numpy as np
import pandas as pd
from sklearn.preprocessing import LabelEncoder
from bert import Bert


def build_features(df, bert_client):
    vectors = bert_client.text2vec(df['text'])
    le = LabelEncoder()
    targets = le.fit_transform(df['media'])
    return vectors, targets


def main():
    BERT_MODEL_PATH = '../../models/bert_jp/'

    # start bert server
    commands = ['bert-serving-start', '-model_dir',
                BERT_MODEL_PATH, '-num_worker=1', '-cpu']
    p = subprocess.Popen(commands, shell=False,
                         stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    # start bert client
    bert = Bert(bert_model_path=BERT_MODEL_PATH, client_ip='0.0.0.0')

    # build train features
    train_dataset = pd.read_csv('../../data/processed/train_dataset.csv')
    train_vectors, train_targets = build_features(train_dataset, bert)
    np.save('../../data/features/train_vectors', train_vectors)
    np.save('../../data/features/train_targets', train_targets)

    # build test features
    test_dataset = pd.read_csv('../../data/processed/test_dataset.csv')
    test_vectors, test_targets = build_features(test_dataset, bert)
    np.save('../../data/features/test_vectors', test_vectors)
    np.save('../../data/features/test_targets', test_targets)

    p.terminate()


if __name__ == '__main__':
    main()

モデルのトレーニング

ニュース媒体を分類するモデルを定義したクラス MediaClassifier を以下のように実装しています。

src/models/classifier.py

import os
import uuid
import pickle
import numpy as np
import lightgbm as lgb
import optuna
from datetime import datetime, timedelta, timezone
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score


class MediaClassifier():
    """ livedoorニュースコーパスの多クラス分類をするモデル """

    def __init__(self, output_dir, use_gpu=False):
        JST = timezone(timedelta(hours=+9), 'JST')
        dt_now = datetime.now(JST)
        training_date = dt_now.strftime("%Y%m%d_%H%M%S")
        self.output_dir = os.path.join(output_dir, training_date)
        os.makedirs(self.output_dir, exist_ok=True)

        self.device = 'gpu' if use_gpu else 'cpu'

    def train(self, features, targets):
        X_train, X_test, y_train, y_test = train_test_split(
            features, targets, test_size=0.2, random_state=0)

        def objectives(trial):
            trial_uuid = str(uuid.uuid4())
            trial.set_user_attr("trial_uuid", trial_uuid)

            # パラメータとコールバックのセッティング
            params = {
                # liverdoorニュースコーパスの媒体数は9なので9つの多クラス分類
                'objective': 'multiclass',
                'num_class': 9,
                'metric': 'multi_logloss',
                'num_leaves': trial.suggest_int("num_leaves", 10, 500),
                'feature_fraction': trial.suggest_uniform("feature_fraction", 0.0, 1.0),
                'class_weight': 'balanced',
                'device': self.device,
                'verbose': -1
            }

            pruning_callback = optuna.integration.LightGBMPruningCallback(
                trial, "multi_logloss")

            # training
            lgb_model = lgb.train(params, lgb.Dataset(X_train, y_train), num_boost_round=100,
                                  valid_sets=lgb.Dataset(X_test, y_test), callbacks=[pruning_callback])

            y_pred_train = np.argmax(lgb_model.predict(X_train), axis=1)
            y_pred_test = np.argmax(lgb_model.predict(X_test), axis=1)
            accuracy_train = accuracy_score(y_train, y_pred_train)
            accuracy_test = accuracy_score(y_test, y_pred_test)

            trial.set_user_attr("accuracy_train", accuracy_train)
            trial.set_user_attr("accuracy_test", accuracy_test)

            # モデル保存
            output_file = os.path.join(self.output_dir, f"{trial_uuid}.pkl")
            with open(output_file, "wb") as fp:
                pickle.dump(lgb_model, fp)

            return 1.0 - accuracy_test

        study = optuna.create_study()
        study.optimize(objectives, n_trials=100)

        result_df = study.trials_dataframe()
        result_csv = os.path.join(self.output_dir, "result.csv")
        result_df.to_csv(result_csv, index=False)

        return study.best_trial.user_attrs

上記のモデルのトレーニングを実行します。

src/models/train_model.py

import numpy as np
from classifier import MediaClassifier


def main():
    train_vectors = np.load('../../data/features/train_vectors.npy')
    train_targets = np.load('../../data/features/train_targets.npy')

    model = MediaClassifier(output_dir='../../models/training_models',
                            use_gpu=False)

    best_result = model.train(train_vectors, train_targets)
    print('best result \n', best_result)


if __name__ == '__main__':
    main()

実行完了後に最も性能のよいモデルのtrial uuidとスコアが以下のように出力されるので、控えておきます。

{'trial_uuid': 'BEST_MODEL_TRIAL_UUID', 'accuracy_train': 1.0, 'accuracy_test': 0.7398190045248869}

BEST_MODEL_TRIAL_UUID の部分には実際にはuuidが入ります。

テストデータによるモデルの評価

テストデータによる評価は以下のように行います。

src/models/predict_model.py

import argparse
import pickle
import numpy as np
from sklearn.metrics import accuracy_score


def main(args):
    test_vectors = np.load('../../data/features/test_vectors.npy')
    test_targets = np.load('../../data/features/test_targets.npy')

    with open(args.best_model, 'rb') as f:
        model = pickle.load(f)

    pred_targets = np.argmax(model.predict(test_vectors), axis=1)
    accuracy = accuracy_score(test_targets, pred_targets)
    print('test accuracy : {:.2f}'.format(accuracy))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('--best_model', help='best model pickle file path.')

    args = parser.parse_args()

    main(args)

モデルのトレーニグ開始時点でディレクトリが自動生成されているので、そのパスを指定します。以下実行例です。

$ cd src/models
$ python predict_model.py --best_model='../../models/training_models/TRINING_DATE/BEST_MODEL_TRIAL_UUID.pkl'

test accuracy : 0.73

実際に私が作成したモデルで正解率が 0.73 でした。あまり高くはないですが、データ量を増やしたりテキストの区切り方など工夫すればもう少し上げれるのではないでしょうか。

コンテナの終了

exit でコンテナから出た後、 docker-compose down で終了します。