Upload train_edit_classifier.py with huggingface_hub

train_edit_classifier.py

@@ -0,0 +1,656 @@
+"""
+Train 31-class Edit Operation Classifier — Neuroswarm Tier 2
+
+Pipeline:
+  Code → HueAI → HSL (H,W,3)
+  → Circular hue encoding (sin/cos) → ViT → HybridRegionPooler (DETR)
+  → Delta fusion + profile_delta(33) + oklab_magnitude(1)
+  → Hierarchical classifier → 31 ops
+
+Usage:
+    python train_edit_classifier.py --epochs 50 --batch-size 128 --lr 3e-4
+    python train_edit_classifier.py --device cuda --fp16
+"""
+
+import argparse
+import json
+import math
+import os
+import sys
+import time
+import random
+from pathlib import Path
+from typing import List, Tuple, Dict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+
+sys.path.insert(0, str(Path(__file__).parent))
+
+from models.edit_ops import (
+    PaletteEditOps, EditAction, OpCode, TRAINABLE_OPS, NUM_OPS,
+    OP_TO_IDX, IDX_TO_OP, OP_LEVEL
+)
+from models.edit_classifier import EditOpClassifier, EditOpLoss
+from models.scope_pooler import ScopePooler
+
+
+# ============================================================
+# Synthetic Dataset Generator
+# ============================================================
+
+class EditOpDatasetGenerator:
+    """
+    Generates (before_palette, after_palette, label) triples by
+    applying each of the 31 ops to random palettes.
+
+    This is the bootstrapping approach — generate synthetic pairs
+    to pre-train, then fine-tune on real git diff pairs.
+    """
+
+    START = PaletteEditOps.START_OF_SCOPE
+    END = PaletteEditOps.END_OF_SCOPE
+    NOOP = PaletteEditOps.NOOP
+
+    def __init__(self, palette_h: int = 8, palette_w: int = 32, vocab_size: int = 256):
+        self.H = palette_h
+        self.W = palette_w
+        self.vocab_size = vocab_size
+        self.ops = PaletteEditOps()
+        self.pooler = ScopePooler(hidden_dim=64)
+
+    def _random_region_tokens(self, min_len: int = 3, max_len: int = 12) -> List[int]:
+        """Generate random content tokens (excluding 0, 1, 2)."""
+        length = random.randint(min_len, max_len)
+        return [random.randint(3, self.vocab_size - 1) for _ in range(length)]
+
+    def _make_palette(self, tokens: List[int]) -> Tuple[torch.Tensor, object]:
+        """Create palette and metadata from flat token list."""
+        total = self.H * self.W
+        if len(tokens) < total:
+            tokens = tokens + [self.NOOP] * (total - len(tokens))
+        tokens = tokens[:total]
+
+        palette = torch.tensor([tokens], dtype=torch.long).view(1, self.H, self.W)
+        features = torch.randn(1, self.H, self.W, 64)
+        _, metadata = self.pooler(features, palette)
+        return palette[0], metadata[0]
+
+    def _make_single_region(self) -> Tuple[List[int], int]:
+        """Create a single-region palette token list."""
+        content = self._random_region_tokens(5, 20)
+        tokens = [self.START] + content + [self.END]
+        # Pad
+        total = self.H * self.W
+        tokens += [self.NOOP] * (total - len(tokens))
+        return tokens[:total], len(content)
+
+    def _make_two_regions(self) -> List[int]:
+        """Create two adjacent region token list."""
+        c1 = self._random_region_tokens(3, 10)
+        c2 = self._random_region_tokens(3, 10)
+        tokens = [self.START] + c1 + [self.END, self.START] + c2 + [self.END]
+        total = self.H * self.W
+        tokens += [self.NOOP] * (total - len(tokens))
+        return tokens[:total]
+
+    def _make_nested_scope(self) -> List[int]:
+        """Create nested scope: outer [inner [content] content]."""
+        inner = self._random_region_tokens(3, 8)
+        outer = self._random_region_tokens(2, 5)
+        block_hue = random.choice([20, 24, 28, 32])  # for/if/while/with hues
+        tokens = [self.START] + outer + [self.START, block_hue] + inner + [self.END] + [self.END]
+        total = self.H * self.W
+        tokens += [self.NOOP] * (total - len(tokens))
+        return tokens[:total]
+
+    def _make_func_palette(self) -> List[int]:
+        """Create palette with function def (hue 12) and call (hue 60) for async ops."""
+        content = self._random_region_tokens(3, 8)
+        tokens = [self.START, 12] + content + [60] + self._random_region_tokens(2, 4) + [self.END]
+        total = self.H * self.W
+        tokens += [self.NOOP] * (total - len(tokens))
+        return tokens[:total]
+
+    def generate_pair(self, op: OpCode) -> Tuple[torch.Tensor, torch.Tensor, int]:
+        """
+        Generate a (before, after) palette pair for a specific op.
+
+        Returns:
+            before_hsl: (H, W, 3) float tensor (normalized HSL)
+            after_hsl: (H, W, 3) float tensor (normalized HSL)
+            label: int in [0, 30]
+        """
+        label = OP_TO_IDX[op]
+        max_attempts = 10
+
+        for attempt in range(max_attempts):
+            try:
+                before_palette, action = self._create_op_scenario(op)
+                palette, metadata = self._make_palette(before_palette)
+
+                after_palette, success = self.ops.apply(palette, action, metadata)
+                if not success:
+                    continue
+
+                # Convert int palettes to fake HSL (for now: map token → hue/sat/light)
+                before_hsl = self._palette_to_hsl(palette)
+                after_hsl = self._palette_to_hsl(after_palette)
+
+                return before_hsl, after_hsl, label
+
+            except Exception:
+                continue
+
+        # Fallback: return identical palettes (will be NO_OP-like, model must learn)
+        tokens, _ = self._make_single_region()
+        palette, _ = self._make_palette(tokens)
+        hsl = self._palette_to_hsl(palette)
+        return hsl, hsl, label
+
+    @staticmethod
+    def compute_profile_delta(before_hsl: torch.Tensor, after_hsl: torch.Tensor) -> torch.Tensor:
+        """
+        Compute a 33-dim structural profile delta from HSL tensors.
+
+        Mirrors PaletteStructuralProfile dimensions:
+          [0:10]  Category distribution delta (hue bands)
+          [10:19] Color stats delta (mean/std/entropy of H,S,L)
+          [19:25] Structural metrics delta (scope, density, etc.)
+          [25:33] Spectral alignment delta (placeholder zeros)
+
+        This is an approximation for synthetic data. Real training
+        will use PaletteProfiler.profile_file() on actual source code.
+        """
+        PROFILE_DIM = 33
+        delta = torch.zeros(PROFILE_DIM)
+
+        # Category distribution via hue bands (10 bins, 36° each)
+        before_h = before_hsl[..., 0].flatten()
+        after_h = after_hsl[..., 0].flatten()
+
+        for i in range(10):
+            lo, hi = i / 10.0, (i + 1) / 10.0
+            before_count = ((before_h >= lo) & (before_h < hi)).float().mean()
+            after_count = ((after_h >= lo) & (after_h < hi)).float().mean()
+            delta[i] = after_count - before_count
+
+        # Color stats: mean/std/entropy of H,S,L
+        for ch in range(3):
+            before_ch = before_hsl[..., ch].flatten()
+            after_ch = after_hsl[..., ch].flatten()
+            delta[10 + ch * 3] = after_ch.mean() - before_ch.mean()
+            delta[11 + ch * 3] = after_ch.std() - before_ch.std()
+            # Entropy approximation: histogram entropy
+            before_hist = torch.histc(before_ch, bins=16, min=0, max=1) + 1e-8
+            after_hist = torch.histc(after_ch, bins=16, min=0, max=1) + 1e-8
+            before_ent = -(before_hist / before_hist.sum() * (before_hist / before_hist.sum()).log()).sum()
+            after_ent = -(after_hist / after_hist.sum() * (after_hist / after_hist.sum()).log()).sum()
+            delta[12 + ch * 3] = after_ent - before_ent
+
+        # Structural metrics: scope marker changes, density changes
+        before_s = before_hsl[..., 1].flatten()
+        after_s = after_hsl[..., 1].flatten()
+        # Scope markers have S=1.0 — count them
+        delta[19] = (after_s > 0.95).float().mean() - (before_s > 0.95).float().mean()
+        # Content density (non-zero L)
+        delta[20] = (after_hsl[..., 2] > 0.01).float().mean() - (before_hsl[..., 2] > 0.01).float().mean()
+        # Mean saturation change
+        delta[21] = after_s.mean() - before_s.mean()
+        # Mean lightness change
+        delta[22] = after_hsl[..., 2].flatten().mean() - before_hsl[..., 2].flatten().mean()
+        # Unique hue ratio change
+        before_unique = before_h[before_h > 0].unique().numel() / max(1, (before_h > 0).sum().item())
+        after_unique = after_h[after_h > 0].unique().numel() / max(1, (after_h > 0).sum().item())
+        delta[23] = after_unique - before_unique
+        # Token count change (non-NOOP)
+        delta[24] = (after_hsl[..., 2] > 0.01).float().sum() - (before_hsl[..., 2] > 0.01).float().sum()
+
+        # [25:33] spectral alignment — zeros for synthetic, real data fills these
+        return delta
+
+    def _palette_to_hsl(self, palette: torch.Tensor) -> torch.Tensor:
+        """Convert integer palette to normalized HSL float tensor (H, W, 3)."""
+        H, W = palette.shape
+        hsl = torch.zeros(H, W, 3)
+        flat = palette.flatten().float()
+
+        # Map token values to HSL:
+        # H = (token_value / vocab_size) * 360 → normalized to [0, 1]
+        # S = 0.7 for content, 0.0 for NOOP, 1.0 for scope markers
+        # L = 0.5 for content, 0.1 for scope markers, 0.0 for NOOP
+        for i in range(H * W):
+            h, w = i // W, i % W
+            val = flat[i].item()
+            if val == self.NOOP:
+                hsl[h, w] = torch.tensor([0.0, 0.0, 0.0])
+            elif val == self.START:
+                hsl[h, w] = torch.tensor([0.0, 1.0, 0.1])
+            elif val == self.END:
+                hsl[h, w] = torch.tensor([0.5, 1.0, 0.1])
+            else:
+                hsl[h, w] = torch.tensor([
+                    val / self.vocab_size,
+                    0.7,
+                    0.5
+                ])
+        return hsl
+
+    def _create_op_scenario(self, op: OpCode) -> Tuple[List[int], EditAction]:
+        """Create appropriate palette and EditAction for a given op."""
+
+        # === LEVEL 1: Primitive ===
+        if op == OpCode.DELETE_RANGE:
+            tokens, n = self._make_single_region()
+            i_end = min(random.randint(0, 2), n - 1)
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=i_end, payload_idx=0)
+
+        elif op == OpCode.INSERT_TOKEN:
+            tokens, n = self._make_single_region()
+            pos = random.randint(0, n)
+            payload = random.randint(3, self.vocab_size - 1)
+            return tokens, EditAction(op_id=op, region_id=0, i_start=pos, i_end=-1, payload_idx=payload)
+
+        elif op == OpCode.REPLACE_TOKEN:
+            tokens, n = self._make_single_region()
+            pos = random.randint(0, n - 1)
+            payload = random.randint(3, self.vocab_size - 1)
+            return tokens, EditAction(op_id=op, region_id=0, i_start=pos, i_end=-1, payload_idx=payload)
+
+        elif op == OpCode.SWAP_TOKENS:
+            tokens, n = self._make_single_region()
+            i_start = random.randint(0, max(0, n - 2))
+            i_end = random.randint(i_start + 1, n - 1) if i_start < n - 1 else i_start
+            return tokens, EditAction(op_id=op, region_id=0, i_start=i_start, i_end=i_end, payload_idx=0)
+
+        elif op == OpCode.MOVE_RANGE:
+            tokens = self._make_two_regions()
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=0,
+                                      payload_idx=0, target_region_id=1)
+
+        elif op == OpCode.COPY_RANGE:
+            tokens = self._make_two_regions()
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=0,
+                                      payload_idx=0, target_region_id=1)
+
+        elif op == OpCode.WRAP_SCOPE:
+            tokens, n = self._make_single_region()
+            i_end = min(random.randint(1, 3), n - 1)
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=i_end, payload_idx=0)
+
+        elif op == OpCode.UNWRAP_SCOPE:
+            tokens = self._make_nested_scope()
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=-1, payload_idx=0)
+
+        # === LEVEL 2: Structural ===
+        elif op == OpCode.INDENT:
+            tokens, n = self._make_single_region()
+            i_end = min(2, n - 1)
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=i_end, payload_idx=0)
+
+        elif op == OpCode.DEDENT:
+            tokens = self._make_nested_scope()
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=0, payload_idx=0)
+
+        elif op == OpCode.EXTRACT:
+            tokens, n = self._make_single_region()
+            i_end = min(2, n - 1)
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=i_end, payload_idx=0)
+
+        elif op == OpCode.INLINE:
+            # Need a palette with ref token and source region
+            c1 = self._random_region_tokens(3, 6)
+            c2 = self._random_region_tokens(3, 6)
+            tokens = [self.START, 3] + c1[1:] + [self.END, self.START] + c2 + [self.END]
+            total = self.H * self.W
+            tokens += [self.NOOP] * (total - len(tokens))
+            tokens = tokens[:total]
+            return tokens, EditAction(op_id=op, region_id=1, i_start=0, i_end=-1,
+                                      payload_idx=0, target_region_id=0)
+
+        elif op == OpCode.SPLIT_REGION:
+            tokens, n = self._make_single_region()
+            split_at = max(1, min(n // 2, n - 1))
+            return tokens, EditAction(op_id=op, region_id=0, i_start=split_at, i_end=-1, payload_idx=0)
+
+        elif op == OpCode.MERGE_REGIONS:
+            tokens = self._make_two_regions()
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=-1,
+                                      payload_idx=0, target_region_id=1)
+
+        elif op == OpCode.REORDER:
+            tokens, n = self._make_single_region()
+            i_end = min(3, n - 1)
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=i_end, payload_idx=0)
+
+        elif op == OpCode.NEST_IN_BLOCK:
+            tokens, n = self._make_single_region()
+            i_end = min(2, n - 1)
+            block_hue = random.choice([20, 24, 28])
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=i_end,
+                                      payload_idx=block_hue)
+
+        elif op == OpCode.UNNEST_FROM_BLOCK:
+            tokens = self._make_nested_scope()
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=-1, payload_idx=0)
+
+        elif op == OpCode.HOIST:
+            tokens = self._make_nested_scope()
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=0, payload_idx=0)
+
+        elif op == OpCode.SINK:
+            tokens = self._make_two_regions()
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=0,
+                                      payload_idx=0, target_region_id=1)
+
+        # === LEVEL 3: Semantic ===
+        elif op == OpCode.RENAME:
+            tokens, n = self._make_single_region()
+            pos = random.randint(0, n - 1)
+            payload = random.randint(3, self.vocab_size - 1)
+            return tokens, EditAction(op_id=op, region_id=0, i_start=pos, i_end=-1, payload_idx=payload)
+
+        elif op == OpCode.RETYPE:
+            tokens, n = self._make_single_region()
+            i_end = min(1, n - 1)
+            new_types = [random.randint(3, self.vocab_size - 1) for _ in range(3)]
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=i_end,
+                                      payload_idx=0, payload_tokens=new_types)
+
+        elif op == OpCode.CONVERT_CONSTRUCT:
+            # Use built-in macro pattern
+            content = [20, 220, 220] + self._random_region_tokens(2, 5)
+            tokens = [self.START] + content + [self.END]
+            total = self.H * self.W
+            tokens += [self.NOOP] * (total - len(tokens))
+            return tokens[:total], EditAction(op_id=op, region_id=0, i_start=0, i_end=-1, payload_idx=0)
+
+        elif op == OpCode.SYNC_TO_ASYNC:
+            tokens = self._make_func_palette()
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=-1, payload_idx=0)
+
+        elif op == OpCode.PARAMETERIZE:
+            tokens, n = self._make_single_region()
+            pos = random.randint(0, n - 1)
+            param_hue = random.randint(3, self.vocab_size - 1)
+            return tokens, EditAction(op_id=op, region_id=0, i_start=pos, i_end=-1, payload_idx=param_hue)
+
+        elif op == OpCode.SPECIALIZE:
+            tokens, n = self._make_single_region()
+            i_end = min(1, n - 1)
+            concrete = [random.randint(3, self.vocab_size - 1) for _ in range(3)]
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=i_end,
+                                      payload_idx=0, payload_tokens=concrete)
+
+        elif op == OpCode.GUARD:
+            tokens, n = self._make_single_region()
+            i_end = min(2, n - 1)
+            guard_hue = random.choice([24, 28, 32])
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=i_end,
+                                      payload_idx=guard_hue)
+
+        elif op == OpCode.UNGUARD:
+            tokens = self._make_nested_scope()
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=-1, payload_idx=0)
+
+        elif op == OpCode.SCATTER:
+            tokens, n = self._make_single_region()
+            # Pick 2-3 positions to scatter to
+            positions = random.sample(range(n), min(3, n))
+            payload = random.randint(3, self.vocab_size - 1)
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=-1,
+                                      payload_idx=payload, positions=positions)
+
+        elif op == OpCode.GATHER:
+            tokens, n = self._make_single_region()
+            palette, metadata = self._make_palette(tokens)
+            positions = PaletteEditOps._get_content_positions(palette, metadata, 0)
+            abs_positions = positions[:min(3, len(positions))]
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=-1,
+                                      payload_idx=0, positions=abs_positions)
+
+        elif op == OpCode.MIRROR:
+            tokens = self._make_two_regions()
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0, i_end=0,
+                                      payload_idx=random.randint(3, self.vocab_size - 1),
+                                      target_region_id=1)
+
+        elif op == OpCode.COMPOSE:
+            tokens = self._make_nested_scope()
+            palette, metadata = self._make_palette(tokens)
+            mask = metadata.masks[0]
+            n_positions = mask.sum().item()
+            return tokens, EditAction(op_id=op, region_id=0, i_start=0,
+                                      i_end=max(0, int(n_positions) - 1), payload_idx=0)
+
+        raise ValueError(f"Unknown op: {op}")
+
+
+class EditOpDataset(Dataset):
+    """PyTorch Dataset for edit op classification training."""
+
+    def __init__(self, num_samples: int = 10000, palette_h: int = 8, palette_w: int = 32):
+        self.generator = EditOpDatasetGenerator(palette_h, palette_w)
+        self.num_samples = num_samples
+        self.samples_per_op = num_samples // NUM_OPS
+
+        # Pre-generate balanced dataset with profile deltas
+        self.data: List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]] = []
+        print(f"Generating {num_samples} training pairs ({self.samples_per_op} per op)...")
+        for op in TRAINABLE_OPS:
+            for _ in range(self.samples_per_op):
+                before, after, label = self.generator.generate_pair(op)
+                profile_delta = self.generator.compute_profile_delta(before, after)
+                self.data.append((before, after, profile_delta, label))
+
+        # Shuffle
+        random.shuffle(self.data)
+        print(f"Generated {len(self.data)} pairs across {NUM_OPS} ops")
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        before, after, profile_delta, label = self.data[idx]
+        return before, after, profile_delta, torch.tensor(label, dtype=torch.long)
+
+
+# ============================================================
+# Training Loop
+# ============================================================
+
+def train(args):
+    device = torch.device(args.device)
+    print(f"Device: {device}")
+    print(f"Training {NUM_OPS}-class edit op classifier")
+    print(f"Ops: {[op.name for op in TRAINABLE_OPS]}")
+
+    # Create datasets
+    train_dataset = EditOpDataset(args.train_samples, args.palette_h, args.palette_w)
+    val_dataset = EditOpDataset(args.val_samples, args.palette_h, args.palette_w)
+
+    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True,
+                              num_workers=0, pin_memory=True)
+    val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False,
+                            num_workers=0, pin_memory=True)
+
+    # Model
+    model = EditOpClassifier(
+        hidden_dim=args.hidden_dim,
+        vit_layers=args.vit_layers,
+        vit_heads=args.vit_heads,
+        num_regions=args.num_regions,
+        patch_size=args.patch_size,
+        dropout=args.dropout,
+    ).to(device)
+
+    param_count = sum(p.numel() for p in model.parameters())
+    print(f"Model parameters: {param_count:,}")
+
+    # Loss
+    criterion = EditOpLoss().to(device)
+
+    # Optimizer
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.01)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
+
+    # FP16 support
+    scaler = torch.amp.GradScaler('cuda') if args.fp16 and device.type == 'cuda' else None
+
+    best_val_acc = 0.0
+    save_dir = Path("trained_models")
+    save_dir.mkdir(exist_ok=True)
+
+    for epoch in range(args.epochs):
+        model.train()
+        epoch_metrics = {'loss': 0, 'op_acc': 0, 'level_acc': 0, 'batches': 0}
+        t0 = time.time()
+
+        for before, after, profile_delta, labels in train_loader:
+            before = before.to(device)
+            after = after.to(device)
+            profile_delta = profile_delta.to(device)
+            labels = labels.to(device)
+
+            optimizer.zero_grad()
+
+            if scaler:
+                with torch.amp.autocast('cuda'):
+                    op_logits, level_logits, _ = model(before, after, profile_delta)
+                    loss, metrics = criterion(op_logits, level_logits, labels)
+                scaler.scale(loss).backward()
+                scaler.unscale_(optimizer)
+                nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                scaler.step(optimizer)
+                scaler.update()
+            else:
+                op_logits, level_logits, _ = model(before, after, profile_delta)
+                loss, metrics = criterion(op_logits, level_logits, labels)
+                loss.backward()
+                nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                optimizer.step()
+
+            epoch_metrics['loss'] += metrics['loss']
+            epoch_metrics['op_acc'] += metrics['op_acc']
+            epoch_metrics['level_acc'] += metrics['level_acc']
+            epoch_metrics['batches'] += 1
+
+        scheduler.step()
+
+        n = epoch_metrics['batches']
+        train_loss = epoch_metrics['loss'] / n
+        train_op_acc = epoch_metrics['op_acc'] / n
+        train_level_acc = epoch_metrics['level_acc'] / n
+        elapsed = time.time() - t0
+
+        # Validation
+        model.eval()
+        val_metrics = {'loss': 0, 'op_acc': 0, 'level_acc': 0, 'consistency': 0, 'batches': 0}
+        per_op_correct = {i: 0 for i in range(NUM_OPS)}
+        per_op_total = {i: 0 for i in range(NUM_OPS)}
+
+        with torch.no_grad():
+            for before, after, profile_delta, labels in val_loader:
+                before = before.to(device)
+                after = after.to(device)
+                profile_delta = profile_delta.to(device)
+                labels = labels.to(device)
+
+                op_logits, level_logits, _ = model(before, after, profile_delta)
+                _, metrics = criterion(op_logits, level_logits, labels)
+
+                preds = op_logits.argmax(dim=-1)
+                for pred, label in zip(preds, labels):
+                    l = label.item()
+                    per_op_total[l] += 1
+                    if pred.item() == l:
+                        per_op_correct[l] += 1
+
+                val_metrics['loss'] += metrics['loss']
+                val_metrics['op_acc'] += metrics['op_acc']
+                val_metrics['level_acc'] += metrics['level_acc']
+                val_metrics['consistency'] += metrics['consistency']
+                val_metrics['batches'] += 1
+
+        vn = val_metrics['batches']
+        val_loss = val_metrics['loss'] / vn
+        val_op_acc = val_metrics['op_acc'] / vn
+        val_level_acc = val_metrics['level_acc'] / vn
+        val_consistency = val_metrics['consistency'] / vn
+
+        print(f"Epoch {epoch+1:3d}/{args.epochs} "
+              f"[{elapsed:.1f}s] "
+              f"train: loss={train_loss:.4f} op={train_op_acc:.1%} level={train_level_acc:.1%} | "
+              f"val: loss={val_loss:.4f} op={val_op_acc:.1%} level={val_level_acc:.1%} "
+              f"consist={val_consistency:.1%}")
+
+        # Per-op breakdown every 10 epochs
+        if (epoch + 1) % 10 == 0 or epoch == args.epochs - 1:
+            print("  Per-op accuracy:")
+            for level in ['primitive', 'structural', 'semantic']:
+                ops_in_level = [op for op in TRAINABLE_OPS if OP_LEVEL[op] == level]
+                print(f"    {level.upper()}:")
+                for op in ops_in_level:
+                    idx = OP_TO_IDX[op]
+                    total = per_op_total[idx]
+                    correct = per_op_correct[idx]
+                    acc = correct / total if total > 0 else 0
+                    print(f"      {op.name:25s} {correct:3d}/{total:3d} = {acc:.1%}")
+
+        # Save best
+        if val_op_acc > best_val_acc:
+            best_val_acc = val_op_acc
+            checkpoint = {
+                'epoch': epoch + 1,
+                'model_state': model.state_dict(),
+                'optimizer_state': optimizer.state_dict(),
+                'val_op_acc': val_op_acc,
+                'val_level_acc': val_level_acc,
+                'val_consistency': val_consistency,
+                'args': vars(args),
+                'num_ops': NUM_OPS,
+                'op_names': [op.name for op in TRAINABLE_OPS],
+            }
+            torch.save(checkpoint, save_dir / 'edit_classifier_best.pt')
+            print(f"  -> Saved best model (op_acc={val_op_acc:.1%})")
+
+    # Save final
+    torch.save({
+        'epoch': args.epochs,
+        'model_state': model.state_dict(),
+        'val_op_acc': val_op_acc,
+        'best_val_acc': best_val_acc,
+        'args': vars(args),
+        'num_ops': NUM_OPS,
+    }, save_dir / 'edit_classifier_final.pt')
+
+    print(f"\nTraining complete. Best val accuracy: {best_val_acc:.1%}")
+    return best_val_acc
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train 31-class Edit Op Classifier")
+    parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
+    parser.add_argument('--epochs', type=int, default=50)
+    parser.add_argument('--batch-size', type=int, default=128)
+    parser.add_argument('--lr', type=float, default=3e-4)
+    parser.add_argument('--hidden-dim', type=int, default=256)
+    parser.add_argument('--vit-layers', type=int, default=4)
+    parser.add_argument('--vit-heads', type=int, default=8)
+    parser.add_argument('--num-regions', type=int, default=8)
+    parser.add_argument('--patch-size', type=int, default=4)
+    parser.add_argument('--dropout', type=float, default=0.1)
+    parser.add_argument('--train-samples', type=int, default=31000)
+    parser.add_argument('--val-samples', type=int, default=6200)
+    parser.add_argument('--fp16', action='store_true')
+    parser.add_argument('--palette-h', type=int, default=8)
+    parser.add_argument('--palette-w', type=int, default=32)
+    args = parser.parse_args()
+
+    train(args)
+
+
+if __name__ == '__main__':
+    main()