Update classifier to support Kilo local server

trm_solver/classify_tasks.py

@@ -1,17 +1,17 @@
 #!/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.
+ARC-AGI Task Classifier via Kilo Code server (local DeepSeek, free tier).
+Also supports DeepSeek API as fallback.
 
-Usage on Kaggle:
-  python classify_tasks.py
+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
 
-Usage locally:
-  python classify_tasks.py --data_dir ARC-AGI/data/training/
+API mode (fallback):
+  python classify_tasks.py --mode api --data_dir /kaggle/input/competitions/neurogolf-2026/
 """
 
-import json, os, glob, time, argparse
+import json, os, glob, time, argparse, requests
 
 SOLVER_NAMES = [
     "identity", "constant", "color_map", "transpose", "flip", "rotate",
@@ -31,11 +31,9 @@ COMPOSITION_PATTERNS = [
 
 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 single solvers: {', '.join(SOLVER_NAMES)}
 
-Available composition solvers (two transforms chained):
-{', '.join(COMPOSITION_PATTERNS)}
+Available composition solvers: {', '.join(COMPOSITION_PATTERNS)}
 
 Solver descriptions:
 - identity: output = input
@@ -69,49 +67,55 @@ Solver descriptions:
 - 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.
+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 pattern description"
-  }}
-}}
+{{"TASK_ID": {{"primary_solver": "solver_name", "fallback_solvers": ["solver1", "solver2"], "grid_size_changed": true/false, "confidence": 1-10, "notes": "brief description"}}}}
 
-Output ONLY JSON. No other text."""
+Output ONLY JSON."""
 
 
 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."""
-
-    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.")
-
+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")
+    print(f"Found {len(all_files)} task files. Mode: {mode}")
 
     classifications = {}
     if os.path.exists(output_file):
@@ -138,24 +142,23 @@ def classify_tasks(data_dir, output_file, api_key=None, base_url=None,
 
         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)
+                if mode == "kilo":
+                    content = call_kilo(prompt, kilo_url, model)
+                else:
+                    content = call_api(prompt, api_key, model=model)
+                
+                # Parse JSON from response
+                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)}] Classified: {list(batch_results.keys())}")
+                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)
 
+    # Generate routing table
     routing = {}
     for tid, data in classifications.items():
         primary = data.get('primary_solver', '')
@@ -167,25 +170,22 @@ def classify_tasks(data_dir, output_file, api_key=None, base_url=None,
             '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}")
+    print(f"\nDone. {len(classifications)} classified. Routing: {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)
+    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)