[🤗 강좌 6.7] WordPiece 토큰화

WordPiece는 Google이 BERT를 사전 학습하기 위해 개발한 토큰화 알고리즘입니다. 그 이후로 DitilBERT, MobileBERT, Funnel Transformers 및 MPNET과 같은 BERT 기반의 상당히 많은 Transformer 모델에서 재사용되었습니다. 학습 측면에서 BPE와 매우 유사하지만 실제 토큰화는 다르게 수행됩니다.

💡 이 섹션에서는 WordPiece를 심층적으로 다루며 전체 구현 과정을 보여줍니다. 토큰화 알고리즘에 대한 일반적인 개요를 원하는 경우 생략해도 좋습니다.

학습 알고리즘

⚠️ Google은 WordPiece의 학습 알고리즘 구현을 오픈 소스로 공개하지 않았으므로 이번 섹션에서는 발표된 논문 내용을 기반으로 구현 과정을 설명합니다. 거의 100% 정확할 것입니다.

BPE와 마찬가지로 WordPiece는 모델에서 사용하는 특수 토큰과 초기 알파벳을 포함한 작은 vocabulary에서 시작합니다. 접두사(예: BERT의 ##)를 추가하여 하위 단어(subwords)를 식별하기 때문에 각 단어는 처음에 해당 접두사를 단어 내부의 모든 문자에 추가하여 분할됩니다. 예를 들어 "word"는 다음과 같이 분할됩니다.

w ##o ##r ##d

따라서 초기 알파벳에는 단어의 시작 부분에 있는 모든 문자들(예: 'w')과 WordPiece 접두사가 선행하는 단어 내부에 있는 문자(예: 'o', 'r', 'd')가 포함됩니다.

그런 다음 BPE와 마찬가지로 WordPiece도 병합 규칙을 학습합니다. 주요 차이점은 병합할 쌍이 선택되는 방식입니다. 가장 빈번하게 출현하는 쌍을 선택하는 대신 WordPiece는 다음 공식을 사용하여 각 쌍에 대한 점수를 계산합니다:

$score = \frac{freq\_of\_pair}{freq\_of\_first\_element \times freq\_of\_second\_element}$

쌍의 빈도를 각 부분의 빈도의 곱으로 나눔으로써, 알고리즘은 각 개별 부분들의 빈도가 낮은 쌍의 병합에 높은 우선순위를 부여합니다. 예를 들어, vocabulary 내에서의 출현 빈도가 높은 ("un", "##able") 쌍을 굳이 병합할 필요는 없는데, 그 이유는 "un"과 "##able" 각각이 다른 단어 내에서 매우 빈번하게 출현하여 높은 빈도를 나타내기 때문입니다. 반면에, "hu"와 "##gging"은 각각이 자주 사용되지 않기 때문에 ("hu", "##gging")과 같은 쌍은 아마도 더 빨리 병합될 것입니다("hugging"이라는 단어가 어휘에 자주 등장한다고 가정).

앞서 BPE 학습 예시에서 사용한 것과 동일한 vocabulary를 사용해서 설명하겠습니다:

("hug", 10), ("pug", 5), ("pun", 12), ("bun", 4), ("hugs", 5)

분할 결과는 다음과 같습니다:

("h" "##u" "##g", 10), ("p" "##u" "##g", 5), ("p" "##u" "##n", 12), ("b" "##u" "##n", 4), ("h" "##u" "##g" "##s", 5)

따라서 초기 vocabulary는 ["b", "h", "p", "##g", "##n", "##s", "##u"]가 됩니다(특수 토큰은 일단 잊어버립시다). 가장 빈번한 쌍은 ("##u", "##g")(현재 20회)이지만 "##u"의 개별 빈도가 매우 높아 점수가 가장 높지는 않습니다(1/36). "##u"가 포함된 모든 쌍은 실제로 동일한 점수(1/36)를 가지므로 가장 좋은 점수는 ("##g", "##s")이 가지고 있습니다(1/20). 이는 "##u"가 없는 유일한 쌍입니다. 그리고 학습된 첫 번째 병합은 ("##g", "##s") -> ("##gs")입니다.

병합할 때 두 토큰 사이의 ##을 제거하므로 vocabulary에 "##gs"를 추가하고 말뭉치의 모든 단어에 해당 병합을 적용합니다:

Vocabulary: ["b", "h", "p", "##g", "##n", "##s", "##u", "##gs"]
Corpus: ("h" "##u" "##g", 10), ("p" "##u" "##g", 5), ("p" "##u" "##n", 12), ("b" "##u" "##n", 4), ("h" "##u" "##gs", 5)

이 시점에서 "##u"는 가능한 모든 쌍에 있으므로 모두 동일한 점수를 가집니다. 이 경우 첫 번째 쌍이 병합되므로 ("h", "##u") -> "hu"가 학습됩니다. 그 결과는 다음과 같습니다:

Vocabulary: ["b", "h", "p", "##g", "##n", "##s", "##u", "##gs", "hu"]
Corpus: ("hu" "##g", 10), ("p" "##u" "##g", 5), ("p" "##u" "##n", 12), ("b" "##u" "##n", 4), ("hu" "##gs", 5)

이제 최고 점수는 ("hu", "##g") 및 ("hu", "##gs")가 동일하게 계산되므로(다른 모든 쌍의 경우 1/21이고 이 두 쌍은 1/15) 가장 큰 점수를 가진 첫 번째 쌍이 병합됩니다.

Vocabulary: ["b", "h", "p", "##g", "##n", "##s", "##u", "##gs", "hu", "hug"]
Corpus: ("hug", 10), ("p" "##u" "##g", 5), ("p" "##u" "##n", 12), ("b" "##u" "##n", 4), ("hu" "##gs", 5)

원하는 어휘 크기에 도달할 때까지 이 단계가 계속 반복됩니다.

✏️ Now your turn! 다음 병합 규칙은 무엇일까요?

토큰화 알고리즘

토큰화는 WordPiece가 학습된 병합 규칙은 제외하고 최종 vocabulary만 저장한다는 점에서 BPE와는 다릅니다. 토큰화할 단어에서 시작하여 WordPiece는 vocabulary에 있는 가장 긴 하위 단어를 찾은 다음 분할합니다. 예를 들어, 위의 예에서 학습한 vocabulary를 사용하는 경우, 단어 "hugs"의 경우 처음부터 시작하는 가장 긴 하위 단어는 vocabulary 내부에 있는 "hug"이므로 거기에서 분할하여 ["hug", "##s"]로 분할됩니다. 그런 다음 "##s"가 vocabulary에 존재하고 이를 계속 사용할 수 있으므로 "hugs"의 토큰화 결과는 ["hug", "##s"]입니다.

BPE를 사용하면 학습된 병합(merges)을 순서대로 적용하고 이를 ["hu", "##gs"]로 토큰화하므로 인코딩이 다르게 됩니다.

또 다른 예로 "bugs"라는 단어가 어떻게 토큰화되는지 봅시다. "b"는 vocabulary에 존재하는 단어의 시작 부분에서 시작하는 가장 긴 하위 단어이므로 거기서 분할하여 ["b", "##ugs"]라는 중간 결과가 도출됩니다. 다음으로 "##u"는 vocabulary에 있는 "##ugs"의 시작 부분에서 시작하는 가장 긴 하위 단어이므로 거기에서 분할하여 ["b", "##u, "##gs"]를 얻습니다. 마지막으로 "##gs"가 vocabulary에 있으므로 ["b", "##u, "##gs"]이 "bugs"의 토큰화 결과입니다.

토큰화가 vocabulary에서 하위 단어(subword)를 더이상 찾을 수 없는 단계에 도달하면 전체 단어를 "unknown"으로 토큰화합니다. 예를 들어 "mug"는 "bum"과 마찬가지로 ["[UNK]"]로 토큰화됩니다("b"와 "##u"로 시작할 수 있더라도 "##m"이 vocabulary에 존재하지 않으므로 결과 토큰화는 ["b", "##u", "[UNK]"]가 아니라 ["[UNK]"]입니다). 이것은 vocabulary에 없는 개별 문자만 "unknwon"으로 분류하는 BPE와의 또 다른 차이점입니다.

✏️ Now your turn! "pugs"라는 단어는 어떻게 토큰화될까요?

WordPiece 구현

이제 WordPiece 알고리즘의 구현을 살펴보겠습니다. BPE와 마찬가지로 아래 코드는 이해를 돕기 위해서 구현한 것이며 대규모 말뭉치에서는 사용할 수 없습니다.

우리는 BPE 예시에서와 동일한 말뭉치를 사용할 것입니다:

corpus = [
    "This is the Hugging Face course.",
    "This chapter is about tokenization.",
    "This section shows several tokenizer algorithms.",
    "Hopefully, you will be able to understand how they are trained and generate tokens.",
]

먼저 말뭉치를 단어로 사전 토큰화(pre-tokenization)해야 합니다. BERT와 같은 WordPiece 토크나이저를 구현하고 있으므로 사전 토큰화에 bert-base-cased 토크나이저를 사용합니다:

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")

그런 다음 사전 토큰화 수행 과정에서 말뭉치에 있는 각 단어의 빈도를 계산합니다:

from collections import defaultdict

word_freqs = defaultdict(int)
for text in corpus:
    words_with_offsets = tokenizer.backend_tokenizer.pre_tokenizer.pre_tokenize_str(text)
    new_words = [word for word, offset in words_with_offsets]
    for word in new_words:
        word_freqs[word] += 1
        
word_freqs

이전에 보았듯이 알파벳은 단어의 모든 첫 글자와 ## 접두사가 붙은 단어에 나타나는 다른 모든 글자로 구성된 고유한 집합입니다:

alphabet = []
for word in word_freqs.keys():
    if word[0] not in alphabet:
        alphabet.append(word[0])
    for letter in word[1:]:
        if f"##{letter}" not in alphabet:
            alphabet.append(f"##{letter}")
            
alphabet.sort()
alphabet

print(alphabet)

또한 해당 vocabulary의 시작 부분에 모델이 사용하는 특수 토큰을 추가합니다. BERT의 경우 ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]입니다:

vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"] + alphabet.copy()

다음으로 vocabulary에 존재하는 접두사가 ##이 아닌 모든 문자를 사용하여 각 단어를 분할해야 합니다:

splits = {
    word: [c if i == 0 else f"##{c}" for i, c in enumerate(word)]
    for word in word_freqs.keys()
}

이제 학습할 준비가 되었으므로 각 쌍의 점수를 계산하는 함수를 작성해 보겠습니다. 학습의 각 단계에서 이 함수를 사용해야 합니다:

def compute_pair_scores(splits):
    letter_freqs = defaultdict(int)
    pair_freqs = defaultdict(int)
    for word, freq in word_freqs.items():
        split = splits[word]
        if len(split) == 1:
            letter_freqs[split[0]] += freq
            continue
        for i in range(len(split) - 1):
            pair = (split[i], split[i + 1])
            letter_freqs[split[i]] += freq
            pair_freqs[pair] += freq
        letter_freqs[split[-1]] += freq
        
    scores = {
        pair: freq / (letter_freqs[pair[0]] * letter_freqs[pair[1]])
        for pair, freq in pair_freqs.items()
    }
    return scores

초기 분할 후 pair_scores의 일부를 살펴보겠습니다:

pair_scores = compute_pair_scores(splits)
for i, key in enumerate(pair_scores.keys()):
    print(f"{key}: {pair_scores[key]}")
    if i >= 5:
        break

이제 최고의 점수를 가진 쌍을 찾는 간단한 루프를 구현합니다:

best_pair = ""
max_score = None
for pair, score in pair_scores.items():
    if max_score is None or max_score < score:
        best_pair = pair
        max_score = score

print(best_pair, max_score)

따라서 학습할 첫 번째 병합은 ('a', '##b') -> 'ab'이고 vocabulary에 'ab'를 추가합니다:

vocab.append("ab")

계속하려면 splits 딕셔너리에 해당 병합을 적용해야 합니다. 이를 위해 다른 함수를 작성해 보겠습니다:

def merge_pair(a, b, splits):
    for word in word_freqs:
        split = splits[word]
        if len(split) == 1:
            continue
        i = 0
        while i < len(split) - 1:
            if split[i] == a and split[i + 1] == b:
                merge = a + b[2:] if b.startswith("##") else a + b
                split = split[:i] + [merge] + split[i + 2 :]
            else:
                i += 1
        splits[word] = split
    return splits

이제 첫 번째 병합의 결과를 볼 수 있습니다:

splits = merge_pair("a", "##b", splits)
splits["about"]

이제 원하는 모든 병합을 모두 학습할때 까지 반복하는데 필요한 모든 것을 구현했습니다. 목표 vocabulary 크기를 70으로 합시다:

vocab_size = 70
while len(vocab) < vocab_size:
    scores = compute_pair_scores(splits)
    best_pair, max_score = "", None
    for pair, score in scores.items():
        if max_score is None or max_score < score:
            best_pair = pair
            max_score = score
    splits = merge_pair(*best_pair, splits)
    new_token = (
        best_pair[0] + best_pair[1][2:]
        if best_pair[1].startswith("##")
        else best_pair[0] + best_pair[1]
    )
    vocab.append(new_token)

생성된 vocabulary를 볼 수 있습니다:

print(vocab)

보시다시피 BPE에 비해 이 토크나이저는 단어의 일부를 토큰으로 더 빨리 학습합니다.

💡 동일한 말뭉치에서 train_new_from_iterator()를 사용하면 똑같은 vocabulary가 나오지 않습니다. 🤗Tokenizers 라이브러리는 학습을 위해 WordPiece를 구현하지 않고(내부에 대해 완전히 확신하지 못하기 때문에) 대신 BPE를 사용하기 때문입니다.

새로운 텍스트를 토큰화하기 위해 사전 토큰화하고(pre-tokenization), 분할한 다음(split), 각 단어에 토큰화 알고리즘을 적용합니다. 즉, 첫번째 단어의 시작 부분에서 시작하는 가장 큰 하위 단어를 찾아 분할한 다음, 두번째 부분에서 프로세스를 반복하고 나머지 단어와 텍스트의 다음 단어에 대해 계속 반복합니다:

def encode_word(word):
    tokens = []
    while len(word) > 0:
        i = len(word)
        while i > 0 and word[:i] not in vocab:
            i -= 1
        if i == 0:
            return ["[UNK]"]
        tokens.append(word[:i])
        word = word[i:]
        if len(word) > 0:
            word = f"##{word}"
    return tokens

Vocabulary에 존재하는 단어와 그렇지 않은 단어에 대해 테스트해 보겠습니다:

print(encode_word("Hugging"))
print(encode_word("HOgging"))

이제 텍스트를 토큰화하는 함수를 작성해 보겠습니다:

def tokenize(text):
    pre_tokenize_result = tokenizer._tokenizer.pre_tokenizer.pre_tokenize_str(text)
    pre_tokenized_text = [word for word, offset in pre_tokenize_result]
    encoded_words = [encode_word(word) for word in pre_tokenized_text]
    return sum(encoded_words, [])

이제 어떤 텍스트를 가지고도 테스트해볼 수 있습니다:

tokenize("This is the Hugging Face course!")

WordPiece 알고리즘에 대한 설명은 여기까지입니다! 이제 Unigram에 대해 알아보겠습니다.