BERT 모델 MLM 기반 pre-train 파이썬 코드 예제

Transformers 라이브러리 BERT corpus pre-train 실습

BERT 모델을 사전 학습하는 방법인 MLM(Masked Language Modeling) 기법을 적용하여

원하는 corpus로 학습된 pre-trained BERT를 생성하는 파이썬 코드 예시를 살펴보겠습니다.

 

Transformers 라이브러리의 기능을 사용하여 진행한 예제이며,

random word 토큰 선별 과정을 제외하고는 아래 사이트의 코드를 참조하였습니다.

Masked-Language Modelling With BERT

 

 

Step 1 : BERT, tokenizer 불러오기

가장 먼저, Transformers 라이브러리 내에서 제공하고 있는 사전 학습된

BERT 모델의 파라미터 정보와 토크나이저를 불러와보도록 하겠습니다.

from transformers import BertTokenizer, BertForMaskedLM
# 패키지 미설치의 경우 : !pip install transformers로 설치

bert_tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForMaskedLM.from_pretrained('bert-base-uncased')

가장 기본적인 BERT 모델을 불러오려는 경우에는 bert-base-uncased로 input을 지정하시고,

 

다른 종류의 사전 학습이 완료된 모델을 가져오고 싶다면 HuggingFace/Models 페이지에서

원하는 모델의 이름을 목록에서 찾아 input으로 지정해주시면 됩니다.

Models - Hugging Face

 

 

Step 2 : input text 토큰화, label 지정

토큰 512개 이하 길이의 텍스트(100~200 단어를 많이 사용하는 듯 합니다.)에 대하여 

토큰화를 진행해주고, 아래와 같이 결과를 확인해보겠습니다.

(단, 여기서는 padding은 생략하였습니다.)

text = """
Squid Game was released worldwide on September 17, 2021, to critical acclaim and international attention. It is Netflix's most-watched series, becoming the top-viewed program in 94 countries and attracting more than 142 million member households and amassing 1.65 billion viewing hours during its first four weeks from launch, surpassing Bridgerton for the title of most watched show. The series has also received numerous accolades, including the Golden Globe Award for Best Supporting Actor – Series, Miniseries or Television Film for O Yeong-su and the Screen Actors Guild Award for Outstanding Performance by a Male Actor in a Drama Series and Outstanding Performance by a Female Actor in a Drama Series for Lee Jung-jae and HoYeon Jung, respectively, with all three making history as the first Korean actors to win in those categories. A second season is in development.
"""

inputs = bert_tokenizer(text, return_tensors='pt')

### max length에 대한 padding을 포함한 경우 ###
# inputs = bert_tokenizer(text, return_tensors='pt', max_length=512, truncation=True, padding='max_length')

print(inputs)

 

MLM 학습 방법에서는 마스킹된 기존 단어를 그대로 맞추는 것이 목표이기에,

labels는 기존 토큰을 그대로 복사하여 가져와주시면 됩니다.

inputs['labels'] = inputs.input_ids.detach().clone()

반응형

Step 3 : 마스킹 토큰 위치 선별 및 토큰 대체

BERT 논문에 따르면, MLM 기법 학습 시 전체 토큰 중 15%의 위치에 대하여

토큰 대체 후보를 정해준다고 합니다.(80%는 Mask, 10%는 랜덤, 10%는 유지)

 

torch.rand 함수를 이용하여 15%의 랜덤 위치들을 선별하겠습니다.

단, 특수 토큰인 CLS(101번)와 SEP(102번) 토큰은 제외하겠습니다.

(padding이 포함된 경우 패딩을 나타내는 0번 토큰도 제외해주세요.)

import torch

rand = torch.rand(inputs.input_ids.shape)

mask_arr = (rand < 0.15) * (inputs.input_ids != 101) * (inputs.input_ids != 102) # 101, 102번 토큰 제외하고 15% 위치 선별

