Buckets:
| """Strict Visual-Jigsaw state-machine DFA for GRPO format reward.""" | |
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| import torch | |
| from shared.cycles import repeating_period | |
| from shared.protocol import parse_image_scaffold | |
| from uniar import SPECIAL_TOKEN_IDS | |
| from .contracts import ProtocolResult, RLBlock, RLTrajectory, RLTransition, Span | |
| class ProtocolLimits: | |
| require_generated_image: bool = True | |
| min_grid_side: int = 1 | |
| max_grid_side: int = 30 | |
| max_generated_images: int | None = None | |
| max_text_phases: int | None = None | |
| max_generated_visual_tokens: int | None = None | |
| cycle_repeats: int = 3 | |
| cycle_max_period: int = 4 | |
| def __post_init__(self) -> None: | |
| if any( | |
| type(value) is not int or value <= 0 | |
| for value in ( | |
| self.max_grid_side, | |
| self.min_grid_side, | |
| self.cycle_max_period, | |
| ) | |
| ): | |
| raise ValueError("protocol limits must be positive integers") | |
| if type(self.cycle_repeats) is not int or self.cycle_repeats < 2: | |
| raise ValueError("protocol cycle repeats must be an integer of at least two") | |
| if any( | |
| value is not None and (type(value) is not int or value <= 0) | |
| for value in ( | |
| self.max_generated_images, | |
| self.max_text_phases, | |
| self.max_generated_visual_tokens, | |
| ) | |
| ): | |
| raise ValueError("optional protocol limits must be positive integers or null") | |
| if self.min_grid_side > self.max_grid_side: | |
| raise ValueError("protocol grid-side range is empty") | |
| def _spans(positions: set[int]) -> tuple[Span, ...]: | |
| result = [] | |
| start = None | |
| limit = max(positions, default=-1) + 2 | |
| for position in range(limit): | |
| active = position in positions | |
| if active and start is None: | |
| start = position | |
| elif not active and start is not None: | |
| result.append(Span(start, position)) | |
| start = None | |
| return tuple(result) | |
| def _decoded_nonempty(tokenizer, token_id: int) -> bool: | |
| return bool( | |
| tokenizer.decode( | |
| [token_id], | |
| skip_special_tokens=False, | |
| clean_up_tokenization_spaces=False, | |
| ).strip() | |
| ) | |
| def _invalid(reason: str) -> ProtocolResult: | |
| return ProtocolResult(valid=False, failure_reason=reason) | |
| def _text_parts( | |
| ids: list[int], | |
| *, | |
| start: int, | |
| end: int, | |
| tokenizer, | |
| final: bool, | |
| limits: ProtocolLimits, | |
| ): | |
| think = SPECIAL_TOKEN_IDS["<think>"] | |
| think_end = SPECIAL_TOKEN_IDS["</think>"] | |
| marker = SPECIAL_TOKEN_IDS["<image_gen>"] | |
| im_end = SPECIAL_TOKEN_IDS["<|im_end|>"] | |
| reserved = { | |
| SPECIAL_TOKEN_IDS["<|vision_end|>"], | |
| SPECIAL_TOKEN_IDS["<|image_pad|>"], | |
| } | |
| vision_start = SPECIAL_TOKEN_IDS["<|vision_start|>"] | |
| reserved_protocol = { | |
| SPECIAL_TOKEN_IDS["<|im_start|>"], | |
| SPECIAL_TOKEN_IDS["<|video_pad|>"], | |
| } | |
| values = ids[start:end] | |
| if not values: | |
| return None, "empty_text_phase" | |
| if any(token in reserved for token in values): | |
| return None, "student_reserved_visual_token" | |
| if any(token in reserved_protocol for token in values): | |
| return None, "student_reserved_protocol_token" | |
| if final: | |
| if values[-1] != im_end or values.count(im_end) != 1: | |
| return None, "invalid_im_end" | |
| if marker in values or vision_start in values: | |
| return None, "image_marker_in_final_text" | |
| body_end = len(values) - 1 | |
| grid = None | |
| else: | |
| if ( | |
| im_end in values | |
| or values.count(marker) != 1 | |
| or values[-1] != vision_start | |
| or values.count(vision_start) != 1 | |
| ): | |
| return None, "invalid_intermediate_endpoint" | |
| marker_offset = values.index(marker) | |
| decoded = tokenizer.decode( | |
| values[marker_offset:], | |
| skip_special_tokens=False, | |
| clean_up_tokenization_spaces=False, | |
| ) | |
| grid = parse_image_scaffold(decoded) | |
| if grid is None: | |
| return None, "invalid_image_scaffold" | |
| if not all( | |
| limits.min_grid_side <= value <= limits.max_grid_side | |
| for value in grid | |
| ): | |
| return None, "invalid_grid" | |
| body_end = marker_offset | |
| payload_positions: set[int] = set() | |
| final_answer_positions: set[int] = set() | |
| inside = False | |
| thought_start = None | |
| thought_nonempty = False | |
| for offset, token in enumerate(values[:body_end]): | |
| absolute = start + offset | |
| if token == think: | |
| if inside: | |
| return None, "nested_think" | |
| inside = True | |
| thought_start = absolute | |
| thought_nonempty = False | |
| elif token == think_end: | |
| if not inside: | |
| return None, "unbalanced_think" | |
| if not thought_nonempty: | |
| return None, "empty_think" | |
| inside = False | |
| thought_start = None | |
| elif inside: | |
| if token in (marker, im_end): | |
| return None, "special_token_inside_think" | |
| payload_positions.add(absolute) | |
| thought_nonempty |= _decoded_nonempty(tokenizer, token) | |
| elif _decoded_nonempty(tokenizer, token): | |
| if not final: | |
| return None, "unwrapped_intermediate_text" | |
| payload_positions.add(absolute) | |
| final_answer_positions.add(absolute) | |
| if inside or thought_start is not None: | |
| return None, "unbalanced_think" | |
| if not final and not payload_positions: | |
| return None, "missing_reasoning_payload" | |
| if final and not final_answer_positions: | |
| return None, "missing_final_answer_payload" | |
| extent = set(range(start, end)) | |
| payload_spans = _spans(payload_positions) | |
| scaffold_spans = _spans(extent - payload_positions) | |
| answer_spans = _spans(final_answer_positions) | |
| answer = None | |
| if final: | |
| answer_ids = [ids[position] for position in sorted(final_answer_positions)] | |
| answer = tokenizer.decode( | |
| answer_ids, | |
| skip_special_tokens=False, | |
| clean_up_tokenization_spaces=False, | |
| ).strip() | |
| return ( | |
| { | |
| "payload_spans": payload_spans, | |
| "scaffold_spans": scaffold_spans, | |
| "answer_spans": answer_spans, | |
| "answer": answer, | |
| "grid": grid, | |
| "scaffold_start": start + marker_offset if not final else None, | |
| }, | |
| None, | |
| ) | |
| def _transitions(blocks: list[RLBlock]) -> tuple[RLTransition, ...]: | |
| events = [] | |
| for source, target in zip(blocks, blocks[1:], strict=False): | |
| prior = tuple( | |
| span | |
| for block in blocks[: source.index] | |
| for span in block.payload_spans | |
| ) | |
| events.append( | |
| RLTransition( | |
| index=len(events), | |
| direction=( | |
| "text_to_visual" | |
| if source.modality == "text" | |
| else "visual_to_text" | |
| ), | |
| source_block_index=source.index, | |
| target_block_index=target.index, | |
| source_payload_spans=source.payload_spans, | |
| prior_generated_payload_spans=prior, | |
| target_sequence_spans=target.target_sequence_spans, | |
| target_visual_spans=target.target_visual_spans, | |
| ) | |
| ) | |
| return tuple(events) | |
| def parse_protocol( | |
| trajectory: RLTrajectory, tokenizer, limits: ProtocolLimits | |
| ) -> ProtocolResult: | |
| """Parse one raw sampled trajectory without trusting rollout stop metadata.""" | |
| ids = trajectory.input_ids.tolist() | |
| owners = trajectory.token_owners | |
| if any(owner == "input" for owner in owners[trajectory.prefix_length :]): | |
| return _invalid("input_owned_generated_token") | |
| visual_by_start = { | |
| span.start: (index, span, trajectory.visual_code_spans[index]) | |
| for index, span in enumerate(trajectory.visual_sequence_spans) | |
| if span.start >= trajectory.prefix_length | |
| } | |
| blocks: list[RLBlock] = [] | |
| position = trajectory.prefix_length | |
| images = phases = visual_tokens = 0 | |
| cycle_signatures: list[tuple] = [] | |
| final_answer = None | |
| final_spans: tuple[Span, ...] = () | |
| while position < len(ids): | |
| phases += 1 | |
| if ( | |
| limits.max_text_phases is not None | |
| and phases > limits.max_text_phases | |
| ): | |
| return _invalid("text_phase_limit") | |
| endpoint = position | |
| final = False | |
| while endpoint < len(ids): | |
| token = ids[endpoint] | |
| owner = owners[endpoint] | |
| if owner == "controller": | |
| return _invalid("unexpected_controller_token") | |
| if token == SPECIAL_TOKEN_IDS["<|vision_start|>"]: | |
| endpoint += 1 | |
| break | |
| if token == SPECIAL_TOKEN_IDS["<|im_end|>"]: | |
| final = True | |
| endpoint += 1 | |
| break | |
| endpoint += 1 | |
| if endpoint > len(ids) or (endpoint == len(ids) and not final): | |
| return _invalid("truncated_text") | |
| parsed, error = _text_parts( | |
| ids, | |
| start=position, | |
| end=endpoint, | |
| tokenizer=tokenizer, | |
| final=final, | |
| limits=limits, | |
| ) | |
| if error: | |
| return _invalid(error) | |
| blocks.append( | |
| RLBlock( | |
| index=len(blocks), | |
| modality="text", | |
| payload_spans=parsed["payload_spans"], | |
| scaffold_spans=parsed["scaffold_spans"], | |
| target_sequence_spans=(Span(position, endpoint),), | |
| ) | |
| ) | |
| if final: | |
| if endpoint != len(ids) or trajectory.stop_reason != "im_end": | |
| return _invalid("trailing_or_false_endpoint") | |
| final_answer = parsed["answer"] | |
| final_spans = parsed["answer_spans"] | |
| position = endpoint | |
| break | |
| images += 1 | |
| if ( | |
| limits.max_generated_images is not None | |
| and images > limits.max_generated_images | |
| ): | |
| return _invalid("image_limit") | |
| grid_h, grid_w = parsed["grid"] | |
| expected = grid_h * grid_w | |
| visual_tokens += expected | |
| if ( | |
| limits.max_generated_visual_tokens is not None | |
| and visual_tokens > limits.max_generated_visual_tokens | |
| ): | |
| return _invalid("visual_token_limit") | |
| vision_start = endpoint - 1 | |
| if ( | |
| ids[vision_start] != SPECIAL_TOKEN_IDS["<|vision_start|>"] | |
| or owners[vision_start] != "student" | |
| ): | |
| return _invalid("vision_start_ownership") | |
| visual_start = endpoint | |
| aligned = visual_by_start.get(visual_start) | |
| if aligned is None: | |
| return _invalid("missing_visual_alignment") | |
| image_index, sequence_span, code_span = aligned | |
| if sequence_span.length != expected or code_span.length != expected: | |
| return _invalid("wrong_visual_length") | |
| if any( | |
| owners[value] != "student" | |
| or ids[value] != SPECIAL_TOKEN_IDS["<|image_pad|>"] | |
| for value in range(sequence_span.start, sequence_span.end) | |
| ): | |
| return _invalid("invalid_visual_payload") | |
| vision_end = sequence_span.end | |
| if ( | |
| vision_end >= len(ids) | |
| or ids[vision_end] != SPECIAL_TOKEN_IDS["<|vision_end|>"] | |
| or owners[vision_end] != "controller" | |
| ): | |
| return _invalid("vision_end_ownership") | |
| expected_grid = (1, grid_h * 2, grid_w * 2) | |
| if tuple(map(int, trajectory.image_grid_thw[image_index])) != expected_grid: | |
| return _invalid("grid_alignment") | |
| blocks.append( | |
| RLBlock( | |
| index=len(blocks), | |
| modality="visual", | |
| payload_spans=(sequence_span,), | |
| scaffold_spans=(Span(vision_end, vision_end + 1),), | |
| target_sequence_spans=(sequence_span,), | |
| target_visual_spans=(code_span,), | |
| ) | |
| ) | |
| cycle_signatures.append( | |
| ( | |
| tuple(ids[position : parsed["scaffold_start"]]), | |
| tuple(ids[parsed["scaffold_start"] : endpoint]), | |
| (grid_h, grid_w), | |
| ) | |
| ) | |
| if ( | |
| repeating_period( | |
| cycle_signatures, | |
| repeats=limits.cycle_repeats, | |
| max_period=limits.cycle_max_period, | |
| ) | |
| is not None | |
| ): | |
| return _invalid("cycle_detected") | |
| position = vision_end + 1 | |
| if position != len(ids) or final_answer is None: | |
| return _invalid("missing_final_endpoint") | |
| if limits.require_generated_image and images == 0: | |
| return _invalid("required_image_missing") | |
| generated_visual_positions = { | |
| value | |
| for span in trajectory.visual_sequence_spans | |
| if span.start >= trajectory.prefix_length | |
| for value in range(span.start, span.end) | |
| } | |
| expected_visual_actions = tuple(sorted(generated_visual_positions)) | |
| expected_text_actions = tuple( | |
| index | |
| for index, owner in enumerate(owners) | |
| if index >= trajectory.prefix_length | |
| and owner == "student" | |
| and index not in generated_visual_positions | |
| ) | |
| if trajectory.visual_action_positions != expected_visual_actions: | |
| return _invalid("visual_action_coverage") | |
| if trajectory.text_action_positions != expected_text_actions: | |
| return _invalid("text_action_coverage") | |
| return ProtocolResult( | |
| valid=True, | |
| failure_reason=None, | |
| blocks=tuple(blocks), | |
| transitions=_transitions(blocks), | |
| final_answer_text=final_answer, | |
| final_answer_spans=final_spans, | |
| ) | |
| __all__ = ["parse_image_scaffold", "ProtocolLimits", "parse_protocol"] | |
Xet Storage Details
- Size:
- 14.1 kB
- Xet hash:
- 2e762c5307253ad9a24b7f1c6fb9155d5b3299171e38fc3dc91dcea8e1a38490
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.