Initial pre-tokenizer and trainer

Author: gbenson

.flake8

@@ -0,0 +1,7 @@
+[flake8]
+exclude = .git,__pycache__,venv*,.venv*,build,dist,.local,.#*,#*,*~
+per-file-ignores =
+    # imported but unused
+    src/dom_tokenizers/**/__init__.py: F401
+    # line too long
+    src/dom_tokenizers/pre_tokenizers/dom_snapshot.py: E501

README.md

@@ -0,0 +1,14 @@
+# DOM tokenizers
+
+HTML DOM-aware tokenizers for Hugging Face language models.
+
+## Setup for development
+
+```sh
+git clone --recursive https://github.com/gbenson/dom-tokenizers.git
+cd dom-tokenizers
+python3 -m venv .venv
+. .venv/bin/activate
+pip install --upgrade pip
+pip install -e .[dev]
+```

pyproject.toml

@@ -0,0 +1,22 @@
+[project]
+name = "dom-tokenizers"
+version = "0.0.1"
+dependencies = [
+    "python-magic",
+    "tokenizers",
+    "transformers",
+]
+
+[project.optional-dependencies]
+dev = [
+    "datasets",
+    "flake8",
+    "pillow",
+]
+
+[project.scripts]
+train-tokenizer = "dom_tokenizers.train:main"
+
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"

src/dom_tokenizers/__init__.py

@@ -0,0 +1 @@
+from .dom_snapshot import DOMSnapshotPreTokenizer

src/dom_tokenizers/pre_tokenizers/__init__.py

