Spaces:
Running on Zero
Running on Zero
| """Slot labels: the placeholder token, the digit labels, and the bin labels each arm uses. | |
| An arm decides how every number is written: | |
| arm coarse slot bins bin labels | |
| A yes 256 ordinary words | |
| F yes 256 bytes | |
| E no 256 bytes (the blog's main arm) | |
| Bplus no 101 single characters (GPC-1's grid) | |
| C no 256 ordinary words (never learns without a coarse slot) | |
| Bplusw no 101 ordinary words | |
| "Bytes" are ASCII characters followed by two-letter tokens (`ab` ... `zz`); the first 101 are exactly | |
| B+'s single-character labels. Every label is one token right after "x: ". | |
| """ | |
| from dataclasses import dataclass, field | |
| PLACEHOLDER = " ___" | |
| # Gemma 4's ids. configure_tokenizer() replaces them for another model's tokenizer. | |
| PH = 52551 # " ___" | |
| DIGITS = [236771, 236770, 236778, 236800, 236812, 236810, 236825, 236832, 236828, 236819] # "0".."9" | |
| ARMS = {"A": (True, 256, "words"), "F": (True, 256, "bytes"), "E": (False, 256, "bytes"), | |
| "Bplus": (False, 101, "bytes"), "C": (False, 256, "words"), "Bplusw": (False, 101, "words")} | |
| class Arm: | |
| name: str | |
| coarse: bool | |
| bins: int | |
| tokens: str | |
| fine: list = field(default_factory=list) # label ids of every fine (number) slot | |
| def configure_tokenizer(tok): | |
| """Derive the placeholder and digit ids from a non-Gemma tokenizer (e.g. GUI-Owl / Qwen3-VL): | |
| the placeholder is the last token of "x: ___", each digit the last token of "x: <d>".""" | |
| global PH, DIGITS | |
| PH = tok("x: ___", add_special_tokens=False)["input_ids"][-1] | |
| assert tok.decode([PH]) == PLACEHOLDER, f"placeholder is not one token: {tok.decode([PH])!r}" | |
| DIGITS = [tok(f"x: {d}", add_special_tokens=False)["input_ids"][-1] for d in "0123456789"] | |
| assert len(set(DIGITS)) == 10 and [tok.decode([i]).strip() for i in DIGITS] == list("0123456789") | |
| def byte_labels(tok, n=256): | |
| """ASCII characters, then two-letter tokens, each a single token after "x: " (GPC-1-style).""" | |
| ids, seen = [], set() | |
| ctx = tok("x: ___", add_special_tokens=False)["input_ids"][:-1] | |
| codes = [chr(c) for c in range(33, 127)] + [f"{chr(a)}{chr(b)}" for a in range(97, 123) for b in range(97, 123)] | |
| for code in codes: | |
| t = tok("x: " + code, add_special_tokens=False)["input_ids"] | |
| if len(t) == len(ctx) + 1 and t[-1] not in seen: | |
| seen.add(t[-1]) | |
| ids.append(t[-1]) | |
| if len(ids) >= n: | |
| break | |
| return ids | |
| def word_labels(tok, n=256): | |
| """n distinct ordinary lowercase words (3-7 letters), each one token after "x: ", spread evenly | |
| over the vocabulary.""" | |
| ctx = tok("x: ___", add_special_tokens=False)["input_ids"][:-1] | |
| cands = [] | |
| for t, i in sorted(tok.get_vocab().items(), key=lambda kv: kv[1]): | |
| w = t[1:] | |
| if 3000 < i < 60000 and t.startswith("▁") and w.isascii() and w.isalpha() and w.islower() and 3 <= len(w) <= 7: | |
| if tok("x: " + w, add_special_tokens=False)["input_ids"] == ctx + [i]: | |
| cands.append(i) | |
| step = max(1, len(cands) // n) | |
| return [cands[k * step] for k in range(n)] | |
| def make_arm(name, tok): | |
| coarse, bins, tokens = ARMS[name] | |
| ids = word_labels(tok, bins) if tokens == "words" else byte_labels(tok, bins)[:bins] | |
| assert len(set(ids)) == bins, "bin labels are not distinct" | |
| if coarse: | |
| assert not set(ids) & set(DIGITS), "a bin label is also a digit label" | |
| return Arm(name, coarse, bins, tokens, list(ids)) | |
| def word_label_ids(tok, prefix, words): | |
| """The label id of each word as the next token after "<prefix>: " (categorical slots).""" | |
| return [tok(f"{prefix}: {w}", add_special_tokens=False)["input_ids"][-1] for w in words] | |