summarize_from_feedback/datasets/encodings.py

import json
import os.path
from dataclasses import dataclass
from functools import lru_cache
from typing import ClassVar

import blobfile as bf
import regex as re
import torch

ENCODINGS_BASE = "https://openaipublic.blob.core.windows.net/summarize-from-feedback/encodings"


def read_file(path):
    with bf.BlobFile(path, "rb") as f:
        return f.read()


@lru_cache()
def bytes_to_unicode():
    """
    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.
    """
    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))


@dataclass
class Encoding:
    name: str
    n_vocab: int
    # End of text token
    eot_token: int = None
    registry: ClassVar[dict] = dict()
    eoprefix_token: int = None

    def __post_init__(self):
        self._register()

    def _register(self):
        assert self.name not in self.registry
        self.registry[self.name] = self

    def __call__(self, text):
        return self.encode(text)

    def encode(self, text):
        """Convert text (or other data) into an array of integer tokens"""
        raise NotImplementedError

    def decode(self, tokens) -> str:
        """Convert array of integer tokens into text (or other data)."""
        raise NotImplementedError


@dataclass
class BPEEncoding(Encoding):
    base_path: str = None
    encoder_path: str = "encoder.json"
    bpe_path: str = "vocab.bpe"
    n_denoise_sentinels: int = 0

    # Should haved added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
    pat = re.compile(
        r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
    )

    def __post_init__(self):
        super().__post_init__()

        self.base_path = self.base_path or os.path.join(ENCODINGS_BASE, self.name)
        self.full_encoder_path = os.path.join(self.base_path, self.encoder_path)
        self.full_bpe_path = os.path.join(self.base_path, self.bpe_path)

        # We don't load the full BPE data until needed
        # Initialize them to 'None' here to make the linter happy
        # Note: Do not move these to @dataclass fields. That will lead to enormous outputs,
        #       since the default dataclass __repr__ prints the values of all fields.
        self._token_str_to_idx = None
        self._token_idx_to_str = None
        self.byte_decoder = None
        self.byte_encoder = None
        self.bpe_ranks = None

        # Without this cache, performance is 5x slower
        self.bpe = lru_cache(maxsize=2 ** 17)(self.bpe)

    def _load(self):
        if self._token_str_to_idx is not None:
            return

        self._token_str_to_idx = json.loads(read_file(self.full_encoder_path).decode())
        bpe_data = read_file(self.full_bpe_path).decode()
        bpe_merges = [tuple(merge_str.split()) for merge_str in bpe_data.split("\n")[1:-1]]

        assert self.eot_token == self._token_str_to_idx["<|endoftext|>"]

        # Add an <|end_of_prefix|> token
        if self.eoprefix_token is not None:
            assert not self._token_str_to_idx.get("<|endofprefix|>")
            self._token_str_to_idx["<|endofprefix|>"] = self.eoprefix_token

        # Add denoise sentinel tokens like <|dn_1|> <|dn_2|> etc.
        # These tokens are added to the end of the vocabulary range
        for denoise_sentinel_idx in range(self.n_denoise_sentinels):
            str_repr = f"<|dn_{denoise_sentinel_idx}|>"
            assert not self._token_str_to_idx.get(str_repr)

            n_non_sentinel_tokens = self.n_vocab - self.n_denoise_sentinels
            sentinel_token = n_non_sentinel_tokens + denoise_sentinel_idx

            self._token_str_to_idx[str_repr] = sentinel_token

        assert len(self._token_str_to_idx) == self.n_vocab

        self._token_idx_to_str = {v: k for k, v in self._token_str_to_idx.items()}
        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))))

    def bpe(self, token):  # pylint: disable=method-hidden
        word = tuple(self.byte_encoder[b] for b in token.encode("utf-8"))

        if len(word) == 1:
            return word

        while True:
            min_pair = None
            min_idxs = []
            min_rank = None
            for i, pair in enumerate(zip(word[:-1], word[1:])):
                try:
                    rank = self.bpe_ranks[pair]
                    if min_rank is None or rank < min_rank:
                        min_rank = rank
                        min_pair = pair
                        del min_idxs[::]
                        min_idxs.append(i)
                    elif min_rank == rank and i > min_idxs[-1] + 1:
                        min_idxs.append(i)
                except KeyError:
                    pass

            if min_pair is None:
                break

            new_word = []
            i_start = 0
            for i in min_idxs:
                new_word.extend(word[i_start:i])
                new_word.append(min_pair[0] + min_pair[1])
                i_start = i + 2

            if i_start < len(word):
                new_word.extend(word[i_start:])

            new_word = tuple(new_word)
            word = new_word
            if len(word) == 1:
                break

        return word

    def encode(self, text):
        self._load()
        bpe_tokens = [
            self._token_str_to_idx[bpe_token]
            for token in re.findall(self.pat, text)
            for bpe_token in self.bpe(token)
        ]
        return bpe_tokens

    def decode(self, tokens, errors="replace"):
        if isinstance(tokens, torch.Tensor):
            if tokens.dim() == 1 and tokens.size(0) == 1:
                tokens = tokens.tolist()
            else:
                tokens = tokens.squeeze().tolist()

        decoded_bytes = self.decode_bytes(tokens)
        text = decoded_bytes.decode("utf-8", errors=errors)
        return text

    def decode_bytes(self, tokens):
        self._load()
        text = "".join([self._token_idx_to_str[token] for token in tokens])
        return bytearray([self.byte_decoder[c] for c in text])


Reversible = BPEEncoding(name="reversible_50000", n_vocab=50257, eot_token=50256)