tiny-cube-value-iter / code /cube_tokenizer.py

Upload folder using huggingface_hub

9eb2e79 verified about 1 month ago

3.84 kB

"""
Tokenizer for the cube task.
Deliberately not a text tokenizer. The task has a tiny closed vocabulary -- six
sticker colours in, eighteen moves out -- so every token is one symbol with one
meaning, and a sequence is a fixed 54-token state followed by a short solution.
A BPE tokenizer would only introduce ambiguity (multi-character merges, casing,
whitespace) into an input that has none.
Sequence layout:
    [BOS] s0 s1 ... s53 [SEP] m0 m1 ... mk [EOS]
Loss is taken only on the move span and [EOS]. The state is the conditioning
prompt -- training the model to predict the scrambled state would spend capacity
learning the distribution of cube states, which is not the task and which the
model gets for free as input anyway.
Both the trainer and the serving shim import this module, so their vocabularies
cannot drift apart. A silent mismatch there would look exactly like a model that
failed to learn.
"""
from typing import List

COLORS = list("URFDLB")
MOVES = [f + s for f in "URFDLB" for s in ("", "'", "2")]
SPECIALS = ["<pad>", "<bos>", "<sep>", "<eos>"]

VOCAB = SPECIALS + COLORS + MOVES
STOI = {t: i for i, t in enumerate(VOCAB)}
ITOS = {i: t for t, i in STOI.items()}

PAD, BOS, SEP, EOS = (STOI[t] for t in SPECIALS)
VOCAB_SIZE = len(VOCAB)

# Colour token ids are contiguous, so a facelet's colour *index* (0..5, as used by
# cube_ops' array form) becomes a token id by adding this. Exposed here rather
# than as a literal in the array-based trainers, for the reason stated above: one
# definition of the encoding, so the two representations cannot drift apart.
# Colour token ids, indexed by the colour *index* (0..5) that cube_ops' array form
# uses. Exposed here rather than recomputed in the array-based trainers, for the
# reason stated above: one definition of the encoding, so the two representations
# cannot drift apart.
#
# These are NOT contiguous from len(SPECIALS), and the reason is a latent wart:
# MOVES contains the bare face turns "U", "R", ... which collide with the colour
# names, and STOI is built by dict comprehension, so the MOVES entries overwrite
# the COLORS ones. Consequences:
#   - a colour and its bare-face move share one embedding (ids 10,13,16,19,22,25),
#   - VOCAB_SIZE is 28 but only 22 ids are distinct, and 4..9 are never emitted.
# Harmless for a value net, which only ever reads state tokens, and a mild
# confound for the policy. Left as-is deliberately: every trained checkpoint
# encodes states this way, so "fixing" it silently invalidates all of them.
# Derive from STOI, never by arithmetic on len(SPECIALS).
COLOR_IDS = [STOI[c] for c in COLORS]

# 54 state tokens + BOS + SEP + a solution (Kociemba stays at or under 22 in the
# half-turn metric, God's Number being 20 for optimal solves) + EOS.
STATE_LEN = 54
MAX_SOLUTION = 26
MAX_SEQ = 1 + STATE_LEN + 1 + MAX_SOLUTION + 1


def encode_state(state: str) -> List[int]:
    """[BOS] + 54 colour tokens + [SEP] -- the prompt the model conditions on."""
    if len(state) != STATE_LEN:
        raise ValueError(f"expected {STATE_LEN} facelets, got {len(state)}")
    return [BOS] + [STOI[c] for c in state] + [SEP]


def encode_solution(moves: List[str]) -> List[int]:
    return [STOI[m] for m in moves] + [EOS]


def decode_solution(ids: List[int]) -> List[str]:
    """Stops at [EOS]; ignores anything that is not a move token."""
    out = []
    for i in ids:
        if i == EOS:
            break
        tok = ITOS.get(int(i))
        if tok in STOI and tok in MOVES:
            out.append(tok)
    return out


