Add Kilo Bridge for LLM-driven ARC task analysis and ONNX export

trm_solver/kilo_bridge.py

@@ -0,0 +1,363 @@
+"""
+Kilo Bridge — call DeepSeek headless to analyze ARC tasks,
+then drive the NN executor to produce ONNX models.
+
+Usage:
+  python trm_solver/kilo_bridge.py --task 007bbfb7 --render
+
+Pipeline:
+  1. Render ARC task as image (or pass raw grid)
+  2. Call `kilo run` with the image + prompt
+  3. Parse markdown output → TransformSpec
+  4. Create NN executor → export ONNX
+"""
+
+import subprocess
+import json
+import os
+import sys
+import argparse
+import tempfile
+from typing import Optional, Dict, List, Tuple
+import numpy as np
+
+# Add parent to path
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from trm_solver.executor import (
+    TransformSpec, create_transform_nn, export_to_onnx, parse_kilo_output
+)
+
+
+# ─── ARC Task Loading ──────────────────────────────────────────
+
+def load_arc_task(task_file: str) -> Dict:
+    """Load an ARC task from JSON file."""
+    with open(task_file) as f:
+        return json.load(f)
+
+
+def render_grid(grid: List[List[int]], cell_size: int = 20) -> np.ndarray:
+    """Render an ARC grid as an RGB image."""
+    # ARC color palette (0-9)
+    palette = {
+        0: [0, 0, 0],        # black
+        1: [0, 116, 217],    # blue
+        2: [255, 65, 54],    # red
+        3: [46, 204, 64],    # green
+        4: [255, 220, 0],    # yellow
+        5: [170, 170, 170],  # gray
+        6: [240, 18, 190],   # magenta
+        7: [255, 133, 27],   # orange
+        8: [127, 219, 255],  # light blue
+        9: [135, 86, 52],    # brown
+    }
+    
+    h, w = len(grid), len(grid[0])
+    img = np.zeros((h * cell_size, w * cell_size, 3), dtype=np.uint8)
+    
+    for r in range(h):
+        for c in range(w):
+            color = palette.get(grid[r][c], [0, 0, 0])
+            img[r*cell_size:(r+1)*cell_size, c*cell_size:(c+1)*cell_size] = color
+    
+    # Add grid lines
+    img[::cell_size, :] = [64, 64, 64]
+    img[:, ::cell_size] = [64, 64, 64]
+    
+    return img
+
+
+def render_task_grids(task: Dict, cell_size: int = 20) -> np.ndarray:
+    """
+    Render all train pairs of an ARC task as a single image.
+    Layout: train pairs side by side, input above output.
+    """
+    train_pairs = task.get("train", [])
+    n_pairs = len(train_pairs)
+    
+    if n_pairs == 0:
+        return None
+    
+    # Find max dimensions
+    max_h = max(len(p["input"]) for p in train_pairs + [{"input": task.get("test", [{}])[0].get("input", [[0]])}])
+    max_w = max(len(p["input"][0]) for p in train_pairs)
+    
+    pair_imgs = []
+    for pair in train_pairs:
+        inp = render_grid(pair["input"], cell_size)
+        out = render_grid(pair["output"], cell_size)
+        # Stack input above output with small gap
+        gap = np.zeros((cell_size//2, inp.shape[1], 3), dtype=np.uint8)
+        pair_img = np.vstack([inp, gap, out])
+        pair_imgs.append(pair_img)
+    
+    # Pad to same height
+    max_pair_h = max(p.shape[0] for p in pair_imgs)
+    padded = []
+    for p in pair_imgs:
+        if p.shape[0] < max_pair_h:
+            pad = np.zeros((max_pair_h - p.shape[0], p.shape[1], 3), dtype=np.uint8)
+            p = np.vstack([p, pad])
+        # Add separator
+        sep = np.zeros((max_pair_h, cell_size//2, 3), dtype=np.uint8)
+        padded.append(p)
+        padded.append(sep)
+    
+    # Remove last separator
+    if padded:
+        padded = padded[:-1]
+    
+    result = np.hstack(padded)
+    return result
+
+
+# ─── Kilo Interface ────────────────────────────────────────────
+
+KILO_PROMPT_TEMPLATE = """Analyze this ARC-AGI task. The image shows training examples: input grids (above) and their output grids (below) for each example pair.
+
+Identify the transformation rule that maps each input to its output.
+
+Output your analysis in this EXACT format:
+
+## Transform
+name: <transform_name>
+
+## Parameters
+- param1: value1
+- param2: value2
+
+Available transform names:
+- identity: output equals input
+- color_map: per-pixel color remapping (params: color_map=[0,2,1,3,...])
+- flip: horizontal or vertical flip (params: direction="horizontal"|"vertical")
+- transpose: matrix transpose
+- rotate: 90/180/270 rotation (params: k=1|2|3)
+- upscale: nearest-neighbor upscale (params: scale=2|3, output_shape=[H,W])
+- kron_self_similar: Kronecker product with own mask
+- tile_repeat: tile input (params: h_repeat=N, w_repeat=N)
+- concat_patterns: concatenate transformed copies (params: axis="horizontal"|"vertical", operations=["identity","flip_h"])
+- pos_color_lut: position-based color lookup (params: lut={"0,0":4,"1,2":3})
+- spatial_gather: pixel rearrangement (params: gather_map={"0,0":"1,2"})
+- onehot_conv: one-hot convolution (params: kernel_h=3, kernel_w=3)
+- onehot_linear: one-hot linear transform (params: weights=[[...]])
+
+Be precise. Output ONLY the structured format above, no extra text."""
+
+
+def call_kilo_headless(image_path: str, prompt: str = None) -> str:
+    """
+    Call DeepSeek via Kilo headless CLI.
+    
+    Args:
+        image_path: Path to rendered ARC task image
+        prompt: Override default prompt
+    
+    Returns:
+        Markdown output from Kilo/DeepSeek
+    """
+    if prompt is None:
+        prompt = KILO_PROMPT_TEMPLATE
+    
+    cmd = [
+        "kilo", "run",
+        prompt,
+        "--image", image_path,
+        "--format", "default"
+    ]
+    
+    print(f"Running: {' '.join(cmd[:2])} ... [prompt + image]")
+    
+    result = subprocess.run(
+        cmd,
+        capture_output=True,
+        text=True,
+        timeout=120  # 2 min timeout per task
+    )
+    
+    if result.returncode != 0:
+        print(f"Kilo error (stderr): {result.stderr[:500]}")
+        raise RuntimeError(f"Kilo failed with code {result.returncode}")
+    
+    return result.stdout.strip()
+
+
+def call_kilo_via_sdk(image_path: str, prompt: str = None, 
+                      server_url: str = "http://localhost:8765") -> str:
+    """
+    Call DeepSeek via Kilo SDK (tunnel to local server).
+    Use this from Kaggle when connected via tunnel.
+    """
+    try:
+        from kilo_sdk import KiloClient
+    except ImportError:
+        raise ImportError("Install kilo_sdk: pip install kilo-sdk")
+    
+    client = KiloClient(server_url=server_url)
+    
+    if prompt is None:
+        prompt = KILO_PROMPT_TEMPLATE
+    
+    response = client.run(
+        prompt=prompt,
+        image=image_path,
+        format="default"
+    )
+    
+    return response
+
+
+# ─── Main Pipeline ─────────────────────────────────────────────
+
+def process_task(task_file: str, output_dir: str, 
+                 use_sdk: bool = False,
+                 server_url: str = "http://localhost:8765",
+                 render: bool = True) -> Tuple[str, TransformSpec, str]:
+    """
+    Full pipeline for one ARC task:
+    1. Render → 2. Kilo → 3. Parse → 4. NN → 5. ONNX
+    
+    Returns: (task_id, TransformSpec, onnx_path)
+    """
+    task = load_arc_task(task_file)
+    task_id = os.path.splitext(os.path.basename(task_file))[0]
+    
+    # Render task as image
+    if render:
+        img = render_task_grids(task)
+        if img is None:
+            raise ValueError(f"No train pairs in {task_file}")
+        from PIL import Image
+        img_path = os.path.join(output_dir, f"{task_id}_render.png")
+        Image.fromarray(img).save(img_path)
+    else:
+        img_path = task_file  # Assume already rendered
+    
+    # Call Kilo
+    print(f"\n{'='*60}")
+    print(f"Task: {task_id}")
+    print(f"{'='*60}")
+    
+    if use_sdk:
+        md_output = call_kilo_via_sdk(img_path, server_url=server_url)
+    else:
+        md_output = call_kilo_headless(img_path)
+    
+    print(f"\nKilo output:\n{md_output[:500]}...\n")
+    
+    # Parse
+    spec = parse_kilo_output(md_output)
+    print(f"Parsed: transform={spec.name}, params={spec.params}")
+    
+    # Create NN
+    model = create_transform_nn(spec)
+    
+    # Get test input shape
+    test_input = task.get("test", [{}])[0].get("input", [[0]])
+    test_h, test_w = len(test_input), len(test_input[0])
+    
+    # Export ONNX
+    os.makedirs(output_dir, exist_ok=True)
+    onnx_path = os.path.join(output_dir, f"{task_id}.onnx")
+    export_to_onnx(model, (test_h, test_w), onnx_path)
+    
+    # Save spec for reference
+    spec_path = os.path.join(output_dir, f"{task_id}_spec.json")
+    with open(spec_path, 'w') as f:
+        json.dump({"name": spec.name, "params": spec.params}, f, indent=2)
+    
+    return task_id, spec, onnx_path
+
+
+def batch_process(data_dir: str, output_dir: str, 
+                  task_ids: Optional[List[str]] = None,
+                  use_sdk: bool = False,
+                  server_url: str = "http://localhost:8765",
+                  max_tasks: int = None) -> List[Tuple[str, TransformSpec, str]]:
+    """
+    Process multiple ARC tasks.
+    
+    Args:
+        data_dir: Path to directory containing task JSON files
+        output_dir: Where to save ONNX files
+        task_ids: Specific task IDs to process (None = all)
+        use_sdk: Use Kilo SDK instead of CLI
+        server_url: SDK server URL
+        max_tasks: Limit number of tasks
+    """
+    results = []
+    task_files = []
+    
+    for f in sorted(os.listdir(data_dir)):
+        if f.endswith('.json'):
+            tid = f.replace('.json', '')
+            if task_ids is None or tid in task_ids:
+                task_files.append(os.path.join(data_dir, f))
+    
+    if max_tasks:
+        task_files = task_files[:max_tasks]
+    
+    print(f"Processing {len(task_files)} tasks...")
+    
+    for i, tf in enumerate(task_files):
+        try:
+            tid, spec, onnx = process_task(tf, output_dir, use_sdk, server_url)
+            results.append((tid, spec, onnx))
+            print(f"  [{i+1}/{len(task_files)}] ✓ {tid} → {spec.name}")
+        except Exception as e:
+            print(f"  [{i+1}/{len(task_files)}] ✗ {os.path.basename(tf)}: {e}")
+    
+    # Summary
+    print(f"\n{'='*60}")
+    print(f"SUMMARY: {len(results)}/{len(task_files)} tasks processed")
+    print(f"{'='*60}")
+    
+    # Save manifest
+    manifest = {
+        "total": len(results),
+        "tasks": {tid: {"transform": spec.name, "params": spec.params} 
+                  for tid, spec, _ in results}
+    }
+    with open(os.path.join(output_dir, "manifest.json"), 'w') as f:
+        json.dump(manifest, f, indent=2)
+    
+    return results
+
+
+# ─── CLI ───────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Kilo Bridge for ARC-AGI")
+    parser.add_argument("--task", help="Task ID or path to task JSON")
+    parser.add_argument("--data-dir", help="Directory of ARC task JSONs")
+    parser.add_argument("--output-dir", default="onnx_models", help="Output directory")
+    parser.add_argument("--use-sdk", action="store_true", help="Use Kilo SDK (tunnel)")
+    parser.add_argument("--server-url", default="http://localhost:8765")
+    parser.add_argument("--max-tasks", type=int, help="Max tasks to process")
+    parser.add_argument("--task-ids", nargs="*", help="Specific task IDs")
+    parser.add_argument("--no-render", action="store_true", help="Skip rendering")
+    parser.add_argument("--dry-run", action="store_true", 
+                       help="Test without calling Kilo (use dummy output)")
+    
+    args = parser.parse_args()
+    
+    if args.dry_run:
+        # Test with dummy Kilo output
+        spec = TransformSpec(name="kron_self_similar", params={"scale": 3})
+        model = create_transform_nn(spec)
+        os.makedirs(args.output_dir, exist_ok=True)
+        onnx_path = os.path.join(args.output_dir, "test_dry_run.onnx")
+        export_to_onnx(model, (3, 3), onnx_path)
+        print(f"Dry run complete: {onnx_path}")
+        sys.exit(0)
+    
+    if args.task:
+        process_task(args.task, args.output_dir, 
+                    use_sdk=args.use_sdk, server_url=args.server_url,
+                    render=not args.no_render)
+    elif args.data_dir:
+        batch_process(args.data_dir, args.output_dir,
+                     task_ids=args.task_ids, use_sdk=args.use_sdk,
+                     server_url=args.server_url, max_tasks=args.max_tasks)
+    else:
+        parser.print_help()