Tokens

final class BPE
{
    public const END = '_';           // marks the end of a word, so "er_" (end) ≠ "er" (middle)
    public array $merges = [];        // in the order they were learned: [left, right, how often]
    private array $rank = [];         // "left right" → merge number (lower = learned earlier)
    public array $trainedWords = [];  // how each training word is split at the end

    // Split text into words and count them.
    public static function wordCounts(string $text): array
    {
        preg_match_all("/[a-z0-9']+/", strtolower($text), $m);
        return array_count_values($m[0]);
    }

    private static function chars(string $word): array
    {
        $c = preg_split('//u', $word, -1, PREG_SPLIT_NO_EMPTY);
        $c[] = self::END;
        return $c;
    }

    // Replace every "a" followed by "b" with the single token "ab".
    private static function merge(array $symbols, string $a, string $b): array
    {
        $out = [];
        for ($i = 0; $i < count($symbols); $i++) {
            if ($i < count($symbols) - 1 && $symbols[$i] === $a && $symbols[$i + 1] === $b) {
                $out[] = $a . $b; $i++;
            } else {
                $out[] = $symbols[$i];
            }
        }
        return $out;
    }

    // TRAINING: repeatedly glue the most common neighbouring pair.
    public function train(string $text, int $numMerges): void
    {
        $words = [];
        foreach (self::wordCounts($text) as $w => $count) $words[] = ['sym' => self::chars((string) $w), 'n' => $count, 'word' => (string) $w];

        for ($m = 0; $m < $numMerges; $m++) {
            $pairs = [];
            foreach ($words as $w) {
                for ($i = 0; $i < count($w['sym']) - 1; $i++) {
                    $key = $w['sym'][$i] . ' ' . $w['sym'][$i + 1];
                    $pairs[$key] = ($pairs[$key] ?? 0) + $w['n'];
                }
            }
            if (!$pairs) break;
            $best = max($pairs);
            if ($best < 2) break;                          // nothing repeats any more
            $ties = array_keys($pairs, $best);
            sort($ties);                                   // break ties alphabetically, so it's repeatable
            [$a, $b] = explode(' ', (string) $ties[0]);
            $this->rank[$a . ' ' . $b] = count($this->merges);
            $this->merges[] = [$a, $b, $best];
            foreach ($words as &$w) $w['sym'] = self::merge($w['sym'], $a, $b);
            unset($w);
        }
        usort($words, fn($x, $y) => $y['n'] <=> $x['n'] ?: strcmp($x['word'], $y['word']));
        $this->trainedWords = $words;
    }

    // ENCODING a new word: apply the learned merges in the order they were
    // learned (earliest first), until no learned pair is left.
    public function encodeWord(string $word): array
    {
        $sym = self::chars($word);
        while (true) {
            $bestRank = PHP_INT_MAX; $bestPair = null;
            for ($i = 0; $i < count($sym) - 1; $i++) {
                $r = $this->rank[$sym[$i] . ' ' . $sym[$i + 1]] ?? PHP_INT_MAX;
                if ($r < $bestRank) { $bestRank = $r; $bestPair = [$sym[$i], $sym[$i + 1]]; }
            }
            if ($bestPair === null) return $sym;
            $sym = self::merge($sym, $bestPair[0], $bestPair[1]);
        }
    }

    public function encode(string $text): array
    {
        preg_match_all("/[a-z0-9']+/", strtolower($text), $m);
        $tokens = [];
        foreach ($m[0] as $w) foreach ($this->encodeWord($w) as $t) $tokens[] = $t;
        return $tokens;
    }

    // The final vocabulary: every starting character, then every merge result.
    // A token's ID is just its position in this list.
    public function vocabulary(string $text): array
    {
        $chars = [];
        foreach (array_keys(self::wordCounts($text)) as $w) foreach (self::chars((string) $w) as $c) $chars[$c] = true;
        $vocab = array_keys($chars);
        sort($vocab);
        foreach ($this->merges as [$a, $b]) $vocab[] = $a . $b;
        return array_values(array_unique(array_map('strval', $vocab)));
    }
}