Buckets:
| import unittest | |
| from dataclasses import replace | |
| import torch | |
| from rl.contracts import RLTrajectory, Span | |
| from rl.protocol import parse_image_scaffold, ProtocolLimits, parse_protocol | |
| from uniar import SPECIAL_TOKEN_IDS | |
| class _Tokenizer: | |
| pieces = { | |
| SPECIAL_TOKEN_IDS["<think>"]: "<think>", | |
| SPECIAL_TOKEN_IDS["</think>"]: "</think>", | |
| SPECIAL_TOKEN_IDS["<image_gen>"]: "<image_gen>", | |
| SPECIAL_TOKEN_IDS["<|vision_start|>"]: "<|vision_start|>", | |
| SPECIAL_TOKEN_IDS["<|vision_end|>"]: "<|vision_end|>", | |
| SPECIAL_TOKEN_IDS["<|image_pad|>"]: "<|image_pad|>", | |
| SPECIAL_TOKEN_IDS["<|im_end|>"]: "<|im_end|>", | |
| 10: "reason", | |
| 20: " generate", | |
| 21: " image", | |
| 22: " ", | |
| 23: "4", | |
| 24: "4", | |
| 25: "5", | |
| 26: "1", | |
| 27: "0", | |
| 30: "Final answer: ", | |
| 31: "A", | |
| 32: "unwrapped", | |
| } | |
| def decode(self, ids, **_kwargs): | |
| return "".join(self.pieces.get(int(value), "?") for value in ids) | |
| def _trajectory(*, unwrapped=False): | |
| token = SPECIAL_TOKEN_IDS | |
| text = [ | |
| token["<think>"], | |
| 10, | |
| token["</think>"], | |
| *([32] if unwrapped else []), | |
| token["<image_gen>"], | |
| 20, | |
| 21, | |
| 22, | |
| 23, | |
| 22, | |
| 24, | |
| token["<|vision_start|>"], | |
| ] | |
| visual_payload = [token["<|image_pad|>"]] * 16 | |
| visual = [*visual_payload, token["<|vision_end|>"]] | |
| final = [30, 31, token["<|im_end|>"]] | |
| ids = torch.tensor([100, *text, *visual, *final]) | |
| visual_start = 1 + len(text) | |
| return RLTrajectory( | |
| sample_id="fixture", | |
| sample_index=7, | |
| group_index=0, | |
| input_ids=ids, | |
| prefix_length=1, | |
| token_owners=( | |
| "input", | |
| *("student" for _ in text), | |
| *("student" for _ in visual_payload), | |
| "controller", | |
| *("student" for _ in final), | |
| ), | |
| sampled_visual_codes=torch.tensor( | |
| [[0, 1], [1, 0], [0, 0], [1, 1]] * 4, dtype=torch.uint8 | |
| ), | |
| image_grid_thw=torch.tensor([[1, 8, 8]]), | |
| visual_sequence_spans=(Span(visual_start, visual_start + 16),), | |
| visual_code_spans=(Span(0, 16),), | |
| text_action_positions=tuple( | |
| index | |
| for index, owner in enumerate( | |
| ( | |
| "input", | |
| *("student" for _ in text), | |
| *("student" for _ in visual_payload), | |
| "controller", | |
| *("student" for _ in final), | |
| ) | |
| ) | |
| if owner == "student" and int(ids[index]) != token["<|image_pad|>"] | |
| ), | |
| text_old_log_probs=torch.full((len(text) + len(final),), -1.0), | |
| visual_action_positions=tuple(range(visual_start, visual_start + 16)), | |
| visual_old_factor_log_probs=torch.full((16, 2), -0.5), | |
| stop_reason="im_end", | |
| ) | |
| class RLProtocolTest(unittest.TestCase): | |
| def test_image_scaffold_parser_is_the_single_grid_grammar(self): | |
| self.assertEqual( | |
| parse_image_scaffold( | |
| "<image_gen> generate image 12 7<|vision_start|>" | |
| ), | |
| (12, 7), | |
| ) | |
| self.assertIsNone( | |
| parse_image_scaffold("<image_gen> generate image 12 7") | |
| ) | |
| self.assertIsNone( | |
| parse_image_scaffold( | |
| "<image_gen> generate image " | |
| + "9" * 5_000 | |
| + " 7<|vision_start|>" | |
| ) | |
| ) | |
| for noncanonical in ( | |
| "<image_gen> image 12 7<|vision_start|>", | |
| "<image_gen> GeNeRaTe image 12 7<|vision_start|>", | |
| "<image_gen> generate\nimage 12 7<|vision_start|>", | |
| "<image_gen> generate image 12 7<|vision_start|>", | |
| "<image_gen>generate image 12 7<|vision_start|>", | |
| "<image_gen> generate image 012 7<|vision_start|>", | |
| ): | |
| with self.subTest(scaffold=noncanonical): | |
| self.assertIsNone(parse_image_scaffold(noncanonical)) | |
| def test_valid_visual_jigsaw_trajectory_exposes_all_autonomous_blocks(self): | |
| result = parse_protocol( | |
| _trajectory(), | |
| _Tokenizer(), | |
| ProtocolLimits(require_generated_image=True), | |
| ) | |
| self.assertTrue(result.valid) | |
| self.assertEqual(result.failure_reason, None) | |
| self.assertEqual([block.modality for block in result.blocks], ["text", "visual", "text"]) | |
| self.assertEqual( | |
| [event.direction for event in result.transitions], | |
| ["text_to_visual", "visual_to_text"], | |
| ) | |
| self.assertEqual(result.final_answer_text, "Final answer: A") | |
| vision_start = SPECIAL_TOKEN_IDS["<|vision_start|>"] | |
| vision_start_position = int( | |
| (result.blocks[0].target_sequence_spans[0].end - 1) | |
| ) | |
| self.assertEqual(int(_trajectory().input_ids[vision_start_position]), vision_start) | |
| self.assertEqual( | |
| _trajectory().token_owners[vision_start_position], "student" | |
| ) | |
| self.assertIn(vision_start_position, _trajectory().text_action_positions) | |
| self.assertTrue( | |
| any( | |
| span.start <= vision_start_position < span.end | |
| for span in result.blocks[0].scaffold_spans | |
| ) | |
| ) | |
| self.assertFalse( | |
| any( | |
| span.start <= vision_start_position < span.end | |
| for span in result.blocks[1].scaffold_spans | |
| ) | |
| ) | |
| covered = [ | |
| position | |
| for block in result.blocks | |
| for span in (*block.payload_spans, *block.scaffold_spans) | |
| for position in range(span.start, span.end) | |
| ] | |
| self.assertEqual( | |
| sorted(covered), | |
| list(range(_trajectory().prefix_length, _trajectory().input_ids.numel())), | |
| ) | |
| self.assertEqual(len(covered), len(set(covered))) | |
| def test_unwrapped_intermediate_reasoning_fails_the_rl_specific_dfa(self): | |
| result = parse_protocol( | |
| _trajectory(unwrapped=True), | |
| _Tokenizer(), | |
| ProtocolLimits(require_generated_image=True), | |
| ) | |
| self.assertFalse(result.valid) | |
| self.assertEqual(result.failure_reason, "unwrapped_intermediate_text") | |
| def test_empty_nested_or_reserved_reasoning_tokens_fail_closed(self): | |
| base = _trajectory() | |
| reason_position = 2 | |
| cases = ( | |
| (22, "empty_think"), | |
| (SPECIAL_TOKEN_IDS["<think>"], "nested_think"), | |
| ( | |
| SPECIAL_TOKEN_IDS["<|vision_start|>"], | |
| "invalid_intermediate_endpoint", | |
| ), | |
| ( | |
| SPECIAL_TOKEN_IDS["<|im_start|>"], | |
| "student_reserved_protocol_token", | |
| ), | |
| ( | |
| SPECIAL_TOKEN_IDS["<|video_pad|>"], | |
| "student_reserved_protocol_token", | |
| ), | |
| ) | |
| for token_id, reason in cases: | |
| with self.subTest(reason=reason): | |
| ids = base.input_ids.clone() | |
| ids[reason_position] = token_id | |
| result = parse_protocol( | |
| replace(base, input_ids=ids), | |
| _Tokenizer(), | |
| ProtocolLimits(), | |
| ) | |
| self.assertFalse(result.valid) | |
| self.assertEqual(result.failure_reason, reason) | |
| def test_grid_visual_length_and_final_payload_are_independently_checked(self): | |
| base = _trajectory() | |
| scaffold = base.input_ids.clone() | |
| scaffold[10] = 25 | |
| wrong_visual = parse_protocol( | |
| replace(base, input_ids=scaffold), | |
| _Tokenizer(), | |
| ProtocolLimits(), | |
| ) | |
| self.assertFalse(wrong_visual.valid) | |
| self.assertEqual(wrong_visual.failure_reason, "wrong_visual_length") | |
| no_answer_ids = base.input_ids.clone() | |
| no_answer_ids[-3:-1] = 22 | |
| no_answer = parse_protocol( | |
| replace(base, input_ids=no_answer_ids), | |
| _Tokenizer(), | |
| ProtocolLimits(), | |
| ) | |
| self.assertFalse(no_answer.valid) | |
| self.assertEqual(no_answer.failure_reason, "missing_final_answer_payload") | |
| def test_action_coverage_excludes_controller_but_covers_every_student_action(self): | |
| base = _trajectory() | |
| missing = replace( | |
| base, | |
| text_action_positions=base.text_action_positions[:-1], | |
| text_old_log_probs=base.text_old_log_probs[:-1], | |
| ) | |
| result = parse_protocol(missing, _Tokenizer(), ProtocolLimits()) | |
| self.assertFalse(result.valid) | |
| self.assertEqual(result.failure_reason, "text_action_coverage") | |
| def test_controller_owned_vision_start_is_rejected(self): | |
| base = _trajectory() | |
| vision_start = SPECIAL_TOKEN_IDS["<|vision_start|>"] | |
| position = int( | |
| torch.nonzero(base.input_ids == vision_start, as_tuple=False)[0].item() | |
| ) | |
| owners = list(base.token_owners) | |
| owners[position] = "controller" | |
| action_offset = base.text_action_positions.index(position) | |
| text_positions = tuple( | |
| value for value in base.text_action_positions if value != position | |
| ) | |
| old_log_probs = torch.cat( | |
| ( | |
| base.text_old_log_probs[:action_offset], | |
| base.text_old_log_probs[action_offset + 1 :], | |
| ) | |
| ) | |
| result = parse_protocol( | |
| replace( | |
| base, | |
| token_owners=tuple(owners), | |
| text_action_positions=text_positions, | |
| text_old_log_probs=old_log_probs, | |
| ), | |
| _Tokenizer(), | |
| ProtocolLimits(), | |
| ) | |
| self.assertFalse(result.valid) | |
| self.assertEqual(result.failure_reason, "unexpected_controller_token") | |
| if __name__ == "__main__": | |
| unittest.main() | |
Xet Storage Details
- Size:
- 9.96 kB
- Xet hash:
- a71630cd237b3432ae17f54523bb56236db37fcefac96f1907b7ef73ad299a7f
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.