Turn text into tokens

"""A fixed character tokenizer. No pretrained vocabulary or arithmetic rules."""

import json
from pathlib import Path


class CharacterTokenizer:
    VERSION = "printable-ascii-v1"
    PAD_ID = 0
    BOS_ID = 1
    EOS_ID = 2
    TOKENS = ("<PAD>", "<BOS>", "<EOS>") + tuple(chr(code) for code in range(32, 127))

    def __init__(self):
        self.token_to_id = {token: index for index, token in enumerate(self.TOKENS)}

    @property
    def vocab_size(self):
        return len(self.TOKENS)

    def encode(self, text, *, add_bos=False, add_eos=False):
        if not isinstance(text, str):
            raise TypeError("Text must be a string")
        ids = [self.BOS_ID] if add_bos else []
        for index, character in enumerate(text):
            if character not in self.token_to_id:
                raise ValueError(f"Unsupported character {character!r} at position {index}; "
                                 "this vocabulary accepts printable ASCII only")
            ids.append(self.token_to_id[character])
        if add_eos:
            ids.append(self.EOS_ID)
        return ids

    def decode(self, ids, *, skip_special_tokens=False):
        pieces = []
        for token_id in ids:
            if type(token_id) is not int or not 0 <= token_id < self.vocab_size:
                raise ValueError(f"Invalid token ID: {token_id!r}")
            if skip_special_tokens and token_id <= self.EOS_ID:
                continue
            pieces.append(self.TOKENS[token_id])
        return "".join(pieces)

    def specification(self):
        return {
            "version": self.VERSION,
            "vocab_size": self.vocab_size,
            "special_ids": {"pad": self.PAD_ID, "bos": self.BOS_ID, "eos": self.EOS_ID},
            "tokens": list(self.TOKENS),
        }

    def save(self, path):
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(self.specification(), indent=2, sort_keys=True) + "\n",
                        encoding="utf-8")

    @classmethod
    def load(cls, path):
        saved = json.loads(Path(path).read_text(encoding="utf-8"))
        tokenizer = cls()
        if saved != tokenizer.specification():
            raise ValueError("Saved vocabulary differs from printable-ascii-v1; "
                             "do not change token IDs for an existing model")
        return tokenizer

"""Prepare shifted next-token examples and right-padded PyTorch batches."""

from dataclasses import dataclass
import json
from pathlib import Path

from .tokenizer import CharacterTokenizer


CONTEXT_LENGTH = 32
SPLITS = ("train", "validation", "calibration", "test", "ood")


@dataclass(frozen=True)
class TrainingExample:
    id: str
    group_id: str
    input_ids: tuple[int, ...]
    target_ids: tuple[int, ...]

    def record(self):
        return {"id": self.id, "group_id": self.group_id,
                "input_ids": list(self.input_ids), "target_ids": list(self.target_ids)}


def encode_record(record, tokenizer, *, context_length=CONTEXT_LENGTH):
    if record["text"] != record["prompt"] + record["answer"]:
        raise ValueError("Completed text must equal prompt plus answer")
    if not record["prompt"] or not record["answer"]:
        raise ValueError("Both prompt and answer must be nonempty")
    if type(context_length) is not int or context_length < 3:
        raise ValueError("Context length must be an integer of at least 3")
    sequence = tokenizer.encode(record["text"], add_bos=True, add_eos=True)
    # Reserve room for the complete example including both boundary markers.
    if len(sequence) > context_length:
        raise ValueError(f"Example {record['id']} needs {len(sequence)} tokens; "
                         f"limit is {context_length}. No text was truncated.")
    return TrainingExample(record["id"], record["group_id"],
                           tuple(sequence[:-1]), tuple(sequence[1:]))


def encode_prompt(prompt, tokenizer, *, context_length=CONTEXT_LENGTH):
    if type(context_length) is not int or context_length < 3:
        raise ValueError("Context length must be an integer of at least 3")
    ids = tokenizer.encode(prompt, add_bos=True)
    if len(ids) > context_length:
        raise ValueError("Prompt exceeds the context length")
    # No answer and no EOS. Empty text is a BOS-only unconditional prompt.
    return ids


def load_examples(directory, *, split="train", tokenizer=None):
    if split not in SPLITS:
        raise ValueError(f"Unknown split {split!r}")
    tokenizer = tokenizer or CharacterTokenizer()
    path = Path(directory) / f"{split}.jsonl"
    return [encode_record(json.loads(line), tokenizer)
            for line in path.read_text(encoding="utf-8").splitlines()]


def make_batch(examples, *, device="cpu"):
    import torch

    examples = list(examples)
    if not examples:
        raise ValueError("Cannot make an empty batch")
    tokenizer = CharacterTokenizer()
    for example in examples:
        if not example.input_ids or len(example.input_ids) != len(example.target_ids):
            raise ValueError("Each example needs equal, nonempty input and target lengths")
        if len(example.input_ids) >= CONTEXT_LENGTH:
            raise ValueError("A complete example must fit within the 32-token context")
        tokenizer.decode(example.input_ids)
        tokenizer.decode(example.target_ids)
        if (example.input_ids[0] != tokenizer.BOS_ID
                or example.target_ids[-1] != tokenizer.EOS_ID
                or example.input_ids[1:] != example.target_ids[:-1]
                or any(token <= tokenizer.EOS_ID for token in example.input_ids[1:])):
            raise ValueError("Example is not a shifted BOS/text/EOS sequence")
    lengths = [len(example.input_ids) for example in examples]
    width = max(lengths)
    inputs = torch.full((len(examples), width), tokenizer.PAD_ID, dtype=torch.long)
    targets = torch.full_like(inputs, tokenizer.PAD_ID)
    for index, example in enumerate(examples):
        inputs[index, :lengths[index]] = torch.tensor(example.input_ids, dtype=torch.long)
        targets[index, :lengths[index]] = torch.tensor(example.target_ids, dtype=torch.long)
    return {"input_ids": inputs.to(device), "target_ids": targets.to(device),
            "lengths": torch.tensor(lengths, dtype=torch.long, device=device),
            "loss_mask": (targets != tokenizer.PAD_ID).to(device),
            "ids": [example.id for example in examples]}