【No.L006】日本語ニュース分類モデルの構築 – Livedoor ニュースコーパスを用いたBERTファインチューニング

!pip install fugashi ipadic

import tarfile

tar = tarfile.open("./ldcc-20140209.tar.gz")
tar.extractall("./data/livedoor/")
tar.close()

import os

cate_list = [name for name in os.listdir("./data/livedoor/text/") if os.path.isdir(os.path.join("./data/livedoor/text/", name))]

from transformers import BertJapaneseTokenizer, BertModel
import fugashi
import ipadic
import glob
import torch

n_label = 2
max_length = 10
cate_list = cate_list[:n_label]
tokenizer = BertJapaneseTokenizer.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')

dataset = []
for label, cate in enumerate(cate_list):
    for f in glob.glob(f'./data/livedoor/text/{cate}/{cate}*'):
        lines = open(f, encoding="utf-8_sig").read().splitlines()
        text = '\n'.join(lines[2:])
        encoding = tokenizer(text, max_length=max_length, padding='max_length', truncation=True)
        encoding['label'] = label
        encoding = {k: torch.tensor(v) for k, v in encoding.items()}
        dataset.append(encoding)

import random
from torch.utils.data import DataLoader

dataset_for_loader = dataset.copy()
random.seed(0)
random.shuffle(dataset_for_loader)
n = len(dataset_for_loader)
n_train, n_valid = int(n*0.6), int(n*0.2)
dataset_train = dataset_for_loader[:n_train]
dataset_valid = dataset_for_loader[n_train:n_train+n_valid]
dataset_test = dataset_for_loader[n_train+n_valid:]

batch_size = 16
dataloader_train = DataLoader(dataset_train, batch_size = batch_size, shuffle = True)
dataloader_valid = DataLoader(dataset_valid, batch_size = batch_size, shuffle = False)
dataloader_test = DataLoader(dataset_test, batch_size = batch_size, shuffle = False)

import torch.nn as nn

class BertClassifier(nn.Module):
    def __init__(self):
        super(BertClassifier, self).__init__()
        self.model_bert = model_bert
        self.linear = nn.Linear(in_features=768, out_features=n_label)
        nn.init.normal_(self.linear.weight, std=0.02)
        nn.init.normal_(self.linear.bias, 0)
        
    def forward(self, input_ids):
        output_bert = self.model_bert(input_ids)
        output_bert = output_bert[0]
        output_bert = output_bert[:, 0, :].view(-1, 768) # [batch_size, hidden_size]に変換
        output = self.linear(output_bert)
        return output

import torch.optim as optim

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model_bert = BertModel.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')
model = BertClassifier()
model.to(device)

for param in model.parameters():
    param.requires_grad = False
for param in model.model_bert.encoder.layer[-1].parameters():
    param.requires_grad = True
for param in model.linear.parameters():
    param.requires_grad = True   


optimizer = optim.Adam([
    {'params': model.model_bert.encoder.layer[-1].parameters(), 'lr': 5e-5},
    {'params': model.linear.parameters(), 'lr': 1e-4}
])
criterion = nn.CrossEntropyLoss()

n_epochs = 5
for epoch in range(n_epochs):
    loss_train = 0
    loss_valid = 0
    accuracy_train = 0
    accuracy_valid = 0
    
    model.train()
    for batch in dataloader_train:
        x = batch['input_ids'].to(device)
        t = batch['label'].to(device)
        batch_size_iter = len(x)
        model.zero_grad()
        optimizer.zero_grad()
        
        output = model(x)
        loss = criterion(output, t)
        loss.backward()
        optimizer.step()
        pred = output.argmax(dim=1)
        
        loss_train += loss.item() * batch_size_iter
        accuracy_train += torch.sum(pred==t.data)
    
    model.eval()
    for batch in dataloader_valid:
        x = batch['input_ids'].to(device)
        t = batch['label'].to(device)
        batch_size_iter = len(x)
        
        output = model(x)
        loss = criterion(output, t)
        pred = output.argmax(dim=1)
        
        loss_valid += loss.item() * batch_size_iter
        accuracy_valid += torch.sum(pred == t.data)
    
    average_loss_train = loss_train / len(dataset_train)
    average_loss_valid = loss_valid / len(dataset_valid)
    total_accuracy_train = accuracy_train / len(dataset_train)
    total_accuracy_valid = accuracy_valid / len(dataset_valid)
    print('| epoch %d | train loss %.2f, train accuracy %.2f | valid loss %.2f, valid accuracy %.2f'
          % (epoch+1, average_loss_train, total_accuracy_train, average_loss_valid, total_accuracy_valid))