byte pair encoding utilities import os import json import regex as re from functools import lrucache lrucache def bytestounicode bs listrangeord ord1listrangeord ord1listrangeord ord1 cs bs n 0 for b in range28 if b not in bs bs appendb cs append28n n 1 cs chrn for n in cs return dictzipbs cs def getpairsword pairs set prevchar word0 for char in word1 pairs addprevchar char prevchar char return pairs class encoder def initself encoder bpemerges errors replace self encoder encoder self decoder v k for k v in self encoder items self errors errors how to handle errors in decoding self byteencoder bytestounicode self bytedecoder v k for k v in self byteencoder items self bperanks dictzipbpemerges rangelenbpemerges self cache should haved added re ignorecase so bpe merges can happen for capitalized versions of contractions self pat re compiler s t re ve m ll d pl pn splpns ss def bpeself token if token in self cache return self cachetoken word tupletoken pairs getpairsword if not pairs return token while true bigram minpairs key lambda pair self bperanks getpair float inf if bigram not in self bperanks break first second bigram newword i 0 while i lenword try j word indexfirst i newword extendwordi j i j except newword extendwordi break if wordi first and i lenword1 and wordi1 second newword appendfirstsecond i 2 else newword appendwordi i 1 newword tuplenewword word newword if lenword 1 break else pairs getpairsword word joinword self cachetoken word return word def encodeself text bpetokens for token in re findallself pat text token joinself byteencoderb for b in token encode utf8 bpetokens extendself encoderbpetoken for bpetoken in self bpetoken split return bpetokens def decodeself tokens text joinself decodertoken for token in tokens text bytearrayself bytedecoderc for c in text decode utf8 errorsself errors return text def getencodermodelname modelsdir with openos path joinmodelsdir modelname encoder json r as f encoder json loadf with openos path joinmodelsdir modelname vocab bpe r encodingutf8 as f bpedata f read bpemerges tuplemergestr split for mergestr in bpedata split n 1 1 return encoder encoderencoder bpemergesbpemerges returns list of utf 8 byte and a corresponding list of unicode strings the reversible bpe codes work on unicode strings this means you need a large of unicode characters in your vocab if you want to avoid unks when you re at something like a 10b token dataset you end up needing around 5k for decent coverage this is a signficant percentage of your normal say 32k bpe vocab to avoid that we want lookup tables between utf 8 bytes and unicode strings and avoids mapping to whitespace control characters the bpe code barfs on return set of symbol pairs in a word word is represented as tuple of symbols symbols being variable length strings how to handle errors in decoding should haved added re ignorecase so bpe merges can happen for capitalized versions of contractions s t re ve m ll d p l p n s p l p n s s s
import os
import json
import regex as re
from functools import lru_cache

@lru_cache()
def bytes_to_unicode():
    
    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
    cs = bs[:]
    n = 0
    for b in range(2**8):
        if b not in bs:
            bs.append(b)
            cs.append(2**8+n)
            n += 1
    cs = [chr(n) for n in cs]
    return dict(zip(bs, cs))

def get_pairs(word):
    
    pairs = set()
    prev_char = word[0]
    for char in word[1:]:
        pairs.add((prev_char, char))
        prev_char = char
    return pairs

class Encoder:
    def __init__(self, encoder, bpe_merges, errors='replace'):
        self.encoder = encoder
        self.decoder = {v:k for k,v in self.encoder.items()}
        self.errors = errors 
        self.byte_encoder = bytes_to_unicode()
        self.byte_decoder = {v:k for k, v in self.byte_encoder.items()}
        self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))
        self.cache = {}

        
        self.pat = re.compile(r)

    def bpe(self, token):
        if token in self.cache:
            return self.cache[token]
        word = tuple(token)
        pairs = get_pairs(word)

        if not pairs:
            return token

        while True:
            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
            if bigram not in self.bpe_ranks:
                break
            first, second = bigram
            new_word = []
            i = 0
            while i < len(word):
                try:
                    j = word.index(first, i)
                    new_word.extend(word[i:j])
                    i = j
                except:
                    new_word.extend(word[i:])
                    break

                if word[i] == first and i < len(word)-1 and word[i+1] == second:
                    new_word.append(first+second)
                    i += 2
                else:
                    new_word.append(word[i])
                    i += 1
            new_word = tuple(new_word)
            word = new_word
            if len(word) == 1:
                break
            else:
                pairs = get_pairs(word)
        word = ' '.join(word)
        self.cache[token] = word
        return word

    def encode(self, text):
        bpe_tokens = []
        for token in re.findall(self.pat, text):
            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
        return bpe_tokens

    def decode(self, tokens):
        text = ''.join([self.decoder[token] for token in tokens])
        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors=self.errors)
        return text

def get_encoder(model_name, models_dir):
    with open(os.path.join(models_dir, model_name, 'encoder.json'), 'r') as f:
        encoder = json.load(f)
    with open(os.path.join(models_dir, model_name, 'vocab.bpe'), 'r', encoding="utf-8") as f:
        bpe_data = f.read()
    bpe_merges = [tuple(merge_str.split()) for merge_str in bpe_data.split('\n')[1:-1]]
    return Encoder(
        encoder=encoder,
        bpe_merges=bpe_merges,
    )