test(pi0_fast): cover shared tokenizer contract

This commit is contained in:
Pepijn
2026-07-28 11:20:02 +02:00
parent 76171662fb
commit ac197d9ad0
2 changed files with 4 additions and 1 deletions
@@ -62,10 +62,11 @@ def test_pi0_fast_action_tokens_have_no_second_bos():
step.action_tokenizer = lambda _: torch.tensor([4, 9]) step.action_tokenizer = lambda _: torch.tensor([4, 9])
step._paligemma_tokenizer = _FakePaliGemmaTokenizer() step._paligemma_tokenizer = _FakePaliGemmaTokenizer()
tokens, mask = step._tokenize_action(torch.zeros(1, 2, 1)) tokens, mask, code_mask = step._tokenize_action(torch.zeros(1, 2, 1))
mapped = [1000 - 1 - 128 - token for token in (4, 9)] mapped = [1000 - 1 - 128 - token for token in (4, 9)]
assert tokens[0, :6].tolist() == [10, 11, *mapped, 12, 1] assert tokens[0, :6].tolist() == [10, 11, *mapped, 12, 1]
assert code_mask[0].tolist() == [False, False, True, True, False, False, False, False]
assert _FakePaliGemmaTokenizer.bos_token_id not in tokens[0, mask[0]].tolist() assert _FakePaliGemmaTokenizer.bos_token_id not in tokens[0, mask[0]].tolist()
@@ -84,6 +85,7 @@ def test_action_tokenizer_serializes_bos_layout():
"max_action_tokens": 8, "max_action_tokens": 8,
"fast_skip_tokens": 128, "fast_skip_tokens": 128,
"paligemma_tokenizer_name": "paligemma", "paligemma_tokenizer_name": "paligemma",
"allow_truncation": True,
"prepend_bos": False, "prepend_bos": False,
"action_tokenizer_name": "fast", "action_tokenizer_name": "fast",
} }
@@ -104,6 +104,7 @@ def test_action_tokenizer_config_preserves_token_mapping():
"fast_skip_tokens": 64, "fast_skip_tokens": 64,
"paligemma_tokenizer_name": "custom/paligemma", "paligemma_tokenizer_name": "custom/paligemma",
"allow_truncation": False, "allow_truncation": False,
"prepend_bos": True,
"action_tokenizer_name": "custom/fast", "action_tokenizer_name": "custom/fast",
} }