### padding이 포함된 경우(0번도 추가 제외) ###
# mask_arr = (rand < 0.15) * (inputs.input_ids != 101) * (inputs.input_ids != 102) * (inputs.input_ids != 0)

selection = torch.flatten((mask_arr[0]).nonzero())

print(selection) # 선별된 위치의 인덱스 번호
"""tensor([  7,  24,  30,  32,  52,  53,  54,  62,  64,  74,  85,  91,  95,  97,
         98, 100, 103, 112, 113, 136, 141, 143, 153, 159])"""

 

선택된 15%의 위치 중 80%는 Mask 토큰으로의 대체(103번), 10%는 랜덤 토큰으로 대체,

10%는 원래 토큰 그대로 유지의 기법을 따르는 예제 코드를 보겠습니다.

 

해당 위치 선별 방법은 넘파이의 where 함수를 사용하여 진행해보겠습니다.

import numpy as np

selection_val = np.random.random(len(selection)) # selection의 위치마다 0~1 값 부여

mask_selection = selection[np.where(selection_val >= 0.2)[0]] # 80% : Mask 토큰 대체
random_selection = selection[np.where(selection_val < 0.1)[0]] # 10% : 랜덤 토큰 대체

print(random_selection) # tensor([ 30,  95, 143])
print(mask_selection)
"""tensor([  7,  24,  32,  52,  53,  54,  62,  64,  74,  85,  91,  97,  98, 100,
        103, 112, 113, 136, 141, 153, 159])"""

 

이제 위에서 고른 input_id의 위치의 토큰 번호를 알맞게 대체해주시면 됩니다.

Mask는 103번이고, bert-base 모델에서의 전체 토큰의 가짓수는 30,522개이므로

아래와 같이 토큰들을 바꿔보도록 하겠습니다.

inputs.input_ids[0, mask_selection] = 103
inputs.input_ids[0, random_selection] = torch.randint(0, 30522, size = random_selection.shape)

print(inputs)

 

 

Step 4 : BERT 모델 학습

위에서 만든 input을 BERT 모델에 넣은 결과 loss는 아래와 같이 간단히 확인이 가능합니다.

outputs = model(**inputs)

outputs.loss # tensor(0.8096, grad_fn=<NllLossBackward0>)

 

글 첫 부분에서 참고한 towardsdatascience 사이트에 나와있는

BERT를 실제로 학습시키는 과정에 대한 예시 코드는 아래와 같았습니다.

# 출처 : https://towardsdatascience.com/masked-language-modelling-with-bert-7d49793e5d2c

# 데이터 셋 클래스
class MeditationsDataset(torch.utils.data.Dataset):
    def __init__(self, encodings):
        self.encodings = encodings
    def __getitem__(self, idx):
        return {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
    def __len__(self):
        return len(self.encodings.input_ids)

from transformers import TrainingArguments
from transformers import Trainer

# 위에서 설정한 inputs를 데이터 셋 클래스로 지정
dataset = MeditationsDataset(inputs)

# argument 설정 예시
args = TrainingArguments(
    output_dir='result', # 결과 정보를 받아볼 디렉토리
    per_device_train_batch_size=16, # gpu 당 batch size 수
    num_train_epochs=2 # epoch 수
)

trainer = Trainer(
    model=model, # BERT 모델이 저장된 변수
    args=args, # 지정한 argument
    train_dataset=dataset # 데이터셋
)

trainer.train() # TrainOutput(global_step=2, training_loss=0.6562116742134094, metrics={'train_runtime': 8.4303, 'train_samples_per_second': 0.237, 'train_steps_per_second': 0.237, 'total_flos': 173756302200.0, 'train_loss': 0.6562116742134094, 'epoch': 2.0})

 

출처 :

James Briggs 님의 towardsdatascience 사이트의 글(맨 윗 부분 링크) 및

해당 글에 첨부된 github gist 사이트의 코드들(별도 주석 추가 및 일부 변경하여 포스팅)