"""Recover model token IDs from text produced by Gemma's fixed decoder. The black-box server is assumed to expose only ``base_tokenizer.decode(ids)``. Gemma's decoder removes token boundaries, so re-encoding that text is not a true inverse. This module builds the inverse decoding lattice, keeps the minimum-token paths, decodes each path with the alien tokenizer, and ranks the resulting natural-language candidates locally. """ from __future__ import annotations from dataclasses import dataclass import math import re import unicodedata _BYTE_TOKEN = re.compile(r"<0x([0-9A-Fa-f]{2})>") _ALLOWED_FORMAT_CHARS = frozenset({"\u200d", "\ufe0e", "\ufe0f"}) @dataclass(frozen=True) class RecoveryResult: text: str token_ids: tuple[int, ...] minimum_tokens: int candidate_count: int class GemmaInverseSegmenter: """Invert the specific Replace/ByteFallback/Fuse decoder used by Gemma 2.""" def __init__( self, base_tokenizer, alien_tokenizer, *, max_candidates: int = 1_024, max_expansions: int = 50_000, ) -> None: self.base_tokenizer = base_tokenizer self.alien_tokenizer = alien_tokenizer self.max_candidates = max_candidates self.max_expansions = max_expansions self._terminal = object() self._trie: dict = {} self._byte_id_by_value: dict[int, int] = {} self._single_byte_decodings: dict[str, list[tuple[int, ...]]] = {} self._build_decoder_index() def _build_decoder_index(self) -> None: special_ids = set(self.base_tokenizer.all_special_ids) # Sorting by ID makes tie-breaking deterministic across processes. vocab = sorted(self.base_tokenizer.get_vocab().items(), key=lambda item: item[1]) for token, token_id in vocab: byte_match = _BYTE_TOKEN.fullmatch(token) if byte_match: byte_value = int(byte_match.group(1), 16) self._byte_id_by_value[byte_value] = token_id continue if token_id in special_ids: # The server drops these with skip_special_tokens=True. They # therefore cannot, and need not, be reconstructed from text. continue # Gemma's first decoder stage is Replace("▁", " "). The following # Fuse stage concatenates these pieces without retaining boundaries. piece = token.replace("▁", " ") if not piece: continue node = self._trie for character in piece: node = node.setdefault(character, {}) node.setdefault(self._terminal, []).append(token_id) for byte_value, token_id in sorted(self._byte_id_by_value.items()): visible = bytes([byte_value]).decode("utf-8", errors="replace") self._single_byte_decodings.setdefault(visible, []).append((token_id,)) def _byte_arcs(self, character: str) -> list[tuple[int, ...]]: """Return byte-token sequences that can visibly decode to a character.""" arcs = list(self._single_byte_decodings.get(character, ())) utf8 = character.encode("utf-8") if all(byte in self._byte_id_by_value for byte in utf8): arcs.append(tuple(self._byte_id_by_value[byte] for byte in utf8)) return list(dict.fromkeys(arcs)) def _edges(self, text: str) -> list[list[tuple[int, tuple[int, ...]]]]: edges: list[list[tuple[int, tuple[int, ...]]]] = [ [] for _ in range(len(text) + 1) ] for start in range(len(text)): node = self._trie end = start while end < len(text) and text[end] in node: node = node[text[end]] end += 1 for token_id in node.get(self._terminal, ()): edges[start].append((end, (token_id,))) for byte_ids in self._byte_arcs(text[start]): edges[start].append((start + 1, byte_ids)) # A token sequence can reach the same node through duplicate byte # representations. Removing exact duplicate arcs keeps the lattice # smaller without removing any decoded candidate. edges[start] = list(dict.fromkeys(edges[start])) return edges @staticmethod def _minimum_distances( edges: list[list[tuple[int, tuple[int, ...]]]], ) -> tuple[list[int], list[int]]: length = len(edges) - 1 infinity = 10**12 forward = [infinity] * (length + 1) backward = [infinity] * (length + 1) forward[0] = 0 backward[length] = 0 for position in range(length): if forward[position] == infinity: continue for next_position, token_ids in edges[position]: forward[next_position] = min( forward[next_position], forward[position] + len(token_ids), ) for position in range(length - 1, -1, -1): for next_position, token_ids in edges[position]: if backward[next_position] != infinity: backward[position] = min( backward[position], len(token_ids) + backward[next_position], ) return forward, backward def _minimum_candidates( self, text: str, ) -> tuple[list[tuple[tuple[int, ...], str, bool]], int]: edges = self._edges(text) forward, backward = self._minimum_distances(edges) minimum_tokens = forward[-1] if minimum_tokens >= 10**12: return [], 0 candidates_by_text: dict[str, tuple[tuple[int, ...], str, bool]] = {} stack: list[tuple[int, tuple[int, ...]]] = [(0, ())] expansions = 0 while stack and len(candidates_by_text) < self.max_candidates: position, token_ids = stack.pop() expansions += 1 if expansions > self.max_expansions: break if position == len(text): recovered = self.alien_tokenizer.decode( token_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False, ) canonical = tuple( self.alien_tokenizer.encode( recovered, add_special_tokens=False, ) ) == token_ids existing = candidates_by_text.get(recovered) if existing is None or (canonical and not existing[2]): candidates_by_text[recovered] = (token_ids, recovered, canonical) continue shortest_edges = [] for next_position, arc_ids in edges[position]: if ( len(token_ids) + len(arc_ids) + backward[next_position] == minimum_tokens ): shortest_edges.append((next_position, arc_ids)) # Preserve deterministic trie order in depth-first traversal. for next_position, arc_ids in reversed(shortest_edges): stack.append((next_position, token_ids + arc_ids)) return list(candidates_by_text.values()), int(minimum_tokens) @staticmethod def _structural_penalty(text: str) -> float: penalty = 0.0 for character in text: category = unicodedata.category(character) if ( category in {"Cc", "Cf", "Cs", "Co", "Cn"} and character not in {"\n", "\r", "\t"} and character not in _ALLOWED_FORMAT_CHARS ): penalty += 8.0 if character == "\ufffd": penalty += 16.0 for opening, closing in ( ("(", ")"), ("[", "]"), ("{", "}"), ("<", ">"), ): penalty += 0.5 * abs(text.count(opening) - text.count(closing)) return penalty def _bpe_rank(self, text: str) -> tuple[int, float]: """Use only the public tokenizer's vocabulary order as a tie-breaker.""" token_ids = self.base_tokenizer.encode(text, add_special_tokens=False) return len(token_ids), sum(math.log1p(token_id) for token_id in token_ids) def recover(self, server_text: str) -> RecoveryResult: if not server_text: return RecoveryResult("", (), 0, 1) candidates, minimum_tokens = self._minimum_candidates(server_text) if not candidates: raise ValueError("No token segmentation can reproduce the server response.") def penalty(candidate: tuple[tuple[int, ...], str, bool]) -> tuple: _token_ids, recovered, canonical = candidate # No language model or language statistics are used. Prefer a path # that is canonical for the alien tokenizer, then structurally # clean text, then tokens with earlier BPE vocabulary ranks. return ( not canonical, self._structural_penalty(recovered), *self._bpe_rank(recovered), ) best = min(candidates, key=penalty) return RecoveryResult( text=best[1], token_ids=best[0], minimum_tokens=minimum_tokens, candidate_count=len(candidates), )