Pytorchを使ってテキスト生成モデルのT5を構築 〜Transformersでの転移学習による手軽な実践〜

スポンサーリンク

 
近年、自然言語処理のディープラーニングの発展はめざましく、Transformer をベースとした BERT、GPT-3、T5 によって次々にこれまでの精度を超えるモデルが構築されています。そして、自然言語処理のタスクのラスボスと言ってもいよいテキスト生成において、人間が作るものと遜色ないレベル、いや、それを上回るレベルまで達していると言われています。


テキスト生成モデルを自由に構築できれば、めんどくさいメールのやりとりから開放されたり、自分に似た Twitter Bot にいいねを稼いでもらえたりなど、楽しい未来が待っているはずです。ということで、今回は Transformer をベースとしたテキスト生成モデルである T5 をネットで入手できるデータを使って構築します。


今回は以下のリンク先を参考にさせてもらい、Pytorch と Transformers だけでモデルを構築できるようにしました。今回のコードはこちらに置いています。

github.com



準備


T5 のモデルをゼロから構築し、事前学習を自前で行うための莫大なリソースは持っていないので、すでに学習されたモデルを使って、さくっとそこそこの精度のテキスト生成モデルを作ります。そのために今回は、Hugging Face の sonoisa/t5-base-japanese のモデルを使わせてもらいます。このモデルを実際に扱うにはどうすればよいかなどは、コードを添えながら後述します。


今回、GPU を使って転移学習をするので、GPU を簡単に使える Google Colab で行います。もちろん CPU 環境でもモデルはできますが、子供の頃の夏休みの自由研究ぐらいのゆったりとした時間は必要になるかと思います。


Googleドライブのマウント


Google Colaboratory でモデルを作るとここで用意された環境にモデルができるため、自分の手元に残りません。自分の手元で自由に学習済みモデルを扱えるようにするために、自身のアカウントの Google Drive に学習済みモデルを残すようにさせましょう。そのために、Google Colaboratory と Google Drive を紐付けるのです。

from google.colab import drive
drive.mount('/content/drive')


上記のコードを実行することにより、自身のアカウントの Google Drive に紐付けられます。Google Drive 上に適当な作業フォルダーを用意して、その作業フォルダーをカレントディレクトリにします。今回は、Google Drive の MyDrive 直下に「text_generation」というフォルダー用意して、そこで作業します。

import os
path = '/content/drive/MyDrive/text_generation/'
if not os.path.exists(path):
    os.makedirs(path)

%cd /content/drive/MyDrive/text_generation/



設定用ファイルsettings.pyの用意


今回のコードで使用するパラメータの値を「settings.py」ファイルにまとめています。一つのファイルにまとめることで、パラメータがある場所が明確になるので、扱いやすくなるでしょう。

import torch

MODEL_NAME = "sonoisa/t5-base-japanese"
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

max_length_src = 30
max_length_target = 300

batch_size_train = 8
batch_size_valid = 8

epochs = 1000
patience = 20



モジュールのインポート


今回使用するモジュールを Google Colaboratory の環境にインストールし、インポートします。

pathlib
pandas
torch
transformers==4.5.0
sentencepiece
neologdn
sklearn
!pip install -r requirements.txt

from pathlib import Path
import re
import math
import time
import copy
from tqdm import tqdm
import pandas as pd
import tarfile
import neologdn
from sklearn.model_selection import train_test_split
import torch
import torch.nn as nn
from torch import optim
from torch.utils.data import DataLoader
from transformers import T5ForConditionalGeneration, T5Tokenizer
import settings



データの取得


今回作るテキスト生成モデルは、ニュースのタイトルを与えると、そのニュースの本文を生成するモデルです。これを作るためにニュースのデータとして、livedoor ニュースコーパスを使わせてもらいます。下記コードで対象データをダウンロードし、圧縮されたファイルを展開しています。

!wget -O ldcc-20140209.tar.gz https://www.rondhuit.com/download/ldcc-20140209.tar.gz

with tarfile.open("ldcc-20140209.tar.gz") as tar:
    tar.extractall()


展開されたファイルは、各ニュースのサイトごとにディレクトリ内に格納されており、その中のニューステキストファイルは、1行目がURL、2行目が日時です。そのため、テキストファイルを読み取る際に、冒頭の2行を飛ばしています。

file_paths = []
for dir_path in Path('text').glob('*/**'):
    for file_path in dir_path.glob('*'):
        if dir_path.name in file_path.name:
            file_paths.append(file_path)

data = []
for file_path in file_paths:
    with open(file_path, 'r') as file:
        lines = file.readlines()[2:]
        title = lines[0]
        body = ''.join(lines[1:])
        data.append((title, body))
data = pd.DataFrame(data, columns=['title', 'body'])



学習データの作成


