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