def encode_pair(state: str, moves: List[str]):
    """Returns (input_ids, labels) with labels masked to -100 over the prompt."""
    prompt = encode_state(state)
    answer = encode_solution(moves)
    ids = prompt + answer
    labels = [-100] * len(prompt) + answer
    return ids, labels

	"""
	Tokenizer for the cube task.

	Deliberately not a text tokenizer. The task has a tiny closed vocabulary -- six
	sticker colours in, eighteen moves out -- so every token is one symbol with one
	meaning, and a sequence is a fixed 54-token state followed by a short solution.
	A BPE tokenizer would only introduce ambiguity (multi-character merges, casing,
	whitespace) into an input that has none.

	Sequence layout:

	[BOS] s0 s1 ... s53 [SEP] m0 m1 ... mk [EOS]

	Loss is taken only on the move span and [EOS]. The state is the conditioning
	prompt -- training the model to predict the scrambled state would spend capacity
	learning the distribution of cube states, which is not the task and which the
	model gets for free as input anyway.

	Both the trainer and the serving shim import this module, so their vocabularies
	cannot drift apart. A silent mismatch there would look exactly like a model that
	failed to learn.
	"""
	from typing import List

	COLORS = list("URFDLB")
	MOVES = [f + s for f in "URFDLB" for s in ("", "'", "2")]
	SPECIALS = ["<pad>", "<bos>", "<sep>", "<eos>"]

	VOCAB = SPECIALS + COLORS + MOVES
	STOI = {t: i for i, t in enumerate(VOCAB)}
	ITOS = {i: t for t, i in STOI.items()}

	PAD, BOS, SEP, EOS = (STOI[t] for t in SPECIALS)
	VOCAB_SIZE = len(VOCAB)

	# Colour token ids are contiguous, so a facelet's colour index (0..5, as used by
	# cube_ops' array form) becomes a token id by adding this. Exposed here rather
	# than as a literal in the array-based trainers, for the reason stated above: one
	# definition of the encoding, so the two representations cannot drift apart.
	# Colour token ids, indexed by the colour index (0..5) that cube_ops' array form
	# uses. Exposed here rather than recomputed in the array-based trainers, for the
	# reason stated above: one definition of the encoding, so the two representations
	# cannot drift apart.
	#
	# These are NOT contiguous from len(SPECIALS), and the reason is a latent wart:
	# MOVES contains the bare face turns "U", "R", ... which collide with the colour
	# names, and STOI is built by dict comprehension, so the MOVES entries overwrite
	# the COLORS ones. Consequences:
	# - a colour and its bare-face move share one embedding (ids 10,13,16,19,22,25),
	# - VOCAB_SIZE is 28 but only 22 ids are distinct, and 4..9 are never emitted.
	# Harmless for a value net, which only ever reads state tokens, and a mild
	# confound for the policy. Left as-is deliberately: every trained checkpoint
	# encodes states this way, so "fixing" it silently invalidates all of them.
	# Derive from STOI, never by arithmetic on len(SPECIALS).
	COLOR_IDS = [STOI[c] for c in COLORS]

	# 54 state tokens + BOS + SEP + a solution (Kociemba stays at or under 22 in the
	# half-turn metric, God's Number being 20 for optimal solves) + EOS.
	STATE_LEN = 54
	MAX_SOLUTION = 26
	MAX_SEQ = 1 + STATE_LEN + 1 + MAX_SOLUTION + 1


	def encode_state(state: str) -> List[int]:
	"""[BOS] + 54 colour tokens + [SEP] -- the prompt the model conditions on."""
	if len(state) != STATE_LEN:
	raise ValueError(f"expected {STATE_LEN} facelets, got {len(state)}")
	return [BOS] + [STOI[c] for c in state] + [SEP]


	def encode_solution(moves: List[str]) -> List[int]:
	return [STOI[m] for m in moves] + [EOS]


	def decode_solution(ids: List[int]) -> List[str]:
	"""Stops at [EOS]; ignores anything that is not a move token."""
	out = []
	for i in ids:
	if i == EOS:
	break
	tok = ITOS.get(int(i))
	if tok in STOI and tok in MOVES:
	out.append(tok)
	return out


	def encode_pair(state: str, moves: List[str]):
	"""Returns (input_ids, labels) with labels masked to -100 over the prompt."""
	prompt = encode_state(state)
	answer = encode_solution(moves)
	ids = prompt + answer
	labels = [-100] * len(prompt) + answer
	return ids, labels