BPE vs. Unigram Tokenization

This small script compares ordered BPE merges with unigram path search on the same word.

from math import inf

word = "unhugs"
merges = [("u", "g"), ("u", "n"), ("h", "ug"), ("p", "un")]
cost = {
    "un": 1.10, "hugs": 1.35, "hug": 1.25, "s": 1.80,
    "u": 3.20, "n": 3.20, "h": 3.20, "g": 3.20, "ug": 2.10,
}

def apply_merge(parts, pair):
    out, i = [], 0
    while i < len(parts):
        if i + 1 < len(parts) and (parts[i], parts[i + 1]) == pair:
            out.append(parts[i] + parts[i + 1])
            i += 2
        else:
            out.append(parts[i])
            i += 1
    return out

def bpe_encode(text):
    parts = list(text)
    for pair in merges:
        parts = apply_merge(parts, pair)
    return parts

def unigram_encode(text):
    dp = [(inf, -1) for _ in range(len(text) + 1)]
    dp[0] = (0.0, 0)
    for end in range(1, len(text) + 1):
        for start in range(end):
            token = text[start:end]
            if token in cost and dp[start][0] + cost[token] < dp[end][0]:
                dp[end] = (dp[start][0] + cost[token], start)
    tokens, pos = [], len(text)
    while pos:
        start = dp[pos][1]
        tokens.append(text[start:pos])
        pos = start
    return list(reversed(tokens)), dp[-1][0]

print("BPE:", bpe_encode(word))
print("Unigram:", unigram_encode(word))

BPE returns ["un", "hug", "s"] because those pieces are created by the ordered merge list. The unigram path search returns ["un", "hugs"] because that two-piece path has lower total cost in the given scored vocabulary.