@@ -0,0 +1,261 @@
+import json
+import re
+
+from base64 import b64decode
+from collections import defaultdict
+from collections.abc import Iterable
+from functools import cached_property
+from itertools import chain
+from posixpath import commonprefix
+from typing import List
+from xml.dom import Node
+
+import magic
+
+from tokenizers import NormalizedString, PreTokenizedString
+
+
+class DOMSnapshotPreTokenizer:
+    """Pre-tokenizer that consumes JSON-serialized DOM snapshots
+    and emits tokenized representations of the snapshotted DOMs.
+    """
+    bos_token = "[BOS]"        # beginning of sequence
+    eos_token = "[EOS]"        # end of sequence
+    sep_token = "[SEP]"        # separator between documents
+    elem_token = "[TAG]"       # beginning of element name
+    attr_token = "[ATTR]"      # beginning of attribute
+    comm_token = "[COMMENT]"   # beginning of comment
+    base64_token = "[BASE64]"  # beginning of some base64
+    long_token = "[LONG]"      # elided long token
+
+    @property
+    def special_tokens(self):
+        return [
+            value
+            for attr, value in self.__class__.__dict__.items()
+            if attr.endswith("token")
+        ]
+
+    def pre_tokenize(self, pretok: PreTokenizedString):
+        """Pre-tokenize a :class:`~tokenizers.PyPreTokenizedString` in-place.
+        """
+        pretok.split(self._split_json)
+
+    def _split_json(self, i: int, s: NormalizedString) -> List[NormalizedString]:
+        snapshot = json.loads(s.normalized)
+        return list(chain.from_iterable(self._split_serialized(snapshot)))
+
+    def _split_serialized(self, snapshot: dict) -> Iterable[List[NormalizedString]]:
+        emitter = TokenEmitter(self, snapshot)
+        elem_token = [NormalizedString(self.elem_token)]
+        attr_token = [NormalizedString(self.attr_token)]
+
+        for document_index, document in enumerate(snapshot["documents"]):
+            token = self.bos_token if document_index == 0 else self.sep_token
+            yield [NormalizedString(token)]
+
+            nodes = document["nodes"]
+            for node_index, node_values in enumerate(zip(
+                    nodes["nodeType"],
+                    nodes["nodeName"],
+                    nodes["nodeValue"],
+                    nodes["attributes"])):
+                node_type, name_index, value_index, attr_indexes = node_values
+
+                match node_type:
+                    case Node.ELEMENT_NODE:
+                        yield elem_token
+                        yield emitter.emit(name_index)
+                        for attr_index in range(0, len(attr_indexes), 2):
+                            yield attr_token
+                            yield emitter.emit(attr_indexes[attr_index])
+                            yield emitter.emit(attr_indexes[attr_index + 1])
+
+                    case Node.TEXT_NODE:
+                        yield emitter.emit(value_index)
+
+                    case Node.COMMENT_NODE:
+                        yield [NormalizedString(self.comm_token)]
+                        yield emitter.emit(value_index)
+
+        yield [NormalizedString(self.eos_token)]
+
+
+_B64_RE_S = r"(?:[A-Za-z0-9+/]{4}){"
+_B64_RE_E = r",}(?:[A-Za-z0-9+/]{3}=|[A-Za-z0-9+/]{2}==)?"
+
+
+def base64_matcher(min_encoded_len=24):
+    min_groups, extra = divmod(min_encoded_len, 4)
+    if extra:
+        min_groups += 1
+    return re.compile(f"{_B64_RE_S}{min_groups}{_B64_RE_E}")
+
+
+class TokenEmitter:
+    MAXWORDLEN = 32
+    WORD_RE = re.compile(
+        r"[a-z0-9]+(?:[a-z0-9']*[a-z0-9])?")  # XXX English only :(
+    ESCAPED_RE = re.compile(
+        r"((?:%|\\x|\\u[0-9a-f]{2})[0-9a-f]{2})", re.I)
+    HEX_RE = re.compile(r"^(?:0x|[0-9a-f]{2})[0-9a-f]{6,}$")
+    DIGIT_RE = re.compile(r"\d")
+    URLISH_RE = re.compile(r"(?:[a-z]+|[0-9a-f]+|[A-Z0-9]+)")
+    SHORTEST_URLISH = 16
+    LONGEST_PHITEST = 85
+    BASE64_RE = base64_matcher()
+    B64_PNG_RE = re.compile(r"iVBORw0KGg[o-r]")
+    XML_HDR_RE = re.compile(r"<([a-z]{3,})\s+[a-z]+")
+
+    def __init__(self, pretokenizer: DOMSnapshotPreTokenizer, snapshot: dict):
+        self._pt = pretokenizer
+        self._strings = snapshot["strings"]
+        self._tokens = {}
+
+    @cached_property
+    def base64_token(self):
+        return self._pt.base64_token
+
+    @cached_property
+    def long_token(self):
+        return self._pt.long_token
+
+    def emit(self, string_index: int) -> Iterable[NormalizedString]:
+        """Emit tokens for one string in a DOM snapshot's string table.
+
+        It splits on any non-alphanumeric character, but also tries
+        to detect (and recurse into) base64-encoded date, of which
+        there's a lot in just the 295 `interesting-dom-snapshots`.
+        (Not dealing with base64 results in a whole load of "words"
+        which are just fragments of base64.  It isn't easy though,
+        lots of regular text is valid base64, we have to sniff.)
+        """
+        if string_index < 0:
+            return []
+        tokens = self._tokens.get(string_index)
+        if tokens is not None:
+            return tokens
+        tokens = [
+            NormalizedString(token)
+            for token in self._postprocess(
+                    chain.from_iterable(
+                        self._split(
+                            self._preprocess(
+                                self._strings[string_index]))))
+        ]
+        self._tokens[string_index] = tokens
+        return tokens
+
+    def _preprocess(self, text):
+        return "".join(
+            self._unescape_char(s) if i & 1 else s
+            for i, s in enumerate(self.ESCAPED_RE.split(text))
+        )
+
+    def _unescape_char(self, escaped):
+        if escaped[0] == "%":
+            escaped = "\\x" + escaped[1:]
+        return eval(f'"{escaped}"')
+
+    def _split(self, text):
+        while text:
+            match = self.BASE64_RE.search(text)
+            if match is not None:
+                start, limit = match.span()
+            else:
+                start = limit = len(text)
+            if start > 0:
+                yield self._split_words(text[:start])
+            if limit > start:
+                encoded = text[start:limit]
+                matched = self._match_urlish_base64(encoded)
+                if matched is not None:
+                    limit = start + len(matched)
+                    yield self._split_words(text[start:limit])
+                else:
+                    yield self._enter_base64(encoded)
+            if limit == len(text):
+                break
+            text = text[limit:]
+
+    def _split_words(self, text):
+        return self.WORD_RE.findall(text.lower())
+
+    def _match_urlish_base64(self, encoded):
+        urlish = "/".join(self.URLISH_RE.findall(encoded))
+        result = commonprefix((encoded, urlish))
+        if len(result) < self.SHORTEST_URLISH:
+            return None
+        return result
+
+    def _enter_base64(self, encoded):
+        # Lots of false-positives here, try sniffing
+        if self.B64_PNG_RE.match(encoded):
+            return [self.base64_token, "png"]
+        data = b64decode(encoded)
+        try:
+            text = data.decode("utf-8")
+        except UnicodeDecodeError:
+            text = None
+        if text is not None:
+            return self._enter_base64_utf8(text)
+        return self._enter_base64_binary(data, encoded)
+
+    def _enter_base64_utf8(self, text):
+        # XXX recurse??
+        match = self.XML_HDR_RE.match(text)
+        if match is not None:
+            if match.group(1) == "svg":
+                return [self.base64_token, "svg"]
+            return [self.base64_token, "xml"]
+        try:
+            _ = json.loads(text)
+            return [self.base64_token, "json"]
+        except json.JSONDecodeError:
+            pass
+        return [self.base64_token, "utf", "8"]
+
+    def _enter_base64_binary(self, data, encoded):
+        # Not out of false-positive territory yet
+        full_magic = magic.from_buffer(data)
+        easy_magic = full_magic.split(maxsplit=1)[0]
+        if easy_magic in {"GIF", "zlib", "JPEG"}:
+            return [self.base64_token, easy_magic.lower()]
+        if " Web/P image" in full_magic:
+            return [self.base64_token, "webp"]
+        if full_magic.startswith("Web Open Font Format"):
+            return [self.base64_token, "woff"]
+        if len(encoded) > self.LONGEST_PHITEST:
+            return [self.base64_token]
+        # phi test for monoalphabeticity
+        hist = defaultdict(int)
+        for symbol in encoded:
+            hist[symbol] += 1
+        phi_o = sum(freq * (freq - 1) for freq in hist.values())
+        N = len(encoded)
+        phi_r = N * (N - 1) / 64
+        # non-standard comparison (observed phi > twice random)
+        if phi_o > phi_r * 2:
+            return self._split_words(encoded)
+        return [self.base64_token]
+
+    def _postprocess(self, tokens: Iterable[str]) -> Iterable[str]:
+        for token in tokens:
+            if self.HEX_RE.match(token):
+                yield self.long_token
+                try:
+                    _ = int(token)
+                except ValueError:
+                    yield "hex"
+                yield "digits"
+                continue
+
+            if len(token) <= self.MAXWORDLEN:
+                yield token
+                continue
+
+            yield self.long_token
+            if self.DIGIT_RE.search(token):
+                yield "alphanumeric"
+            else:
+                yield "alphabetic"

