File size: 3,338 Bytes
0ae2d39 | 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 | #!/usr/bin/env python3
"""Bucket the 8192-token vocab into coarse types.
Used to log the mixing scalar's distribution over token types: knowing that
alpha averages 0.5 overall says nothing, knowing that it sits at 0.8 on digits
and 0.4 on word continuations says a lot about what each tower is doing.
The Mini tokenizer is byte-level BPE (with digits split out by a pre-tokenizer),
so token strings are byte-level encoded and have to be mapped back through the
GPT-2 byte<->unicode table before they can be classified.
"""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
TYPE_NAMES = (
"special",
"whitespace",
"digit",
"word_start",
"word_cont",
"punct",
"other",
)
TYPE_INDEX = {name: i for i, name in enumerate(TYPE_NAMES)}
def _byte_decoder() -> dict[str, int]:
bs = (
list(range(ord("!"), ord("~") + 1))
+ list(range(ord("¡"), ord("¬") + 1))
+ list(range(ord("®"), ord("ÿ") + 1))
)
cs = bs[:]
n = 0
for b in range(256):
if b not in bs:
bs.append(b)
cs.append(256 + n)
n += 1
return {chr(c): b for b, c in zip(bs, cs)}
def _decode(token: str, byte_decoder: dict[str, int]) -> str:
try:
raw = bytes(byte_decoder[ch] for ch in token)
except KeyError:
return token
return raw.decode("utf-8", errors="replace")
def classify(text: str) -> int:
if text == "":
return TYPE_INDEX["other"]
core = text.lstrip(" \t")
leading_space = core != text
if core.strip() == "":
return TYPE_INDEX["whitespace"]
if any(ch.isdigit() for ch in core):
return TYPE_INDEX["digit"]
letters = core.replace("'", "").replace("’", "")
if letters and all(ch.isalpha() for ch in letters):
return TYPE_INDEX["word_start"] if leading_space else TYPE_INDEX["word_cont"]
if all(not ch.isalnum() for ch in core):
return TYPE_INDEX["punct"]
return TYPE_INDEX["other"]
def build_token_type_table(tokenizer_path: Path, vocab_size: int) -> np.ndarray:
"""Returns an int64 array of length vocab_size mapping token id -> type index."""
path = Path(tokenizer_path)
if path.is_dir():
path = path / "tokenizer.json"
with path.open("r", encoding="utf-8") as f:
data = json.load(f)
vocab: dict[str, int] = data["model"]["vocab"]
specials = {entry["content"] for entry in data.get("added_tokens", [])}
byte_decoder = _byte_decoder()
table = np.full(vocab_size, TYPE_INDEX["other"], dtype=np.int64)
for token, idx in vocab.items():
if idx >= vocab_size:
continue
if token in specials:
table[idx] = TYPE_INDEX["special"]
continue
table[idx] = classify(_decode(token, byte_decoder))
return table
if __name__ == "__main__":
import sys
from collections import Counter
tok = Path(
sys.argv[1]
if len(sys.argv) > 1
else "/home/banaxi/Desktop/BananaMind/BananaMind-2/tokenizers/fineweb_edu_first_50gib_8k_digits/tokenizer.json"
)
table = build_token_type_table(tok, 8192)
counts = Counter(table.tolist())
for i, name in enumerate(TYPE_NAMES):
print(f"{name:12s} {counts.get(i, 0):5d} {counts.get(i, 0) / len(table):6.2%}")
|