code/cube_tokenizer.py

"""
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)

# 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