BERT로 뉴스 기사 카테고리 분류하기

작업 인원 : 1명, 소요기간 : 1일, 작업환경 : google colab, python, KoBERT, 데이터 : 뉴스기사 텍스크와 해당 카테고리

정확도 0.847, F1 Score - 0.845

 

자연어처리에서 분류에 사용되는 모델은 간단한 RNN, LSTM부터 BERT , 등 다양한 모델이 있습니다. 이번에는 KoBERT를 사용해서 Fine-tuning해보았는데, KoBERT는 BERT에 한국어 데이터를 추가로 더 학습시킨 모델로 한국어에 더 좋은 성능을 보이는 것이 특징입니다. 아래 링크에서 모델을 볼 수 있습니다. 

https://github.com/SKTBrain/KoBERT

분류 카테고리는 총 6가지 입니다. ('경제', '사회', '생활', '스포츠', '연예/문화', '정치')

 

 

 

실행 전 코랩에서 런타임 유형을 GPU로 변경합니다.

1. 필요한 라이브러리 불러오기, 모델 다운로드

 !pip install mxnet
 !pip install gluonnlp pandas tqdm
 !pip install sentencepiece
 !pip install transformers==3.0.2
 !pip install torch
 
 # KoBERT 다운, 로드
!pip install git+https://git@github.com/SKTBrain/KoBERT.git@master

from kobert.utils import get_tokenizer
from kobert.pytorch_kobert import get_pytorch_kobert_model

from transformers import AdamW
from transformers.optimization import get_cosine_schedule_with_warmup''

import numpy as np
import random
import pandas as pd

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import gluonnlp as nlp

# 시드 고정
RANDOM_SEED = 2022
torch.manual_seed(RANDOM_SEED)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.random.seed(RANDOM_SEED)
random.seed(RANDOM_SEED)

#BERT 모델, Vocabulary 불러오기
bertmodel, vocab = get_pytorch_kobert_model()

 

2. 전처리하기 (한글과 영어만 남기고 모두 삭제했습니다.) 

import re
df['text'] = [re.sub('[^A-Za-z0-9가-힣]', '', s) for s in df['text']]

뉴스기사의 특성은 특수기호, 한자가 들어가는 경우가 있고, 마지막에 기자의 메일주소도 있는 경우가 있습니다. 정확도를 높이려면 이러한 노이즈 또한 정규표현식을 사용하여 일일이 제거해주면 좋을 것 같습니다. 

+ 데이터 증강으로 수집한 데이터를 각 카테고리별로 8000~8500 개로 맞췄습니다. 참고!

https://getacherryontop.tistory.com/109

 

3. 라벨 인코딩 

from sklearn.preprocessing import LabelEncoder

encoder = LabelEncoder()
encoder.fit(df['category'])
df['category'] = encoder.transform(df['category'])

mapping = dict(zip(range(len(encoder.classes_)), encoder.classes_))

문자인 카테고리를 숫자로 변경해줍니다. 

data_list = []
for q, label in zip(df['text'], df['category'])  :
    data = []
    data.append(q)
    data.append(str(label))
    data_list.append(data)

 

4. 파라미터 설정

Bert의 특성은 512토큰까지밖에 들어가지 않는다는 점입니다.(약 2~3문장 정도) 아무리 긴 텍스트를 넣어도 해당 길이밖에 학습이 되지 않습니다. 따라서 max_length로 앞에서부터 잘라 넣어줍니다. 배치사이즈는 아무래도 gpu의 할당량이 있어 최대한 작게 넣었습니다. 

max_len = 128
batch_size = 32

class CustomDataset(Dataset):

    def __init__(self, dataset, sent_idx, label_idx, bert_tokenizer, max_len,
                 pad, pair):
        transform = nlp.data.BERTSentenceTransform(
            bert_tokenizer, max_seq_length=max_len, pad=pad, pair=pair)
        
        self.sentences = [transform([i[sent_idx]]) for i in dataset]
        self.labels = [np.int32(i[label_idx]) for i in dataset]

    def __getitem__(self, i):
        return (self.sentences[i] + (self.labels[i], ))

    def __len__(self):
        return (len(self.labels))

__init__ : 초기화 ,  __len__ : 데이터셋(input)의 길이 반환, __getitem__ : 데이터셋을 인덱스로 불러옴

4. 검증 데이터셋 나누기

from sklearn.model_selection import train_test_split
                                                         
dataset_train, dataset_val = train_test_split(data_list, test_size=0.1, random_state=42)

train , 44066 개 , test 4897 개입니다. 

5. 토큰화

#토큰화
tokenizer = get_tokenizer()
tok = nlp.data.BERTSPTokenizer(tokenizer, vocab, lower=False)

data_train = CustomDataset(dataset_train, 0, 1, tok, max_len, True, False)
data_val = CustomDataset(dataset_val, 0, 1, tok, max_len, True, False)

train_dataloader = torch.utils.data.DataLoader(data_train, batch_size=batch_size, shuffle=True, num_workers=5)
valid_dataloader = torch.utils.data.DataLoader(data_val, batch_size=batch_size, shuffle=True, num_workers=5)

 

6. 하이퍼파라미터, 분류모델 설정

learning_rate =  2e-5
num_epochs = 4
log_interval = 200
warmup_ratio = 0.1

# 경로 설정
ROOT_PATH = '/content/drive/MyDrive'
DATA_DIR = '/content/drive/MyDrive'

