opentau.datasets.grounding.tokenizer_utils

Tokenizer-side support for PaliGemma <loc0000>..`<loc1023>` tokens.

Two cases must be handled, both via the same call:

  1. PaliGemma (`google/paligemma-3b-pt-224`). The 1024 loc strings are already in the SentencePiece vocab at IDs 256000..``257023``. They are NOT, however, registered as added tokens, so the bare HF tokenizer BPE-fragments any <loc0000>-shaped string into seven pieces (['<', 'loc', '0', '0', '0', '0', '>']) instead of matching it as one unit at ID 256000. Calling add_tokens with an AddedToken whose string already exists in the vocab is the documented HF mechanism to promote an existing entry to single-token-match status without reassigning its ID. No new vocab slots are created and no embedding resize is needed.

  2. Gemma 3 (`google/gemma-3-4b-pt`). The strings are absent. The same add_tokens call appends 1024 new IDs at the end of the vocab; the model’s embedding table and tied LM head must be resized to match. The new rows are random-init — they learn from the grounding data on first use. There is no PaliGemma loc-embedding transfer.

The utility below covers both cases idempotently, so it can be wired into every policy __init__ defensively.

Functions

ensure_loc_tokens(tokenizer[, model])

Idempotently make <loc0000>..`<loc1023>` available as single tokens.
opentau.datasets.grounding.tokenizer_utils.ensure_loc_tokens(tokenizer, model=None) → int[source]

Idempotently make <loc0000>..`<loc1023>` available as single tokens.

Always promotes the 1024 loc strings to added/special tokens via tokenizer.add_tokens. For PaliGemma this is a no-op vocab-size-wise: the strings already live at the reserved IDs 256000..``257023``, and the call only flips them into single-token match mode. For Gemma 3 the 1024 strings are appended as new IDs, and the model’s embedding table and tied LM head are resized via model.resize_token_embeddings when a model handle is supplied.

The embedding resize is wrapped in a snapshot/restore of the global torch RNG (CPU + all visible CUDA devices) and re-seeded with a fixed constant inside that block. This guarantees the 1024 new rows are bit-identical across runs regardless of when in construction the helper is called, and leaves the outer RNG state untouched so downstream consumers (the loop seed, dataset shuffler, dropout, etc.) are not perturbed. Without this, construction-time embedding init would couple to the active RNG and silently violate CLAUDE.md hard rule #3 (deterministic seeded reruns).

Safe to call multiple times — once the strings are registered as added tokens, subsequent calls neither grow the vocab nor resize.

Parameters:

  • tokenizer – A HuggingFace PreTrainedTokenizer / PreTrainedTokenizerFast.

  • model – Optional PreTrainedModel whose embeddings should be resized when new IDs are assigned. Pass the top-level VLM (e.g. the Gemma3ForConditionalGeneration / PaliGemmaForConditionalGeneration instance) — HF’s resize_token_embeddings dispatches through get_input_embeddings / set_input_embeddings to the language model and updates the tied LM head as well.

Returns:

The number of NEW IDs appended to the tokenizer vocab. Always 0 for PaliGemma; 1024 on the first call against a fresh Gemma 3 tokenizer; 0 for any subsequent call.