BPE Encode & Decode

The problem

Training BPE (byte-pair encoding) produced merges, a dict (a, b) -> new_id in the order the merges were learned: a lower id means learned earlier. Now use it, in both directions:

decode: ids → text;
encode: text → ids, cutting any new text the same way the training text was cut.

The examples below use the merges of the previous exercise: {(97, 97): 256, (256, 97): 257, (257, 98): 258}, where 97 is "a" and 98 is "b".

Part 1: the vocabulary

Every id stands for a string of bytes:

ids 0 to 255 are single bytes: id 97 is b"a";
each merge (a, b) -> n glues together the bytes of its two parts: vocab[n] = vocab[a] + vocab[b].

Example. 256 is b"a" + b"a" = b"aa". 257 is vocab[256] + vocab[97] = b"aaa". 258 is vocab[257] + vocab[98] = b"aaab".

Process the merges in id order: 257 needs 256 to exist already.

Part 2: decoding

Join the bytes of every id, then decode them as UTF-8 (Unicode Transformation Format, 8-bit) to get a string.

Example. [97, 98, 258] is b"a" + b"b" + b"aaab" = b"abaaab", which decodes to "abaaab".

A sequence of ids can end in the middle of a character: for example the first two of the four bytes of an emoji. Decoding those bytes normally raises an error, so decode with errors="replace", which turns them into the replacement character "�".

Part 3: encoding

Start from the bytes of the text. Then repeat: look at the pairs of neighbours present, and among those that were learned, apply the one with the lowest id (the earliest learned) everywhere. Stop when no pair present was learned, or when a single id is left.

Example. Encode "abaaab".

Bytes: [97, 98, 97, 97, 97, 98]. Pairs present: (97, 98), (98, 97), (97, 97). Only (97, 97) was learned (256): merge it, [97, 98, 256, 97, 98].

Pairs present: (97, 98), (98, 256), (256, 97). Only (256, 97) was learned (257): [97, 98, 257, 98].

Now (257, 98) (258): [97, 98, 258].

Pairs present: (97, 98), (98, 258). Neither was learned: stop.

The result is [97, 98, 258], and decoding it gives back "abaaab".

Why the earliest merge, and not the most frequent pair in the new text? Because training applied the merges in that order, and encoding must cut every text exactly as training would have. With "aaaa", merge 256 gives [256, 256]; no merge (256, 256) was learned, so encoding stops there, even though aaa (257) is also a token.

What to implement

build_vocab(merges): a dict mapping every id to its bytes, as in part 1.
decode(ids, merges): the string of part 2.
encode(text, merges): the list of ids of part 3.

The merge function from the earlier exercise is given in the starter.

Requirements

decode(encode(text, merges), merges) == text for any string, including text the tokenizer was never trained on, accents and emoji.
Apply merges in the order they were learned, not by how frequent a pair is in the text being encoded.

Hint

To find the earliest learned pair present, take the minimum of the pairs by merges.get(pair, float("inf")): pairs that were never learned count as infinitely late. If even the minimum is not in merges, no learned pair is left.