src/dom_tokenizers/pre_tokenizers/dom_snapshot.py

@@ -0,0 +1,89 @@
+import json
+import warnings
+
+from datasets import load_dataset
+from tokenizers.pre_tokenizers import PreTokenizer, WhitespaceSplit
+from transformers import AutoTokenizer
+
+from .pre_tokenizers import DOMSnapshotPreTokenizer
+
+FULL_DATASET = "gbenson/webui-dom-snapshots"
+TEST_DATASET = "gbenson/interesting-dom-snapshots"
+
+
+def train_tokenizer(
+        *args,
+        training_dataset=None,
+        base_tokenizer="bert-base-uncased",
+        vocab_size=1024,  # XXX including all tokens and alphabet
+        **kwargs):
+    """
+    XXX
+    base_tokenizer
+    all other args passed to load_dataset for XXX...
+    """
+
+    # Load the training data we'll train our new tokenizer with.
+    if training_dataset is None:
+        training_dataset = load_dataset(*args, **kwargs)
+
+    # Create the base tokenizer we'll train our new tokenizer from.
+    if isinstance(base_tokenizer, str):
+        base_tokenizer = AutoTokenizer.from_pretrained(base_tokenizer)
+
+    # Create the custom pretokenizer our new tokenizer will use.
+    new_pretokenizer = DOMSnapshotPreTokenizer()
+
+    # List the custom special tokens that need adding to our tokenizer.
+    new_special_tokens = [
+        special_token
+        for special_token in new_pretokenizer.special_tokens
+        if base_tokenizer.tokenize(special_token) != [special_token]
+    ]
+
+    # It's not possible to train using a custom pre-tokenizer, the Rust
+    # code raises "Exception: Custom PreTokenizer cannot be serialized"
+    # (see https://github.com/huggingface/tokenizers/issues/269) so we
+    # have to run our pre-tokenizer manually, then join its output with
+    # whitespace and hope the regular pretokenizer takes it back apart
+    # how we need it to.
+
+    base_tokenizer.backend_tokenizer.pre_tokenizer = WhitespaceSplit()
+    base_pretokenizer = base_tokenizer.backend_tokenizer.pre_tokenizer
+    new_pretokenizer = PreTokenizer.custom(new_pretokenizer)
+
+    def futz_input(real_input):
+        pretokenized = new_pretokenizer.pre_tokenize_str(real_input)
+        want_tokens = [token for token, offsets in pretokenized]
+        futzed_input = " ".join(want_tokens)
+        pretokenized = base_pretokenizer.pre_tokenize_str(futzed_input)
+        got_tokens = [token for token, offsets in pretokenized]
+        assert got_tokens == want_tokens
+        return futzed_input
+
+    def get_training_corpus():
+        for row in training_dataset:
+            yield futz_input(json.dumps(row["dom_snapshot"]))
+
+    # Train the new tokenizer.
+    new_tokenizer = base_tokenizer.train_new_from_iterator(
+        text_iterator=get_training_corpus(),
+        vocab_size=vocab_size,
+        new_special_tokens=new_special_tokens,
+        length=len(training_dataset),  # used for progress tracking
+        show_progress=True,
+    )
+
+    return new_tokenizer
+
+
+def main(save_directory="pretrained", use_full_dataset=False):
+    warnings.filterwarnings("ignore", message=r".*resume_download.*")
+
+    if use_full_dataset:
+        dataset, kwargs = FULL_DATASET, dict(streaming=True)
+    else:
+        dataset, kwargs = TEST_DATASET, {}
+
+    tokenizer = train_tokenizer(dataset, split="train", **kwargs)
+    tokenizer.save_pretrained(save_directory)