|
|
| """Exact linear-time RWKV World tokenizer for Transformers AutoTokenizer."""
|
|
|
| from __future__ import annotations
|
|
|
| import json
|
| import os
|
| import shutil
|
| from typing import Optional
|
|
|
| from transformers.tokenization_utils import PreTrainedTokenizer
|
|
|
|
|
| VOCAB_FILES_NAMES = {"tokenizer_file": "tokenizer.json"}
|
| END_TOKEN = "<|endoftext|>"
|
| INVALID_TOKEN_PREFIX = "\ue000"
|
| INVALID_TOKEN_SUFFIX = "\ue001"
|
|
|
|
|
| class _CharEncoding:
|
| """Minimal Encoding surface used by BatchEncoding.char_to_token."""
|
|
|
| def __init__(self, offsets):
|
| self.offsets = offsets
|
|
|
| def char_to_token(self, char_index, sequence_index=0):
|
| if sequence_index != 0:
|
| return None
|
| for token_index, offset in enumerate(self.offsets):
|
| if offset is not None and offset[0] <= char_index < offset[1]:
|
| return token_index
|
| return None
|
|
|
|
|
| def _bytes_to_unicode():
|
| values = list(range(ord("!"), ord("~") + 1))
|
| values += list(range(ord("¡"), ord("¬") + 1))
|
| values += list(range(ord("®"), ord("ÿ") + 1))
|
| characters = list(values)
|
| extra = 0
|
| for byte in range(256):
|
| if byte not in values:
|
| values.append(byte)
|
| characters.append(256 + extra)
|
| extra += 1
|
| return dict(zip(values, map(chr, characters)))
|
|
|
|
|
| class Rwkv7Tokenizer(PreTrainedTokenizer):
|
| """RWKV World longest-prefix tokenizer backed only by tokenizer.json."""
|
|
|
| vocab_files_names = VOCAB_FILES_NAMES
|
| model_input_names = ["input_ids", "attention_mask"]
|
| padding_side = "left"
|
|
|
| def __init__(
|
| self,
|
| tokenizer_file,
|
| eos_token=END_TOKEN,
|
| pad_token=END_TOKEN,
|
| unk_token=END_TOKEN,
|
| add_bos_token=False,
|
| **kwargs,
|
| ):
|
| if not tokenizer_file or not os.path.isfile(tokenizer_file):
|
| raise ValueError(f"tokenizer.json does not exist: {tokenizer_file}")
|
| if add_bos_token:
|
| raise ValueError("RWKV World has no BOS token")
|
| self.tokenizer_file = tokenizer_file
|
| self.add_bos_token = False
|
| value = json.loads(open(tokenizer_file, encoding="utf-8").read())
|
| model = value.get("model")
|
| if not isinstance(model, dict) or model.get("type") != "WordPiece":
|
| raise ValueError("tokenizer.json is not the locked RWKV WordPiece model")
|
| if (
|
| model.get("unk_token") != END_TOKEN
|
| or model.get("continuing_subword_prefix") != ""
|
| or value.get("normalizer") is not None
|
| ):
|
| raise ValueError("tokenizer.json does not use exact RWKV semantics")
|
| pre_tokenizer = value.get("pre_tokenizer")
|
| decoder = value.get("decoder")
|
| if (
|
| not isinstance(pre_tokenizer, dict)
|
| or pre_tokenizer.get("type") != "ByteLevel"
|
| or pre_tokenizer.get("add_prefix_space") is not False
|
| or pre_tokenizer.get("use_regex") is not False
|
| or not isinstance(decoder, dict)
|
| or decoder.get("type") != "ByteLevel"
|
| ):
|
| raise ValueError("tokenizer.json does not use exact ByteLevel semantics")
|
| vocabulary = model.get("vocab")
|
| if not isinstance(vocabulary, dict) or not vocabulary:
|
| raise ValueError("tokenizer.json has no vocabulary")
|
| if any(not isinstance(token, str) or not isinstance(index, int) for token, index in vocabulary.items()):
|
| raise TypeError("tokenizer.json vocabulary entry is invalid")
|
| if set(vocabulary.values()) != set(range(len(vocabulary))):
|
| raise ValueError("tokenizer.json vocabulary ids are not contiguous")
|
| if vocabulary.get(END_TOKEN) != 0:
|
| raise ValueError("tokenizer.json must reserve id 0 for end-of-text")
|
|
|
| byte_decoder = {character: byte for byte, character in _bytes_to_unicode().items()}
|
| self.encoder = dict(vocabulary)
|
| self.decoder = {index: token for token, index in vocabulary.items()}
|
| self._token_bytes = {}
|
| reverse = {}
|
| for token, index in vocabulary.items():
|
| if index == 0:
|
| continue
|
| placeholder = f"{INVALID_TOKEN_PREFIX}{index}{INVALID_TOKEN_SUFFIX}"
|
| if token == placeholder:
|
| continue
|
| try:
|
| raw = bytes(byte_decoder[character] for character in token)
|
| except KeyError as error:
|
| raise ValueError("tokenizer.json vocabulary is not ByteLevel encoded") from error
|
| if raw in reverse:
|
| raise ValueError("tokenizer.json contains duplicate byte tokens")
|
| reverse[raw] = index
|
| self._token_bytes[token] = raw
|
| missing = [byte for byte in range(256) if bytes([byte]) not in reverse]
|
| if missing:
|
| raise ValueError(f"tokenizer.json is missing singleton bytes: {missing}")
|
| for raw, index in reverse.items():
|
| for endpoint in range(1, len(raw)):
|
| prefix = reverse.get(raw[:endpoint])
|
| if prefix is not None and prefix > index:
|
| raise ValueError("token ranks do not support longest-prefix encoding")
|
|
|
| self._children = [{}]
|
| self._terminal = [None]
|
| self._max_token_bytes = 1
|
| for raw, index in sorted(reverse.items(), key=lambda item: item[1]):
|
| node = 0
|
| for byte in raw:
|
| child = self._children[node].get(byte)
|
| if child is None:
|
| child = len(self._children)
|
| self._children[node][byte] = child
|
| self._children.append({})
|
| self._terminal.append(None)
|
| node = child
|
| self._terminal[node] = index
|
| self._max_token_bytes = max(self._max_token_bytes, len(raw))
|
|
|
| if "additional_special_tokens" not in kwargs:
|
| appended = []
|
| for item in sorted(value.get("added_tokens", []), key=lambda item: item.get("id", -1)):
|
| if not isinstance(item, dict):
|
| raise TypeError("tokenizer.json added token is invalid")
|
| content = item.get("content")
|
| index = item.get("id")
|
| if not isinstance(content, str) or not isinstance(index, int):
|
| raise TypeError("tokenizer.json added token is invalid")
|
| if content != END_TOKEN:
|
| if not item.get("special") or index < len(vocabulary):
|
| raise ValueError("only append-only special tokens are supported")
|
| appended.append(content)
|
| if appended:
|
| kwargs["additional_special_tokens"] = appended
|
| super().__init__(
|
| eos_token=eos_token,
|
| pad_token=pad_token,
|
| unk_token=unk_token,
|
| add_bos_token=self.add_bos_token,
|
| **kwargs,
|
| )
|
|
|
| @property
|
| def vocab_size(self):
|
| return len(self.encoder)
|
|
|
| def get_vocab(self):
|
| vocabulary = dict(self.encoder)
|
| vocabulary.update(self.added_tokens_encoder)
|
| return vocabulary
|
|
|
| def _encode_bytes(self, data):
|
| output = []
|
| position = 0
|
| while position < len(data):
|
| node = 0
|
| cursor = position
|
| best_id = None
|
| best_end = position
|
| limit = min(len(data), position + self._max_token_bytes)
|
| while cursor < limit:
|
| child = self._children[node].get(data[cursor])
|
| if child is None:
|
| break
|
| node = child
|
| cursor += 1
|
| token_id = self._terminal[node]
|
| if token_id is not None and (best_id is None or token_id > best_id):
|
| best_id = token_id
|
| best_end = cursor
|
| if best_id is None:
|
| raise RuntimeError("singleton-byte vocabulary invariant was violated")
|
| output.append(best_id)
|
| position = best_end
|
| return output
|
|
|
| def _encode_bytes_with_offsets(self, data, char_offset, byte_to_char):
|
| output = []
|
| position = 0
|
| while position < len(data):
|
| node = 0
|
| cursor = position
|
| best_id = None
|
| best_end = position
|
| limit = min(len(data), position + self._max_token_bytes)
|
| while cursor < limit:
|
| child = self._children[node].get(data[cursor])
|
| if child is None:
|
| break
|
| node = child
|
| cursor += 1
|
| token_id = self._terminal[node]
|
| if token_id is not None and (best_id is None or token_id > best_id):
|
| best_id = token_id
|
| best_end = cursor
|
| if best_id is None:
|
| raise RuntimeError("singleton-byte vocabulary invariant was violated")
|
| start_char = char_offset + byte_to_char[position]
|
| end_char = char_offset + byte_to_char[best_end - 1] + 1
|
| output.append((best_id, (start_char, end_char)))
|
| position = best_end
|
| return output
|
|
|
| def _encode_with_offsets(self, text):
|
| special = sorted(self.added_tokens_encoder, key=lambda token: (-len(token), token))
|
| output = []
|
| ordinary_start = 0
|
| position = 0
|
|
|
| def append_ordinary(segment, char_offset):
|
| if not segment:
|
| return
|
| data = segment.encode("utf-8")
|
| byte_to_char = []
|
| for char_index, character in enumerate(segment):
|
| byte_to_char.extend([char_index] * len(character.encode("utf-8")))
|
| output.extend(
|
| self._encode_bytes_with_offsets(data, char_offset, byte_to_char)
|
| )
|
|
|
| while position < len(text):
|
| matched = next(
|
| (token for token in special if text.startswith(token, position)), None
|
| )
|
| if matched is None:
|
| position += 1
|
| continue
|
| append_ordinary(text[ordinary_start:position], ordinary_start)
|
| output.append(
|
| (
|
| self.added_tokens_encoder[matched],
|
| (position, position + len(matched)),
|
| )
|
| )
|
| position += len(matched)
|
| ordinary_start = position
|
| append_ordinary(text[ordinary_start:], ordinary_start)
|
| return output
|
|
|
| def _tokenize(self, text, **kwargs):
|
| del kwargs
|
| return [self.decoder[index] for index in self._encode_bytes(text.encode("utf-8"))]
|
|
|
| def __call__(self, text=None, text_pair=None, **kwargs):
|
| output = super().__call__(text=text, text_pair=text_pair, **kwargs)
|
| if text_pair is not None or text is None:
|
| return output
|
| texts = [text] if isinstance(text, str) else list(text)
|
| if any(not isinstance(item, str) for item in texts):
|
| return output
|
| input_ids = output["input_ids"]
|
| attention_mask = output.get("attention_mask")
|
| raw_ids = input_ids.tolist() if hasattr(input_ids, "tolist") else input_ids
|
| if raw_ids and isinstance(raw_ids[0], int):
|
| rows = [raw_ids]
|
| masks = [attention_mask] if attention_mask is not None else None
|
| else:
|
| rows = raw_ids
|
| if attention_mask is None:
|
| masks = None
|
| else:
|
| masks = (
|
| attention_mask.tolist()
|
| if hasattr(attention_mask, "tolist")
|
| else attention_mask
|
| )
|
| encodings = []
|
| for index, source in enumerate(texts):
|
| exact = self._encode_with_offsets(source)
|
| row = rows[index]
|
| active = len(row) if masks is None else sum(int(item) for item in masks[index])
|
| exact = exact[:active]
|
| offsets = [offset for _token_id, offset in exact]
|
| missing = len(row) - len(offsets)
|
| if self.padding_side == "left":
|
| offsets = [None] * missing + offsets
|
| else:
|
| offsets.extend([None] * missing)
|
| encodings.append(_CharEncoding(offsets))
|
| output._encodings = encodings
|
| return output
|
|
|
| def _convert_token_to_id(self, token):
|
| return self.encoder.get(token, 0)
|
|
|
| def _convert_id_to_token(self, index):
|
| return self.decoder.get(index, END_TOKEN)
|
|
|
| def convert_tokens_to_string(self, tokens):
|
| decoded = bytearray()
|
| for token in tokens:
|
| raw = self._token_bytes.get(token)
|
| if raw is not None:
|
| decoded.extend(raw)
|
| else:
|
| decoded.extend(str(token).encode("utf-8"))
|
| return bytes(decoded).decode("utf-8", errors="replace")
|
|
|
| def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
|
| bos = [self.bos_token_id] if self.add_bos_token else []
|
| output = bos + list(token_ids_0)
|
| if token_ids_1 is not None:
|
| output.extend(bos + list(token_ids_1))
|
| return output
|
|
|
| def save_vocabulary(self, save_directory, filename_prefix: Optional[str] = None):
|
| if not os.path.isdir(save_directory):
|
| raise ValueError("save directory does not exist")
|
| filename = (f"{filename_prefix}-" if filename_prefix else "") + "tokenizer.json"
|
| destination = os.path.join(save_directory, filename)
|
| if os.path.abspath(destination) != os.path.abspath(self.tokenizer_file):
|
| shutil.copyfile(self.tokenizer_file, destination)
|
| return (destination,)
|
|
|
|
|
| __all__ = ["Rwkv7Tokenizer"]
|
|
|
|
|
