| |
| """ |
| ARC-AGI Task Classifier via Kilo Code server (local DeepSeek, free tier). |
| Also supports DeepSeek API as fallback. |
| |
| Kilo server mode (preferred, faster on free tier): |
| 1. Start Kilo server: kilo serve --port 8765 |
| 2. Run: python classify_tasks.py --mode kilo --kilo_url http://127.0.0.1:8765 |
| |
| API mode (fallback): |
| python classify_tasks.py --mode api --data_dir /kaggle/input/competitions/neurogolf-2026/ |
| """ |
|
|
| import json, os, glob, time, argparse, requests |
|
|
| 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: {', '.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. |
| |
| 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 description"}}}} |
| |
| Output ONLY JSON.""" |
|
|
|
|
| def format_grid(grid): |
| return "\n".join([f"R{i}: {row}" for i, row in enumerate(grid)]) |
|
|
|
|
| def call_kilo(prompt, kilo_url, model="deepseek-ai/deepseek-chat", timeout=120): |
| """Call Kilo local server (OpenAI-compatible API).""" |
| payload = { |
| "model": model, |
| "messages": [ |
| {"role": "system", "content": SYSTEM_PROMPT}, |
| {"role": "user", "content": prompt} |
| ], |
| "temperature": 0.3, |
| } |
| resp = requests.post(f"{kilo_url}/v1/chat/completions", json=payload, timeout=timeout) |
| data = resp.json() |
| return data['choices'][0]['message']['content'] |
|
|
|
|
| def call_api(prompt, api_key, base_url="https://api.deepseek.com", model="deepseek-chat"): |
| """Call DeepSeek API (fallback, slower).""" |
| from openai import OpenAI |
| client = OpenAI(api_key=api_key, base_url=base_url) |
| response = client.chat.completions.create( |
| model=model, |
| messages=[ |
| {"role": "system", "content": SYSTEM_PROMPT}, |
| {"role": "user", "content": prompt} |
| ], |
| response_format={'type': 'json_object'}, |
| temperature=0.3, |
| ) |
| return response.choices[0].message.content |
|
|
|
|
| def classify_tasks(data_dir, output_file, mode="kilo", kilo_url="http://127.0.0.1:8765", |
| api_key=None, model="deepseek-ai/deepseek-chat", batch_size=5): |
| 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. Mode: {mode}") |
|
|
| classifications = {} |
| if os.path.exists(output_file): |
| with open(output_file) as f: |
| classifications = json.load(f) |
| print(f"Resuming: {len(classifications)} already classified") |
|
|
| 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: |
| if mode == "kilo": |
| content = call_kilo(prompt, kilo_url, model) |
| else: |
| content = call_api(prompt, api_key, model=model) |
| |
| |
| batch_results = json.loads(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)}] OK: {list(batch_results.keys())}") |
| break |
| except Exception as e: |
| print(f" Retry {attempt+1}: {e}") |
| time.sleep(3) |
|
|
| |
| 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)} classified. Routing: {routing_file}") |
| return routing |
|
|
|
|
| if __name__ == "__main__": |
| p = argparse.ArgumentParser() |
| p.add_argument('--mode', default='kilo', choices=['kilo', 'api']) |
| p.add_argument('--data_dir', default='/kaggle/input/competitions/neurogolf-2026/') |
| p.add_argument('--output_file', default='arc_task_routes.json') |
| p.add_argument('--kilo_url', default='http://127.0.0.1:8765') |
| p.add_argument('--api_key', default='') |
| p.add_argument('--model', default='deepseek-ai/deepseek-chat') |
| p.add_argument('--batch_size', type=int, default=5) |
| args = p.parse_args() |
| classify_tasks(args.data_dir, args.output_file, args.mode, |
| args.kilo_url, args.api_key, args.model, args.batch_size) |
|
|
