MartinNav/compliantLLM

1"""Exact 256-entry byte-level tokenizer for compliantLLM."""2 3import json4import os5 6from transformers import PreTrainedTokenizer7 8 9def _byte_token(value):10    return f"<0x{value:02X}>"11 12 13class CompliantLLMTokenizer(PreTrainedTokenizer):14    """The zero-merge form of byte-level BPE."""15 16    model_input_names = ["input_ids", "attention_mask"]17    vocab_files_names = {"vocab_file": "vocab.json"}18 19    def __init__(self, vocab_file=None, **kwargs):20        del vocab_file21        self.encoder = {_byte_token(value): value for value in range(256)}22        self.decoder = {value: token for token, value in self.encoder.items()}23        kwargs.setdefault("pad_token", _byte_token(0))24        kwargs.setdefault("model_max_length", 1024)25        kwargs.setdefault("padding_side", "right")26        kwargs.setdefault("truncation_side", "left")27        super().__init__(**kwargs)28 29    @property30    def vocab_size(self):31        return 25632 33    def get_vocab(self):34        return dict(self.encoder)35 36    def _tokenize(self, text, **kwargs):37        del kwargs38        return [_byte_token(value) for value in text.encode("utf-8")]39 40    def _convert_token_to_id(self, token):41        return self.encoder.get(token, 0)42 43    def _convert_id_to_token(self, index):44        return self.decoder.get(index, _byte_token(0))45 46    def convert_tokens_to_string(self, tokens):47        values = [self.encoder[token] for token in tokens if token in self.encoder]48        return bytes(values).decode("utf-8", errors="replace")49 50    def save_vocabulary(self, save_directory, filename_prefix=None):51        os.makedirs(save_directory, exist_ok=True)52        filename = ((filename_prefix + "-") if filename_prefix else "") + "vocab.json"53        path = os.path.join(save_directory, filename)54        with open(path, "w", encoding="utf-8") as handle:55            json.dump(self.encoder, handle, indent=2, sort_keys=True)56            handle.write("\n")57        return (path,)58