Tokenization & Vocabulary Design

from collections import Counter
import math

def learn_bpe_merges(corpus, num_merges=3):
    vocab = {tuple(word): freq for word, freq in corpus.items()}
    merges = []
    for _ in range(num_merges):
        counts = Counter()
        for symbols, freq in vocab.items():
            for pair in zip(symbols, symbols[1:]):
                counts[pair] += freq
        if not counts:
            break
        best = max(counts, key=counts.get)
        merges.append(best)
        merged_vocab = {}
        for symbols, freq in vocab.items():
            out, i = [], 0
            while i < len(symbols):
                if i + 1 < len(symbols) and (symbols[i], symbols[i + 1]) == best:
                    out.append(symbols[i] + symbols[i + 1])
                    i += 2
                else:
                    out.append(symbols[i])
                    i += 1
            merged_vocab[tuple(out)] = freq
        vocab = merged_vocab
    return merges, vocab

def unigram_map(x, tok_logp):
    n, NEG = len(x), -1e30
    dp, back = [NEG] * (n + 1), [None] * (n + 1)
    dp[0] = 0.0
    buckets = {}
    for tok, lp in tok_logp.items():
        buckets.setdefault(tok[0], []).append((tok, lp))
    for i in range(n):
        if dp[i] <= NEG / 2:
            continue
        for tok, lp in buckets.get(x[i], []):
            if x.startswith(tok, i):
                j, s = i + len(tok), dp[i] + lp
                if s > dp[j]:
                    dp[j], back[j] = s, (i, tok)
    if back[n] is None:
        return None
    out, j = [], n
    while j > 0:
        i, tok = back[j]
        out.append(tok)
        j = i
    return out[::-1], dp[n]

merges, vocab = learn_bpe_merges({"low": 5, "lower": 2, "new": 6, "newer": 3}, num_merges=4)
print("bpe merges:", merges)
print("bpe vocab:", vocab)

probs = {"the": 0.08, "there": 0.02, "re": 0.05, "th": 0.06, "e": 0.07, "r": 0.03, "t": 0.02, "h": 0.02}
tok_logp = {t: math.log(p) for t, p in probs.items()}
x = "there"
res = unigram_map(x, tok_logp)
assert res is not None, "No valid segmentation"
seg, score = res
print("x:", x)
print("seg:", seg, "tokens:", len(seg), "logp:", round(score, 3))