File size: 5,559 Bytes
9b994a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | from __future__ import annotations
import json
import re
from pathlib import Path
from typing import Iterable
BASE_VOCAB_SIZE = 256
ATOM_RE = re.compile(
r"<!--[\s\S]*?-->|<!doctype\s+html>|</?[A-Za-z][^<>]*>|"
r"\b(?:display|place-items|align-items|justify-content|min-height|margin|font|padding|border|border-radius|background)\s*:\s*[^;{}]+;?|"
r"\s+|[A-Za-z_$][A-Za-z0-9_$-]*|[0-9]+(?:\.[0-9]+)?(?:%|px|rem|em|vh|vw)?|[^\x00-\x7F]|[^\w\s]",
re.IGNORECASE | re.ASCII,
)
class ConstrainedBPE:
"""Byte-fallback BPE matching DivCenterer's syntax-aware Node tokenizer."""
def __init__(self, model: dict):
if model.get("format") != "divcenterer-constrained-byte-bpe-v1":
raise ValueError("Unrecognized tokenizer model format")
self.special_tokens: list[str] = model["special_tokens"]
self.special_to_id = {
token: BASE_VOCAB_SIZE + index
for index, token in enumerate(self.special_tokens)
}
self.special_by_id = {value: key for key, value in self.special_to_id.items()}
self.merge_ranks: dict[tuple[int, int], int] = {}
self.token_bytes: list[bytes | None] = [bytes([index]) for index in range(BASE_VOCAB_SIZE)]
self.token_bytes.extend([None] * len(self.special_tokens))
for rank, (left, right) in enumerate(model["merges"]):
if self.token_bytes[left] is None or self.token_bytes[right] is None:
raise ValueError(f"Merge {rank} refers to a special or missing token")
token_id = BASE_VOCAB_SIZE + len(self.special_tokens) + rank
self.merge_ranks[(left, right)] = rank
self.token_bytes.append(self.token_bytes[left] + self.token_bytes[right])
self._special_re = re.compile("(" + "|".join(re.escape(t) for t in self.special_tokens) + ")")
self._atom_cache: dict[str, list[int]] = {}
@classmethod
def from_file(cls, filename: str | Path) -> "ConstrainedBPE":
with Path(filename).open("r", encoding="utf-8") as stream:
return cls(json.load(stream))
def _atoms(self, text: str) -> list[str]:
atoms: list[str] = []
cursor = 0
for match in ATOM_RE.finditer(text):
if match.start() != cursor:
raise ValueError(f"Pretokenizer left a gap at character {cursor}")
atoms.append(match.group(0))
cursor = match.end()
if cursor != len(text):
raise ValueError(f"Pretokenizer left a gap at character {cursor}")
return atoms
@staticmethod
def _apply_merge(ids: list[int], left: int, right: int, merged: int) -> list[int]:
result: list[int] = []
index = 0
while index < len(ids):
if index + 1 < len(ids) and ids[index] == left and ids[index + 1] == right:
result.append(merged)
index += 2
else:
result.append(ids[index])
index += 1
return result
def _encode_atom(self, atom: str) -> list[int]:
cached = self._atom_cache.get(atom)
if cached is not None:
return cached.copy()
ids = list(atom.encode("utf-8"))
while len(ids) > 1:
best_rank: int | None = None
best_pair: tuple[int, int] | None = None
for left, right in zip(ids, ids[1:]):
rank = self.merge_ranks.get((left, right))
if rank is not None and (best_rank is None or rank < best_rank):
best_rank = rank
best_pair = (left, right)
if best_pair is None or best_rank is None:
break
merged_id = BASE_VOCAB_SIZE + len(self.special_tokens) + best_rank
ids = self._apply_merge(ids, best_pair[0], best_pair[1], merged_id)
self._atom_cache[atom] = ids.copy()
return ids
def encode(self, text: str) -> list[int]:
result: list[int] = []
for part in self._special_re.split(text):
if not part:
continue
special_id = self.special_to_id.get(part)
if special_id is not None:
result.append(special_id)
else:
for atom in self._atoms(part):
result.extend(self._encode_atom(atom))
return result
def decode(self, ids: Iterable[int]) -> str:
output: list[str] = []
byte_buffer = bytearray()
def flush_bytes() -> None:
if byte_buffer:
output.append(bytes(byte_buffer).decode("utf-8", errors="replace"))
byte_buffer.clear()
for token_id in ids:
special = self.special_by_id.get(int(token_id))
if special is not None:
flush_bytes()
output.append(special)
elif 0 <= int(token_id) < len(self.token_bytes) and self.token_bytes[int(token_id)] is not None:
byte_buffer.extend(self.token_bytes[int(token_id)])
else:
raise ValueError(f"Unknown token id: {token_id}")
flush_bytes()
return "".join(output)
def special_id(self, token: str) -> int:
return self.special_to_id[token]
@property
def vocab_size(self) -> int:
return len(self.token_bytes)
def conversation_prefix(prompt: str) -> str:
return f"<|user|>\n{prompt}\n<|assistant|>\n"
def conversation_text(prompt: str, answer: str) -> str:
return f"{conversation_prefix(prompt)}{answer}\n<|end|>"
|