Tokeny: jak BPE buduje słownik

BPE na 28 zdaniach labu w czystym Pythonie i numpy: słownik startowy, pierwsze scalenie, cel labu i koszt słownika.

import numpy as np
from collections import Counter

corpus = [
    "kot siedzi na parapecie", "kot je rybę na kanapie", "pies siedzi na kanapie", "pies je kość pod stołem",
    "koty lubią spać na parapecie", "psy lubią biegać po parku", "mały kot i mały pies śpią razem", "dzieci lubią koty i psy",
    "kot śpi na kanapie a pies na podłodze", "dzieci biegają po parku z psem", "mały pies je kość a kot je rybę", "koty i psy śpią razem na kanapie",
    "pies biega po parku a kot siedzi na parapecie", "koty lubią rybę a psy lubią kość", "dzieci karmią koty i psy", "mały kot śpi na podłodze",
    "pies śpi pod stołem", "dzieci lubią biegać z psem po parku", "kot patrzy na rybę", "mały pies patrzy na kota",
    "ptak siedzi na drzewie", "ptaki śpiewają na drzewie", "kot patrzy na ptaki", "babcia gotuje obiad dla dzieci",
    "dzieci jedzą obiad przy stole", "babcia karmi ptaki w ogrodzie", "pies biega po ogrodzie", "babcia lubi ptaki",
]

def letters(word):  # "kot" -> ["▁k", "o", "t"]: the first letter carries the word-start mark
    return ["▁" + word[0]] + list(word[1:])

def merge(syms, a, b):
    out, i = [], 0
    while i < len(syms):
        if i + 1 < len(syms) and syms[i] == a and syms[i + 1] == b:
            out.append(a + b); i += 2
        else:
            out.append(syms[i]); i += 1
    return out

def train(sentences):
    freq = Counter(w for s in sentences for w in s.split())
    vocab = [(letters(w), f) for w, f in freq.items()]
    merges = []
    while True:
        pairs = Counter()
        for syms, f in vocab:
            for a, b in zip(syms, syms[1:]):
                pairs[(a, b)] += f
        (a, b), n = max(pairs.items(), key=lambda kv: kv[1])  # ties: the pair met first
        if n < 2:
            return merges
        merges.append((a, b, n))
        vocab = [(merge(syms, a, b), f) for syms, f in vocab]

def encode(sentence, merges, k):
    out = []
    for w in sentence.split():
        syms = letters(w)
        for a, b, _ in merges[:k]:
            syms = merge(syms, a, b)
        out += syms
    return out

merges = train(corpus)
start = sorted({t for s in corpus for t in encode(s, [], 0)})
a, b, n = merges[0]
print(f"start vocabulary: {len(start)} symbols; merges possible: {len(merges)}, first: {a} + {b} ({n}x)")

s = "pies je kość pod stołem"
fewest = next(k for k in range(len(merges) + 1) if len(encode(s, merges, k)) <= len(s.split()))
print(f"'{s}': {len(encode(s, merges, 0))} tokens at 0 merges, {len(encode(s, merges, 80))} at 80, fewest for <= words: {fewest}")

for k in [0, 80]:
    tokens = sum(len(encode(x, merges, k)) for x in corpus)
    E = np.zeros((len(start) + k, 16))  # one 16-number embedding per token of the vocabulary
    print(f"{k} merges: corpus {tokens} tokens, vocabulary {len(start) + k}, embedding matrix {E.size} numbers")

Wynik

start vocabulary: 43 symbols; merges possible: 92, first: i + e (36x)
'pies je kość pod stołem': 19 tokens at 0 merges, 5 at 80, fewest for <= words: 68
0 merges: corpus 620 tokens, vocabulary 43, embedding matrix 688 numbers
80 merges: corpus 195 tokens, vocabulary 123, embedding matrix 1968 numbers

Wynik jest ten sam co w labie, bo kod robi te same kroki w tej samej kolejności, łącznie z remisami. Uruchomione w numpy 2.1.3.

Pobierz notebook (.ipynb). Otwórz go w Jupyterze albo wgraj do Google Colab; brakujące dane pobierze sam.