| 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|>"
|
|
|