これまでの処理でデータを入手できたので、モデルの学習データとして使えるように、テキストの特殊文字を消して、正規化し、数字ベクトルに変換します。


テキストの前処理


テキストの特殊文字を消し、正規化し、ローマ字を小文字に統一します。

def preprocess_text(text):
    text = re.sub(r'[\r\t\n\u3000]', '', text)
    text = neologdn.normalize(text)
    text = text.lower()
    text = text.strip()
    return text

data = data.assign(
    title=lambda x: x.title.map(lambda y: preprocess_text(y)),
    body=lambda x: x.body.map(lambda y: preprocess_text(y))
)
data.head()

f:id:dskomei:20210927215134p:plain:w600


データを学習用データと確認用データに分離し、ベクトル化


テキストを単語に分解し、単語をインデックス化させるトークナイザーを用意します。

tokenizer = T5Tokenizer.from_pretrained(settings.MODEL_NAME, is_fast=True)


次に、学習用データと確認用データに分割し、それぞれを数字ベクトル化します。

def convert_batch_data(train_data, valid_data, tokenizer):

    def generate_batch(data):

        batch_src, batch_tgt = [], []
        for src, tgt in data:
            batch_src.append(src)
            batch_tgt.append(tgt)

        batch_src = tokenizer(
            batch_src, max_length=settings.max_length_src, truncation=True, padding="max_length", return_tensors="pt"
        )
        batch_tgt = tokenizer(
            batch_tgt, max_length=settings.max_length_target, truncation=True, padding="max_length", return_tensors="pt"
        )

        return batch_src, batch_tgt

    train_iter = DataLoader(train_data, batch_size=settings.batch_size_train, shuffle=True, collate_fn=generate_batch)
    valid_iter = DataLoader(valid_data, batch_size=settings.batch_size_valid, shuffle=True, collate_fn=generate_batch)

    return train_iter, valid_iter

X_train, X_test, y_train, y_test = train_test_split(
    data['title'], data['body'], test_size=0.2, random_state=42, shuffle=True
)

train_data = [(src, tgt) for src, tgt in zip(X_train, y_train)]
valid_data = [(src, tgt) for src, tgt in zip(X_test, y_test)]

train_iter, valid_iter = convert_batch_data(train_data, valid_data, tokenizer)



ニュースタイトルから本文を生成するテキスト生成モデルの構築


ここまでの処理で学習データができたので、このデータを使ってタイトルから本文を作るテキスト生成モデルを作っていきます。


モデル定義クラス


学習済みのT5のモデルの「sonoisa/t5-base-japanese」をダウンロードし、Pytorch のディープラーニングのクラスに当てはめます。

class T5FineTuner(nn.Module):
    
    def __init__(self):
        super().__init__()

        self.model = T5ForConditionalGeneration.from_pretrained(settings.MODEL_NAME)

    def forward(
        self, input_ids, attention_mask=None, decoder_input_ids=None,
        decoder_attention_mask=None, labels=None
    ):
        return self.model(
            input_ids,
            attention_mask=attention_mask,
            decoder_input_ids=decoder_input_ids,
            decoder_attention_mask=decoder_attention_mask,
            labels=labels
        )



学習処理関数の定義


モデルを与えられたデータで学習する関数と、確認データで損失値を測定する関数を作ります。

def train(model, data, optimizer, PAD_IDX):
    
    model.train()
    
    loop = 1
    losses = 0
    pbar = tqdm(data)
    for src, tgt in pbar:
                
        optimizer.zero_grad()
        
        labels = tgt['input_ids'].to(settings.device)
        labels[labels[:, :] == PAD_IDX] = -100

        outputs = model(
            input_ids=src['input_ids'].to(settings.device),
            attention_mask=src['attention_mask'].to(settings.device),
            decoder_attention_mask=tgt['attention_mask'].to(settings.device),
            labels=labels
        )
        loss = outputs['loss']

        loss.backward()
        optimizer.step()
        losses += loss.item()
        
        pbar.set_postfix(loss=losses / loop)
        loop += 1
        
    return losses / len(data)

def evaluate(model, data, PAD_IDX):
    
    model.eval()
    losses = 0
    with torch.no_grad():
        for src, tgt in data:

            labels = tgt['input_ids'].to(settings.device)
            labels[labels[:, :] == PAD_IDX] = -100

            outputs = model(
                input_ids=src['input_ids'].to(settings.device),
                attention_mask=src['attention_mask'].to(settings.device),
                decoder_attention_mask=tgt['attention_mask'].to(settings.device),
                labels=labels
            )
            loss = outputs['loss']
            losses += loss.item()
        
    return losses / len(data)



モデルの学習


