pytorch로 BERT 구현하기 - 토크나이즈

이전글에서는 BERT의 이론을 정리해보았다.

이번글에서는 BERT 학습 이전에 토크나이즈를 정리할 예정이다.

wget https://dumps.wikimedia.org/kowiki/latest/kowiki-latest-pages-articles.xml.bz2
pip install wikiextractor
python -m wikiextractor.WikiExtractor kowiki-latest-pages-articles.xml.bz2

위 스크립트를 실행하면 위키데이터를 다운로드 받아 아래와 같은 형태의 데이터로 정리해 준다

그중 하나의 파일을 열어보면

위 사진과 같이 xml 코드를 잘 분해해서 텍스트로 저장이 된다.

def load_text_data(path):
    doc_pattern = re.compile("^<doc id|</doc>")
    data = []
    for folder in sorted(os.listdir(path)):
        for file_name in sorted(os.listdir(path + f"/{folder}")):
            file_path = path + f"/{folder}/{file_name}"
            with open(file_path, "r") as f:
                while True:
                    line = f.readline()
                    if not line: break
                    if bool(doc_pattern.search(line)): continue
                    line = line.replace("\n", "")
                    if line:
                        data.append(line)
    return data
 
 def save_text_data(data):
    mecab = Mecab()
    with open("kowiki.txt", "w") as f:
        for line in tqdm(data):
            f.write(" ".join(mecab.morphs(line)) + "\n")

data = load_text_data("text")
save_text_data(data)

위 함수는 text/AA, text/AB 등과 같이 정리되어 있는 텍스트 파일들을 하나로 합쳐준다. 또한 토크나이즈로 학습하기 이전에 영어는 단순히 split 으로 단어를 분리할 수 있지만 한글은 부사와 같은 단어들이 합쳐져 있기 때문에 분리하기 위해 konlpy의 mecab을 사용하여 저장한다.

from tokenizers import BertWordPieceTokenizer

def train_bert_tokenizer(corpus, min_frequency, limit_alphabet, vocab_size=32000):
    tokenizer = BertWordPieceTokenizer(strip_accents=False, lowercase=False)
    tokenizer.train(files=[corpus], vocab_size=vocab_size, min_frequency=min_frequency, limit_alphabet=limit_alphabet, show_progress=True)

    tokenizer.save_model("tokenize")

train_bert_tokenizer("kowiki.txt", 5, 6000)

위 코드를 실행하면 BERT tokenize 형태로 학습한 vocab.txt 파일을 생성한다.