512 토큰 이상의 한국어 처리를 위한 BERT+Longformer Build

작성 동기

1. 한국어 데이터를 처리하기 위해 처음에 BERT 모델을 고려하였다.

2. BERT의 최대 토큰 수는 512토큰이고 내 데이터는 대부분 1000~3000 토큰 사이였다.

3. 따라서 BERT를 사용할 경우 512토큰 이내인 문장 혹은 문단 단위로 나눠 처리해야하는데 이 경우엔 전체 문맥을 온전히 고려할 수 없다.

4. 긴 문서 처리를 위해 BERT처럼 Transformer 기반의 모델로 설계되고 4096 토큰까지 처리할 수 있는 longformer 모델이 있다. (github 주소 : https://github.com/allenai/longformer)

5. 하지만 한국어 처리는 지원하지 않는다.

6. longformer 모델을 한국어로 사전학습 시켜야 하나 했지만 github에 기존 사전학습 모델을 longformer로 build 해주는 코드를 제공하였다 !(https://github.com/allenai/longformer/blob/master/scripts/convert_model_to_long.ipynb)

7. 위 코드를 응용하여 다국어(한국어 포함)를 지원하는 BERT를 longformer로 build하면 된다.

8. 필자는 정리된 게시글이 딱히 없어서 해당 깃허브 issue들을 뒤적거리며 해결하였고 그 내용을 정리하여 공유하고자 한다.

기본적으로 아래 github issue를 참고했다.

https://github.com/dcaled/convert_bert_to_long/blob/main/convert_bert_to_long.ipynb

가상환경 생성

conda create --name longformer python=3.7
conda activate longformer
git clone https://github.com/allenai/longformer.git
cd longformer
pip install -r requirements.txt

아래 부터는 jupyter notebook으로 실행한다.

%pip install transformers==3.0.2

모듈 import

BertModel과 BertTokenizerFast를 추가로 불러온다. 참고로 BERT 뿐만아니라 원하는 사전학습 모델을 longformer로 build할 수 있다.

import logging
import os
import math
import copy
import torch
from dataclasses import dataclass, field
from transformers import RobertaForMaskedLM, RobertaTokenizerFast, TextDataset, DataCollatorForLanguageModeling, Trainer
from transformers import TrainingArguments, HfArgumentParser
from transformers.modeling_longformer import LongformerSelfAttention
from transformers import BertModel, BertTokenizerFast
import tensorflow as tf
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

class 정의

longformer SelfAttention을 BERTModel에 Build하는 class 선언

class BertLongSelfAttention(LongformerSelfAttention):
    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        output_attentions=False,
    ):
        return super().forward(hidden_states, attention_mask=attention_mask, output_attentions=output_attentions)


class BertLong(BertModel):
    def __init__(self, config):
        super().__init__(config)
        for i, layer in enumerate(self.encoder.layer):
            # replace the `modeling_bert.BertSelfAttention` object with `LongformerSelfAttention`
            layer.attention.self = BertLongSelfAttention(config, layer_id=i)

self.encoder.layer 부분만 상속받은 모델과 맞춰주면 된다. 아래 코드부터도 동일하다.

for i, layer in enumerate(self.encoder.layer):

build function 정의

본인이 사용할 model과 tokenizer를 load 해준다.

다국어 BERT Model인 'bert-base-multilingual-cased'를 load 하였다.

def create_long_model(save_model_to, attention_window, max_pos):
    model = BertModel.from_pretrained('bert-base-multilingual-cased')
    tokenizer = BertTokenizerFast.from_pretrained('bert-base-multilingual-cased', model_max_length=max_pos)
    config = model.config

    print(max_pos)
    # extend position embeddings
    tokenizer.model_max_length = max_pos
    tokenizer.init_kwargs['model_max_length'] = max_pos
    current_max_pos, embed_size = model.embeddings.position_embeddings.weight.shape
    config.max_position_embeddings = max_pos

    assert max_pos > current_max_pos
    # allocate a larger position embedding matrix
    new_pos_embed = model.embeddings.position_embeddings.weight.new_empty(max_pos, embed_size)
    print(new_pos_embed.shape)
    print(model.embeddings.position_embeddings)
    # copy position embeddings over and over to initialize the new position embeddings
    k = 0
    step = current_max_pos
    while k < max_pos - 1:
        new_pos_embed[k:(k + step)] = model.embeddings.position_embeddings.weight
        k += step
    print(new_pos_embed.shape)
    model.embeddings.position_ids = torch.from_numpy(tf.range(new_pos_embed.shape[0], dtype=tf.int32).numpy()[tf.newaxis, :])
    model.embeddings.position_embeddings = torch.nn.Embedding.from_pretrained(new_pos_embed)

    # replace the `modeling_bert.BertSelfAttention` object with `LongformerSelfAttention`
    config.attention_window = [attention_window] * config.num_hidden_layers
    for i, layer in enumerate(model.encoder.layer):
        longformer_self_attn = LongformerSelfAttention(config, layer_id=i)
        longformer_self_attn.query = layer.attention.self.query
        longformer_self_attn.key = layer.attention.self.key
        longformer_self_attn.value = layer.attention.self.value

        longformer_self_attn.query_global = copy.deepcopy(layer.attention.self.query)
        longformer_self_attn.key_global = copy.deepcopy(layer.attention.self.key)
        longformer_self_attn.value_global = copy.deepcopy(layer.attention.self.value)

        layer.attention.self = longformer_self_attn
    print(model.embeddings.position_ids.shape)
    logger.info(f'saving model to {save_model_to}')
    model.save_pretrained(save_model_to)
    tokenizer.save_pretrained(save_model_to)
    return model, tokenizer

아래 코드에서 핵심은 max_pos (4096)이다.

나머지는 추가 사전교육을 위한 부분이다. 원하면 전체 코드를 참고하기 바란다.(https://github.com/allenai/longformer/blob/master/scripts/convert_model_to_long.ipynb)

@dataclass
class ModelArgs:
    attention_window: int = field(default=512, metadata={"help": "Size of attention window"})
    max_pos: int = field(default=4096, metadata={"help": "Maximum position"})

parser = HfArgumentParser((TrainingArguments, ModelArgs,))

training_args, model_args = parser.parse_args_into_dataclasses(look_for_args_file=False, args=[
    '--output_dir', 'tmp',
    '--warmup_steps', '500',
    '--learning_rate', '0.00003',
    '--weight_decay', '0.01',
    '--adam_epsilon', '1e-6',
    '--max_steps', '3000',
    '--logging_steps', '500',
    '--save_steps', '500',
    '--max_grad_norm', '5.0',
    '--per_gpu_eval_batch_size', '8',
    '--per_gpu_train_batch_size', '2',  # 32GB gpu with fp32
    '--gradient_accumulation_steps', '32',
    '--evaluate_during_training',
    '--do_train',
    '--do_eval',
])

# Choose GPU
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

Build

마지막 코드이다. 오류 없이 실행된다면 longformer build에 성공한 것이고, bert-base-4096이라는 폴더에 build 되었을 것이다.

model_path = f'{training_args.output_dir}/bert-base-{model_args.max_pos}'
if not os.path.exists(model_path):
    os.makedirs(model_path)

logger.info(f'Converting bert-base into bert-base-{model_args.max_pos}')
model, tokenizer = create_long_model(
    save_model_to=model_path, attention_window=model_args.attention_window, max_pos=model_args.max_pos)

Final

build한 longformer 모델을 사용하고자 할때 아래 코드처럼 불러오면 된다.

from transformers import BertModel, BertTokenizerFast
from transformers.modeling_longformer import LongformerSelfAttention


class BertLongSelfAttention(LongformerSelfAttention):
    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        output_attentions=False,
    ):
        return super().forward(hidden_states, attention_mask=attention_mask, output_attentions=output_attentions)


class BertLong(BertModel):
    def __init__(self, config):
        super().__init__(config)
        for i, layer in enumerate(self.encoder.layer):
            # replace the `modeling_bert.BertSelfAttention` object with `LongformerSelfAttention`
            layer.attention.self = BertLongSelfAttention(config, layer_id=i)

model = BertLong.from_pretrained('bert-base-4096/')
tokenizer = BertTokenizerFast.from_pretrained('bert-base-4096/')