モデルの学習は、通常のディープラーニングの学習フローと同じです。エポックごとに先程定義した学習関数を叩いて、確認データの損失値を出します。また、学習時間の短縮のために「patience」で指定した回数損失値が改善されないと、学習のループから抜け出すようにしています。その際、「best_model」に格納されているのは、精度が最も良かったモデルです。

model = T5FineTuner()
model = model.to(settings.device)

optimizer = optim.Adam(model.parameters())

PAD_IDX = tokenizer.pad_token_id
best_loss = float('Inf')
best_model = None
counter = 1

for loop in range(1, settings.epochs + 1):

    start_time = time.time()

    loss_train = train(model=model, data=train_iter, optimizer=optimizer, PAD_IDX=PAD_IDX)

    elapsed_time = time.time() - start_time

    loss_valid = evaluate(model=model, data=valid_iter, PAD_IDX=PAD_IDX)

    print('[{}/{}] train loss: {:.4f}, valid loss: {:.4f} [{}{:.0f}s] counter: {} {}'.format(
        loop, settings.epochs, loss_train, loss_valid,
        str(int(math.floor(elapsed_time / 60))) + 'm' if math.floor(elapsed_time / 60) > 0 else '',
        elapsed_time % 60,
        counter,
        '**' if best_loss > loss_valid else ''
    ))

    if best_loss > loss_valid:
        best_loss = loss_valid
        best_model = copy.deepcopy(model)
        counter = 1
    else:
        if counter > settings.patience:
            break

        counter += 1



モデルの保存


今回使用したトークナイザーと学習済みモデルを保存します。

model_dir_path = Path('model')
if not model_dir_path.exists():
    model_dir_path.mkdir(parents=True)

tokenizer.save_pretrained(model_dir_path)
model.model.save_pretrained(model_dir_path)



学習したモデルを使ってニュース本文を生成


学習したモデルを使ってテキストを生成するために、テキストと学習済みモデルを与えて、生成テキストを返す関数を作ります。

def generate_text_from_model(title, trained_model, tokenizer, num_return_sequences=1):

    trained_model.eval()
    
    title = preprocess_text(title)
    batch = tokenizer(
        [title], max_length=settings.max_length_src, truncation=True, padding="longest", return_tensors="pt"
    )

    # 生成処理を行う
    outputs = trained_model.generate(
        input_ids=batch['input_ids'].to(settings.device),
        attention_mask=batch['attention_mask'].to(settings.device),
        max_length=settings.max_length_target,
        repetition_penalty=8.0,   # 同じ文の繰り返し(モード崩壊)へのペナルティ
        # temperature=1.0,  # 生成にランダム性を入れる温度パラメータ
        # num_beams=10,  # ビームサーチの探索幅
        # diversity_penalty=1.0,  # 生成結果の多様性を生み出すためのペナルティパラメータ
        # num_beam_groups=10,  # ビームサーチのグループ
        num_return_sequences=num_return_sequences,  # 生成する文の数
    )

    generated_texts = [
        tokenizer.decode(ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) for ids in outputs
    ]

    return generated_texts


この関数を使って、生成結果を確認します。

index = 0
title = valid_data[index][0]
body = valid_data[index][1]
generated_texts = generate_text_from_model(
    title=title, trained_model=trained_model, tokenizer=tokenizer, num_return_sequences=1
)
print('□ タイトル')
print(title)
print()
print('□ 生成本文')
print(generated_texts[0])
print()
print('□ 教師データ本文')
print(body)

f:id:dskomei:20210928083851p:plain:w600


生成したテキストは、教師テキストとはだいぶ異なっています。冒頭の「みなさん」から始まる文は、読者の気をひくのになかなか良い気がします。ただ
、ところどころ不自然な文言も混じってはいますが、人間が書いたと言われたらすんなり受け入れられるのではないかと思います。

index = 12
title = valid_data[index][0]
body = valid_data[index][1]
generated_texts = generate_text_from_model(
    title=title, trained_model=trained_model, tokenizer=tokenizer, num_return_sequences=1
)
print('□ タイトル')
print(title)
print()
print('□ 生成本文')
print(generated_texts[0])
print()
print('□ 教師データ本文')
print(body)

f:id:dskomei:20210928083927p:plain:w600


冒頭は教師テキストと同じですが、途中から異なる文になっています。箇条書きが箇条書きではなく文になってしまっているところがありますが、それぞの文としては意味が伝わります。AIが作ったテキストと知らずに見れば、すんなりと読んでしまいそうです。


終わりに


今回は T5 のモデルを構築しました。出来上がったモデルの結果を見てみると、多少の粗さはあるものの、人間と遜色ないレベル感であることはひしひしと感じました。Transformer 恐るべし!今回は、テキスト生成モデルだけに関して取り上げましたが、BERT を使ったテキスト分類や文章構成の方法に関しては以下の文献が参考になりました。