""" Scope-Aware Pooler Extracts semantic regions from palette using scope markers (0=START, 1=END). Implements exact scope matching via stack-based algorithm. """ import logging import torch import torch.nn as nn from typing import List, Tuple, NamedTuple from dataclasses import dataclass class RegionMetadata(NamedTuple): """ Metadata about detected semantic regions Fields: - masks: BoolTensor[R, H, W] - spatial masks for each region - starts: List[int] - flattened start indices - ends: List[int] - flattened end indices - depths: List[int] - nesting depth of each region - types: List[str] - region type hints """ masks: torch.Tensor starts: List[int] ends: List[int] depths: List[int] types: List[str] class ScopeImbalanceError(Exception): """Raised when scope markers are critically unbalanced""" pass class ScopePooler(nn.Module): """ Extract semantic regions from palette using scope markers This module identifies code scopes (functions, loops, classes, etc.) by matching START_OF_SCOPE (0) and END_OF_SCOPE (1) tokens. Algorithm: 1. Flatten palette to 1D sequence 2. Stack-based matching of scope markers 3. Extract features for each matched region 4. Pool features via mean+max aggregation Edge Cases Handled: - Unbalanced scopes (warning + best-effort matching) - Nested scopes (via stack depth tracking) - No scopes found (fallback to uniform grid) - Empty regions (skip + warning) """ def __init__( self, hidden_dim: int = 768, min_region_size: int = 2, fallback_grid_size: int = 4 ): """ Args: hidden_dim: Feature dimension min_region_size: Minimum tokens per region fallback_grid_size: Grid size when no scopes found """ super().__init__() self.hidden_dim = hidden_dim self.min_region_size = min_region_size self.fallback_grid_size = fallback_grid_size # Learned pooling projection # Concat [mean, max] then project back to hidden_dim self.pool_proj = nn.Linear(hidden_dim * 2, hidden_dim) def forward( self, features: torch.Tensor, # (B, H, W, D) palette: torch.Tensor # (B, H, W) ) -> Tuple[torch.Tensor, List[RegionMetadata]]: """ Extract semantic regions and pool features Args: features: (B, H, W, D) - ViT output features palette: (B, H, W) - palette indices Returns: regions: (B, R, D) - per-region pooled features metadata: List[RegionMetadata] - one per batch item Guarantees: - R >= 1 always (at least one region) - All regions non-empty - Features normalized (unit norm) """ B, H, W, D = features.shape assert palette.shape == (B, H, W), f"Shape mismatch: features{features.shape} vs palette{palette.shape}" assert D == self.hidden_dim, f"Hidden dim mismatch: {D} != {self.hidden_dim}" all_regions = [] all_metadata = [] for b in range(B): feat_b = features[b] # (H, W, D) pal_b = palette[b] # (H, W) # Extract regions for this sample regions_b, meta_b = self._extract_regions_single(feat_b, pal_b, H, W) all_regions.append(regions_b) # (R_b, D) all_metadata.append(meta_b) # Pad to max number of regions in batch max_regions = max(r.shape[0] for r in all_regions) padded_regions = [] for regions_b in all_regions: R_b = regions_b.shape[0] if R_b < max_regions: # Pad with zeros padding = torch.zeros( max_regions - R_b, D, device=regions_b.device, dtype=regions_b.dtype ) regions_b = torch.cat([regions_b, padding], dim=0) padded_regions.append(regions_b) batched_regions = torch.stack(padded_regions, dim=0) # (B, R_max, D) return batched_regions, all_metadata def _extract_regions_single( self, features: torch.Tensor, # (H, W, D) palette: torch.Tensor, # (H, W) H: int, W: int ) -> Tuple[torch.Tensor, RegionMetadata]: """ Extract regions from a single sample Returns: regions: (R, D) - pooled features metadata: RegionMetadata """ # 1. Flatten to sequence seq = palette.flatten() # (H*W,) features_flat = features.view(-1, self.hidden_dim) # (H*W, D) # 2. Match scopes try: scope_pairs, depths = self._match_scopes(seq) except ScopeImbalanceError as e: # Critical error - scopes too broken to recover logging.warning(f"{e}. Using fallback uniform grid.") scope_pairs, depths = self._fallback_uniform_grid(H, W) # 3. Filter invalid regions valid_pairs = [] valid_depths = [] for (start, end), depth in zip(scope_pairs, depths): if (end - start + 1) >= self.min_region_size: valid_pairs.append((start, end)) valid_depths.append(depth) if not valid_pairs: # No valid regions - use full sequence valid_pairs = [(0, H*W - 1)] valid_depths = [0] # 4. Extract features for each region region_features = [] region_masks = [] starts = [] ends = [] for (start, end) in valid_pairs: # Extract features in range region_feat = features_flat[start:end+1] # (L, D) # Pool: mean + max mean_pool = region_feat.mean(dim=0) # (D,) max_pool = region_feat.max(dim=0)[0] # (D,) # Concatenate and project combined = torch.cat([mean_pool, max_pool], dim=0) # (2D,) pooled = self.pool_proj(combined) # (D,) # Normalize pooled = torch.nn.functional.normalize(pooled, dim=0) region_features.append(pooled) # Create mask mask = torch.zeros(H * W, dtype=torch.bool, device=palette.device) mask[start:end+1] = True mask_2d = mask.view(H, W) region_masks.append(mask_2d) starts.append(start) ends.append(end) # Stack regions regions = torch.stack(region_features, dim=0) # (R, D) masks = torch.stack(region_masks, dim=0) # (R, H, W) # Create metadata types = ['scope'] * len(valid_pairs) # Generic type for now metadata = RegionMetadata( masks=masks, starts=starts, ends=ends, depths=valid_depths, types=types ) return regions, metadata def _match_scopes( self, seq: torch.Tensor # (N,) ) -> Tuple[List[Tuple[int, int]], List[int]]: """ Stack-based scope matching Returns: pairs: List of (start_idx, end_idx) tuples depths: List of nesting depths Algorithm: - Maintain stack of open scope indices - When seeing START (0), push index - When seeing END (1), pop and create pair - Track depth = current stack size Edge Cases: - Unmatched START: close at sequence end - Unmatched END: skip with warning - No scopes: return empty list (caller handles) """ START_OF_SCOPE = 0 END_OF_SCOPE = 1 stack = [] # Stack of (index, depth) pairs = [] depths = [] seq_np = seq.cpu().numpy() # Faster iteration for i, token in enumerate(seq_np): if token == START_OF_SCOPE: # Open new scope depth = len(stack) stack.append((i, depth)) elif token == END_OF_SCOPE: # Close scope if stack: start_idx, depth = stack.pop() pairs.append((start_idx, i)) depths.append(depth) else: # Unmatched END - skip logging.warning(f"Unmatched END_OF_SCOPE at position {i}") # Handle unmatched STARTs if stack: logging.warning(f"{len(stack)} unmatched START_OF_SCOPE tokens") # Close them at sequence end seq_len = len(seq_np) for start_idx, depth in stack: pairs.append((start_idx, seq_len - 1)) depths.append(depth) # Validate: check for severe imbalance num_starts = (seq == START_OF_SCOPE).sum().item() num_ends = (seq == END_OF_SCOPE).sum().item() if abs(num_starts - num_ends) > max(num_starts, num_ends) * 0.5: # More than 50% imbalance - critical error raise ScopeImbalanceError( f"Severe scope imbalance: {num_starts} starts vs {num_ends} ends" ) return pairs, depths def _fallback_uniform_grid( self, H: int, W: int ) -> Tuple[List[Tuple[int, int]], List[int]]: """ Fallback when scope matching fails Returns uniform grid of regions Args: H, W: palette dimensions Returns: pairs: List of (start, end) for grid cells depths: All depth=0 (flat) """ total = H * W grid_size = self.fallback_grid_size region_size = total // grid_size pairs = [] for i in range(grid_size): start = i * region_size end = (i + 1) * region_size - 1 if i < grid_size - 1 else total - 1 pairs.append((start, end)) depths = [0] * grid_size return pairs, depths def visualize_regions( self, palette: torch.Tensor, # (H, W) metadata: RegionMetadata ) -> str: """ Generate human-readable visualization of regions Returns: String representation """ H, W = palette.shape output = [] output.append(f"Detected {len(metadata.starts)} regions:") for i, (start, end, depth) in enumerate(zip( metadata.starts, metadata.ends, metadata.depths )): region_size = end - start + 1 indent = " " * depth output.append( f"{indent}Region {i}: [{start:4d}, {end:4d}] " f"(size={region_size:3d}, depth={depth})" ) return "\n".join(output)