class BertClassifier(nn.Module):
    def __init__(self,
                 bert,
                 hidden_size = 768,
                 num_classes=6,   ##클래스 수 조정##
                 dr_rate=None,
                 params=None):
        super(BertClassifier, self).__init__()
        self.bert = bert
        self.dr_rate = dr_rate
                 
        self.classifier = nn.Linear(hidden_size , num_classes)
        if dr_rate:
            self.dropout = nn.Dropout(p=dr_rate)
    
    def gen_attention_mask(self, token_ids, valid_length):
        attention_mask = torch.zeros_like(token_ids)
        for i, v in enumerate(valid_length):
            attention_mask[i][:v] = 1
        return attention_mask.float()

    def forward(self, token_ids, valid_length, segment_ids):
        attention_mask = self.gen_attention_mask(token_ids, valid_length)
        
        _, pooler = self.bert(input_ids = token_ids, token_type_ids = segment_ids.long(), attention_mask = attention_mask.float().to(token_ids.device))
        if self.dr_rate:
            out = self.dropout(pooler)
        return self.classifier(out)
    
    
    #BERT 모델 불러오기
model = BertClassifier(bertmodel,  dr_rate=0.5).to(device)

#optimizer와 schedule 설정
no_decay = ['bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
    {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
    {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
]

optimizer = AdamW(optimizer_grouped_parameters, lr=learning_rate)
loss_fn = nn.CrossEntropyLoss()

t_total = len(train_dataloader) * num_epochs
warmup_step = int(t_total * warmup_ratio)

scheduler = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=warmup_step, num_training_steps=t_total)

#정확도 측정을 위한 함수 정의
def calc_accuracy(X,Y):
    max_vals, max_indices = torch.max(X, 1)
    train_acc = (max_indices == Y).sum().data.cpu().numpy()/max_indices.size()[0]
    return train_acc
    
train_dataloader

 

 

 

 7. train

criterion = 0 #initialization
print('Training Start...')

for e in range(num_epochs):
    train_acc = 0.0
    val_acc = 0.0
    model.train()
    
    for batch_id, (token_ids, valid_length, segment_ids, label) in enumerate(tqdm(train_dataloader)):
        optimizer.zero_grad()
        token_ids = token_ids.long().to(device)
        segment_ids = segment_ids.long().to(device)
        valid_length= valid_length
        label = label.long().to(device)
        out = model(token_ids, valid_length, segment_ids)
        loss = loss_fn(out, label)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1)
        optimizer.step()
        scheduler.step()  # Update learning rate schedule
        train_acc += calc_accuracy(out, label)
        if batch_id % log_interval == 0:
            print("epoch {} batch id {} loss {} train acc {}".format(e+1, batch_id+1, loss.data.cpu().numpy(), train_acc / (batch_id+1)))
    print("epoch {} train loss {} train acc {}".format(e+1, loss.data.cpu().numpy(), train_acc / (batch_id+1)))
    
    model.eval()
    for batch_id, (token_ids, valid_length, segment_ids, label) in enumerate(tqdm(valid_dataloader)):
        token_ids = token_ids.long().to(device)
        segment_ids = segment_ids.long().to(device)
        valid_length= valid_length
        label = label.long().to(device)
        out = model(token_ids, valid_length, segment_ids)
        loss = loss_fn(out, label)
        val_acc += calc_accuracy(out, label)
        
    print("epoch {} valid loss {} valid acc {}".format(e+1, loss.data.cpu().numpy(), val_acc / (batch_id+1)))
    if val_acc / (batch_id+1) > criterion:
        # 모델이 개선됨 -> 검증 점수와 weight 갱신
        criterion = val_acc / (batch_id+1)
        check_point = {
            'model': model.state_dict(),
            'optimizer': optimizer.state_dict(),
            'scheduler': scheduler.state_dict()
        }
        # 모델이 개선됨 -> 모델 저장
        print("Update model save...")
        torch.save(check_point, os.path.join(ROOT_PATH, 'model_best.pt'))
        
print('Training Finish!')

 

epoch 4 valid loss 0.03084656223654747 valid acc 0.8474025974025974

8. 학습된 모델 적용

class TestDataset(Dataset):
    def __init__(self, dataset, sent_idx, bert_tokenizer, max_len,
                 pad, pair):
        transform = nlp.data.BERTSentenceTransform(
            bert_tokenizer, max_seq_length=max_len, pad=pad, pair=pair)
        self.categories = mapping
        self.num_labels = len(self.categories)
        
        self.sentences = [transform([i]) for i in dataset]
        
    def label_decoder(self, labels):
        try:
            labels = list(map(lambda x : self.categories[x], labels))
            return labels
        except:
            assert 'Invalid intent'
            
    def __len__(self):
        return (len(self.sentences))
    
    def __getitem__(self, i):
        return (self.sentences[i])
 
 dataset_test = testdf['text'].tolist()
data_test = TestDataset(dataset_test, 0, tok, max_len, True, False)
test_dataloader = torch.utils.data.DataLoader(data_test, batch_size=batch_size, num_workers=5)

#MODEL_DIR = os.path.join(ROOT_PATH, 'model_best.pt')
#model = BertClassifier(bertmodel,  dr_rate=0.5).to(device)
#model.load_state_dict(torch.load(MODEL_DIR)['model'])

model.eval()
pred = []

for batch_id, (token_ids, valid_length, segment_ids) in enumerate(test_dataloader):
    token_ids = token_ids.long().to(device)
    segment_ids = segment_ids.long().to(device)
    valid_length = valid_length.long().to(device)

    out = model(token_ids, valid_length, segment_ids)
    pred.extend(out.argmax(dim=1).tolist())
    
    # 진행과정 출력
    if batch_id % 10 == 0:
        print(f'Prediction: {batch_id}/{len(test_dataloader)} completed')

pred = data_test.label_decoder(pred)
print('decode Completed!')

결과

 

나름 잘 나왔음을 볼 수 있습니다.