own-solver/neurogolf_solver/classify_tasks.py

#!/usr/bin/env python3
"""
ARC-AGI Task Classifier — Routes tasks to NeuroGolf solvers via DeepSeek API.
Output: JSON mapping task_id -> ordered solver list to try first.
The LLM call is OFFLINE (model generation time only). Zero ONNX cost.

Usage on Kaggle:
  python -m neurogolf_solver.classify_tasks

Usage locally:
  python -m neurogolf_solver.classify_tasks --data_dir ARC-AGI/data/training/
"""

import json, os, glob, time, argparse

# --- Solver names matching solver_registry.py ---
SOLVER_NAMES = [
    "identity", "constant", "color_map", "transpose", "flip", "rotate",
    "shift", "tile", "upscale", "kronecker", "nonuniform_scale",
    "mirror_h", "mirror_v", "quad_mirror", "concat", "concat_enhanced",
    "diagonal_tile", "fixed_crop", "spatial_gather",
    "varshape_spatial_gather", "gravity_unrolled", "edge_detect",
    "mode_fill", "downsample_stride", "symmetry_complete",
    "extract_inner", "add_border", "sparse_fill", "channel_filter",
]

COMPOSITION_PATTERNS = [
    "transform_then_recolor",
    "crop_then_transform",
    "recolor_then_tile",
]

SYSTEM_PROMPT = f"""You are a world-class ARC-AGI pattern classifier. Analyze grid transformations and predict which solver would produce the correct output.

Available single solvers:
{', '.join(SOLVER_NAMES)}

Available composition solvers (two transforms chained):
{', '.join(COMPOSITION_PATTERNS)}

Solver descriptions:
- identity: output = input
- constant: output is a fixed grid regardless of input
- color_map: per-pixel color remapping
- transpose: matrix transpose
- flip: horizontal or vertical flip
- rotate: 90/180/270 rotation
- shift: translate grid by offset
- tile: repeat input to fill output
- upscale: nearest-neighbor pixel-repeat zoom
- kronecker: kron(mask, input) self-similar
- nonuniform_scale: non-integer scale
- mirror_h/v: mirror and tile horizontally/vertically
- quad_mirror: 4-way kaleidoscope
- concat: concatenate transformed copies
- concat_enhanced: concat with color-dependent selection
- diagonal_tile: tile along diagonal
- fixed_crop: crop a rectangular region
- spatial_gather: arbitrary pixel rearrangement
- varshape_spatial_gather: spatial_gather with variable shapes
- gravity_unrolled: directional pixel compaction
- mode_fill: fill grid with most common color
- downsample_stride: subsample at regular stride
- symmetry_complete: complete partial symmetry
- extract_inner: remove outer border/frame
- add_border: add constant-color border
- sparse_fill: expand non-zero pixels into blocks
- channel_filter: keep only certain color channels
- transform_then_recolor: any spatial transform THEN color_map
- crop_then_transform: crop THEN apply spatial transform
- recolor_then_tile: color_map THEN tile/upscale

IMPORTANT: Look at ALL training pairs together. The pattern must be consistent across all pairs.

Output a valid JSON object mapping each task ID to:
{{
  "TASK_ID": {{
    "primary_solver": "solver_name",
    "fallback_solvers": ["solver1", "solver2"],
    "grid_size_changed": true/false,
    "confidence": 1-10,
    "notes": "brief pattern description"
  }}
}}

Output ONLY JSON. No other text."""


def format_grid(grid):
    return "\n".join([f"R{i}: {row}" for i, row in enumerate(grid)])


def classify_tasks(data_dir, output_file, api_key=None, base_url=None,
                   model="deepseek-chat", batch_size=5):
    """Classify all ARC tasks using DeepSeek API."""

    # --- API Setup ---
    if api_key:
        from openai import OpenAI
        client = OpenAI(api_key=api_key, base_url=base_url or "https://api.deepseek.com")
    else:
        try:
            from kaggle_secrets import UserSecretsClient
            from openai import OpenAI
            user_secrets = UserSecretsClient()
            client = OpenAI(
                api_key=user_secrets.get_secret("Deepseek_api_key"),
                base_url="https://api.deepseek.com"
            )
        except ImportError:
            raise RuntimeError("No API key provided and not on Kaggle.")

    # --- Load tasks ---
    all_files = sorted(glob.glob(os.path.join(data_dir, "task*.json")))
    if not all_files:
        all_files = sorted(glob.glob(os.path.join(data_dir, "*.json")))
    print(f"Found {len(all_files)} task files")

    classifications = {}

    # Resume from previous run
    if os.path.exists(output_file):
        with open(output_file) as f:
            classifications = json.load(f)
        print(f"Resuming: {len(classifications)} already classified")

    # --- Process in batches ---
    for i in range(0, len(all_files), batch_size):
        batch_files = all_files[i : i + batch_size]
        batch_ids = [os.path.basename(f).replace('.json','') for f in batch_files]
        if all(bid in classifications for bid in batch_ids):
            continue

        prompt = "Classify these ARC tasks:\n"
        for f in batch_files:
            tid = os.path.basename(f).replace('.json','')
            with open(f) as fh:
                task = json.load(fh)
            prompt += f"\n### TASK: {tid}\n"
            for idx, pair in enumerate(task.get('train', [])):
                prompt += f"--- Example {idx} ---\nIN:\n{format_grid(pair['input'])}\nOUT:\n{format_grid(pair['output'])}\n"
            for idx, pair in enumerate(task.get('test', [])):
                prompt += f"--- Test Input {idx} ---\nIN:\n{format_grid(pair['input'])}\n"

        for attempt in range(3):
            try:
                response = client.chat.completions.create(
                    model=model,
                    messages=[
                        {"role": "system", "content": SYSTEM_PROMPT},
                        {"role": "user", "content": prompt}
                    ],
                    response_format={'type': 'json_object'}
                )
                batch_results = json.loads(response.choices[0].message.content)
                classifications.update(batch_results)
                with open(output_file, 'w') as f:
                    json.dump(classifications, f, indent=2)
                print(f"  [{i+1}-{i+len(batch_files)}] Classified: {list(batch_results.keys())}")
                break
            except Exception as e:
                print(f"  Retry {attempt+1}: {e}")
                time.sleep(3)

    # --- Generate routing table ---
    routing = {}
    for tid, data in classifications.items():
        primary = data.get('primary_solver', '')
        fallbacks = data.get('fallback_solvers', [])
        solvers = [primary] + [s for s in fallbacks if s != primary]
        routing[tid] = {
            'solvers': solvers,
            'confidence': data.get('confidence', 5),
            'grid_changed': data.get('grid_size_changed', False),
            'notes': data.get('notes', '')
        }

    routing_file = output_file.replace('.json', '_routing.json')
    with open(routing_file, 'w') as f:
        json.dump(routing, f, indent=2)

    print(f"\nDone. {len(classifications)} tasks classified.")
    print(f"Classifications: {output_file}")
    print(f"Routing table: {routing_file}")
    return routing


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_dir', default='/kaggle/input/competitions/neurogolf-2026/')
    parser.add_argument('--output_file', default='/kaggle/working/arc_task_routes.json')
    parser.add_argument('--api_key', default='')
    parser.add_argument('--base_url', default='')
    parser.add_argument('--model', default='deepseek-chat')
    parser.add_argument('--batch_size', type=int, default=5)
    args = parser.parse_args()
    classify_tasks(args.data_dir, args.output_file, args.api_key,
                   args.base_url, args.model, args.batch_size)