KoBERT 쉽게 따라하고 간단한 fine-tuning 하기

"데이터셋에 대한 문의가 많습니다. 해당 데이터셋은 제가 프로젝트의 일환으로 "하이닥"이란 웹 사이트에서 크롤링으로 수집한 것입니다. 아시다시피 제 3자 데이터를 수집한 것을 다시 공유하는 것이 불가하기 때문에 데이터셋 공개는 어려울 것 같습니다. 다만, 데이터셋 형태를 본문에 첨부해놓았으니 참고해주세요. 이해해주셔서 감사합니다."

BERT를 처음부터 학습하기란 쉽지 않은 문제죠, 컴퓨터 사양도 좋아야 하고, 무엇보다 댓글 수집에도 많은 시간이 소요되니까요. 저는 한국어 데이터셋을 바탕으로 사전 학습된 모델인 KoBERT를 사용해보기로 했습니다. KoBERT는 SKT Brain에서 구축한 오픈 소스 라이브러리이며, 제가 알고 있는 한국어 BERT는 ETRI의 KorBERT, 이준범님의 KcBERT 정도가 있습니다. 각 모델마다 사전 학습 데이터셋이 다르기 때문에 구현하려는 목적에 따라 비교해보는 것도 의미가 있을 것 같습니다.

스스로 복습도 하고 여러분도 쉽게 따라할 수 있다는 자신감을 부여하기 위해 제가 KoBERT를 사용해본 경험을 소개하고자 합니다. 특히, 많은 오픈 소스가 이중분류(binary-classification)를 디폴트로 했기 때문에, 다중분류(multi-classification)를 하려면 어느 부분을 건드려야 하는지 모를 수 있는데요. 저는 일단 다중분류를 택했고, 참고해주시면 되겠습니다. 

동작 환경은 구글 코랩이며, SKT Brain의 깃허브에서 이미 코랩 환경 기반의 튜토리얼을 제공하고 있습니다. 따라서 저는 똑같은 코드는 작성하지 않고, 제가 바꾸거나 추가한 부분만 이야기하겠습니다. 

# 기존
# pip install transformers

# 변경: 최신 버전으로 설치하면 "Input: must be Tensor, not str" 라는 에러 발생
pip install transformers==3

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

import pandas as pd
dataset_train1 = pd.read_csv('/content/drive/My Drive/Colab Notebooks/nlp/dataset.csv')
dataset_train1.head() # 약 4500개의 증상 발화 데이터

여러 증상이 있지만, 데이터가 부족하다고 생각하여 아래와 같이 총 5개의 대표 질병만 선택해서 5중분류 모델을 구축했습니다.

data1 = dataset_train1.loc[dataset_train1['질병명'] == '수족냉증']
data2 = dataset_train1.loc[dataset_train1['질병명'] == '식중독']
data3 = dataset_train1.loc[dataset_train1['질병명'] == '소화불량']
data4 = dataset_train1.loc[dataset_train1['질병명'] == '질염']
data5 = dataset_train1.loc[dataset_train1['질병명'] == '비염']
new_data = data1.append([data2, data3, data4, data5], sort=False)
new_data = pd.DataFrame(new_data)
new_data.head()

from sklearn.preprocessing import LabelEncoder

encoder = LabelEncoder()
encoder.fit(new_data['질병명'])
new_data['질병명'] = encoder.transform(new_data['질병명'])
new_data.head()

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

from sklearn.model_selection import train_test_split
train, test = train_test_split(new_data, test_size=0.2, random_state=42) # new_df 오탈자 수정
print("train shape is:", len(train))
print("test shape is:", len(test))

## Setting parameters
max_len = 64 # 해당 길이를 초과하는 단어에 대해선 bert가 학습하지 않음
batch_size = 64
warmup_ratio = 0.1
num_epochs = 5
max_grad_norm = 1
log_interval = 200
learning_rate = 5e-5

# 데이터 형식 확인 
data_train[0]

class BERTClassifier(nn.Module):
    def __init__(self,
                 bert,
                 hidden_size = 768,
                 # 기존 -> num_classes = 2
                 num_classes = 5, # softmax 사용(다중분류)
                 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)

여기까지 진행하면 KoBERT 학습이 완료됩니다. validation 데이터셋에 대한 정확도는 약 0.84를 기록해서 이렇게만 보면 나쁘지 않은 성능인 듯합니다.

마지막으로 새로운 임의의 데이터를 생성하여 테스트를 해봤습니다. "음식, 발열, 구토, 복통, 설사"라고 사용자가 입력하면 인덱스 1(=소화불량)과 인덱스 3(=식중독)의 스코어가 가장 높게 나옵니다. 

unseen_test = pd.DataFrame([['음식, 발열, 구토, 복통, 설사', 7]], columns = [['질문 내용', '질병명']])
unseen_values = unseen_test.values
test_set = BERTDataset(unseen_values, 0, 1, tok, max_len, True, False)
test_input = torch.utils.data.DataLoader(test_set, batch_size=1, num_workers=5)

for batch_id, (token_ids, valid_length, segment_ids, label) in enumerate(tqdm_notebook(test_input)):
  token_ids = token_ids.long().to(device)
  segment_ids = segment_ids.long().to(device)
  valid_length= valid_length
  out = model(token_ids, valid_length, segment_ids)
  print(out)

튜토리얼은 본래 쉬워야 하는데, 딥러닝 분야는 튜토리얼 자체를 이해하는 것도 힘들더군요.ㅠ 본인의 데이터로 KoBERT를 사용해보고 싶은 분들에게 도움이 되었으면 좋겠습니다. 

---

마지막으로 제가 돌린 코드 전체를 궁금해하시는 분들이 있어 깃허브에 올려놓았으니 참고해주시길 바랍니다.

aeddung/ML-DL (github.com)

aeddung/ML-DL