Final single-file submission: vLLM + SOAR-14B + heuristics (target 33.47%+)

submission.py

@@ -0,0 +1,690 @@
+#!/usr/bin/env python3
+"""
+ARC-AGI-2 Kaggle Submission — INTERNET OFF, 4× L4 GPUs, 12h
+Target: 33.47%+
+
+BEFORE SUBMITTING — you need these attached to your Kaggle notebook:
+  1. Kaggle Model: julien31/Soar-qwen-14b  (Import from HuggingFace)
+  2. Kaggle Dataset: your pip wheels (vllm + peft)
+  
+In your Kaggle notebook, just 2 cells:
+
+  Cell 1:
+    !pip install --no-index --find-links /kaggle/input/arc-pip-wheels vllm peft --quiet
+
+  Cell 2:
+    !python submission.py
+"""
+
+import os, sys, json, time, copy, random, traceback, gc
+from typing import List, Dict, Tuple, Optional, Any
+from collections import defaultdict, Counter
+import numpy as np
+
+os.environ["TRANSFORMERS_OFFLINE"] = "1"
+os.environ["HF_HUB_OFFLINE"] = "1"
+os.environ["HF_DATASETS_OFFLINE"] = "1"
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+# ======================== CONFIG ========================
+
+# Model — try these paths in order until one exists
+SOAR_PATHS = [
+    "/kaggle/input/soar-qwen-14b/transformers/default/1",
+    "/kaggle/input/soar-qwen-14b",
+    "/kaggle/input/soar-qwen-7b/transformers/default/1",
+    "/kaggle/input/soar-qwen-7b",
+]
+# Competition data
+DATA_PATHS = [
+    "/kaggle/input/arc-prize-2026-arc-agi-2",
+    "/kaggle/input/arc-prize-2025",
+]
+OUTPUT = "/kaggle/working/submission.json"
+
+# Budget per task
+N_SAMPLES    = 200    # programs to sample
+N_REFINE     = 100    # refinement attempts
+TEMP_SAMPLE  = 0.9
+TEMP_REFINE  = 0.7
+MAX_TOKENS   = 2048
+TOTAL_HOURS  = 11.5
+TP_SIZE      = 4      # tensor parallel across all 4 L4 GPUs
+
+START = time.time()
+def tleft(): return TOTAL_HOURS*3600 - (time.time()-START)
+
+def find(paths):
+    for p in paths:
+        if os.path.exists(p): return p
+    return paths[-1]
+
+# ======================== GRID UTILS ========================
+
+def geq(a,b):
+    if a is None or b is None: return False
+    if len(a)!=len(b): return False
+    return all(list(r1)==list(r2) for r1,r2 in zip(a,b))
+
+def ghash(g):
+    return tuple(tuple(r) for r in g) if g else None
+
+def gnp(g): return str(np.array(g))
+
+# ======================== SAFE EXEC ========================
+
+def safe_exec(code, inp):
+    try:
+        ns = {}
+        exec("import numpy as np\nfrom collections import Counter,defaultdict\nimport copy,itertools,math\n"+code, ns)
+        if "transform" not in ns: return None
+        r = ns["transform"](copy.deepcopy(inp))
+        if isinstance(r, np.ndarray): r = r.tolist()
+        if not isinstance(r, list) or not r: return None
+        out = []
+        for row in r:
+            if isinstance(row, np.ndarray): row = row.tolist()
+            if not isinstance(row, list): return None
+            out.append([int(c) for c in row])
+        if any(any(c<0 or c>9 for c in row) for row in out): return None
+        return out
+    except: return None
+
+def eval_code(code, task):
+    ok = 0; results = []
+    for p in task["train"]:
+        pred = safe_exec(code, p["input"])
+        c = pred is not None and geq(pred, p["output"])
+        if c: ok += 1
+        results.append({"output":pred,"correct":c,"is_test":False})
+    acc = ok/len(task["train"]) if task["train"] else 0
+    to = safe_exec(code, task["test"][0]["input"]) if task.get("test") else None
+    results.append({"output":to,"correct":None,"is_test":True})
+    return acc, results, to
+
+def extract_code(text):
+    if not text: return None
+    if "```python" in text:
+        for part in text.split("```python")[1:]:
+            end = part.find("```")
+            c = part[:end].strip() if end!=-1 else part.strip()
+            if "def transform" in c: return c
+    if "```" in text:
+        parts = text.split("```")
+        for i in range(1, len(parts), 2):
+            c = parts[i].strip()
+            if c.startswith("python\n"): c = c[7:]
+            if "def transform" in c: return c
+    if "def transform" in text:
+        s = text.index("def transform")
+        lines = text[s:].split("\n"); fl = [lines[0]]
+        for l in lines[1:]:
+            if l.strip() and not l[0].isspace() and l.startswith(("def ","class ","```")): break
+            fl.append(l)
+        return "\n".join(fl).rstrip()
+    return None
+
+# ======================== SOAR PROMPTS ========================
+
+CINFO = ("The number in the input grid can be mapped to the following colors: "
+         "0:Black; 1:Blue; 2:Red; 3:Green; 4:Yellow; 5:Grey; 6:Pink; "
+         "7:Orange; 8:Purple; 9:Brown\n")
+
+def fmt(task):
+    p = ["# Task to solve:"]
+    for i,pr in enumerate(task["train"]):
+        inp,out = pr["input"],pr["output"]
+        p.append(f"## Input {i+1} (grid shape: {len(inp)} by {len(inp[0])}):")
+        p.append(gnp(inp))
+        p.append(f"## Output {i+1} (grid shape: {len(out)} by {len(out[0])}):")
+        p.append(gnp(out))
+    for i,tp in enumerate(task["test"]):
+        inp = tp["input"]
+        p.append(f"## Test Input {i+1} (grid shape: {len(inp)} by {len(inp[0])}):")
+        p.append(gnp(inp))
+    return "\n".join(p)
+
+def sample_prompt(task):
+    return (
+        "You are an AI assistant specialized in solving Abstract Reasoning Corpus "
+        "(ARC-AGI) tasks by generating Python code.\n"
+        "Your goal is to analyze input-output grid pairs. The outputs were produced "
+        "by applying a transformation rule to the inputs. Implement the transformation "
+        "rules as a Python function.\n"
+        "You should only write the implemented the transformation in code.\n"
+        "You must write code in triple backticks (```python and then ```). "
+        "You must write a function called `transform` which takes a single argument, "
+        "the input grid as `list[list[int]]`, and returns the transformed grid "
+        "(also as `list[list[int]]`).\n"
+        "You should make sure that you implement a version of the transformation "
+        "that works in general (at least for all given input-output pairs and test input pairs).\n"
+        f"{CINFO}\nNow, solve the following ARC-AGI task:\n\n{fmt(task)}"
+    )
+
+def refine_prompt(task, code, results):
+    ts = fmt(task)
+    nc = sum(1 for r in results if r.get("correct"))
+    nt = sum(1 for r in results if not r.get("is_test"))
+    parts = [f"```python\n{code}\n```",
+             f"This implementation correctly worked on {nc}/{nt} train pairs.",
+             "Detailed results:"]
+    bad = []
+    for i,r in enumerate(results):
+        if r.get("is_test"):
+            o = gnp(r["output"]) if r.get("output") else "EXECUTION ERROR"
+            parts.append(f"## Test Output (unknown correctness):\n{o}")
+        elif r.get("correct"):
+            parts.append(f"## Output {i+1}: CORRECT")
+        else:
+            o = gnp(r["output"]) if r.get("output") else "EXECUTION ERROR"
+            parts.append(f"## Output {i+1}: INCORRECT\n{o}")
+            bad.append(f"Output {i+1}")
+    if bad:
+        parts.append(f"\nFix code for: {', '.join(bad)}")
+    return (
+        "You are an AI assistant specialized in solving Abstract Reasoning Corpus "
+        "(ARC-AGI) tasks by repairing Python code implementations.\n"
+        "Fix the `transform` function to work correctly for all inputs.\n"
+        f"{CINFO}\n**Task:**\n{ts}\n\n"
+        f"**Previous implementation:**\n" + "\n".join(parts)
+    )
+
+# ======================== HEURISTICS ========================
+
+class Heur:
+    def solve(self, t):
+        for fn in [self._id,self._cmap,self._rot,self._flip,self._trans,
+                   self._crop,self._scale,self._tile,self._grav,self._fill,
+                   self._overlay,self._rmcol,self._mirror]:
+            try:
+                r = fn(t)
+                if r and len(r)>0 and all(len(row)>0 for row in r):
+                    if all(all(isinstance(c,int) and 0<=c<=9 for c in row) for row in r):
+                        return r
+            except: pass
+        return None
+
+    def _id(s,t): return copy.deepcopy(t["test"][0]["input"]) if all(p["input"]==p["output"] for p in t["train"]) else None
+    def _cmap(s,t):
+        i0,o0=t["train"][0]["input"],t["train"][0]["output"]
+        if len(i0)!=len(o0) or len(i0[0])!=len(o0[0]): return None
+        cm={}
+        for r in range(len(i0)):
+            for c in range(len(i0[0])):
+                k,v=i0[r][c],o0[r][c]
+                if k in cm and cm[k]!=v: return None
+                cm[k]=v
+        for p in t["train"][1:]:
+            if len(p["input"])!=len(p["output"]) or len(p["input"][0])!=len(p["output"][0]): return None
+            for r in range(len(p["input"])):
+                for c in range(len(p["input"][0])):
+                    if cm.get(p["input"][r][c])!=p["output"][r][c]: return None
+        return [[cm.get(c,c) for c in row] for row in t["test"][0]["input"]]
+    def _rot(s,t):
+        for k in [1,2,3]:
+            if all(np.rot90(np.array(p["input"]),k=-k).tolist()==p["output"] for p in t["train"]):
+                return np.rot90(np.array(t["test"][0]["input"]),k=-k).tolist()
+        return None
+    def _flip(s,t):
+        for fn in [np.fliplr,np.flipud]:
+            if all(fn(np.array(p["input"])).tolist()==p["output"] for p in t["train"]):
+                return fn(np.array(t["test"][0]["input"])).tolist()
+        return None
+    def _trans(s,t):
+        return np.array(t["test"][0]["input"]).T.tolist() if all(np.array(p["input"]).T.tolist()==p["output"] for p in t["train"]) else None
+    def _crop(s,t):
+        for p in t["train"]:
+            a=np.array(p["input"]); nz=np.argwhere(a!=0)
+            if len(nz)==0: return None
+            r1,c1=nz.min(0); r2,c2=nz.max(0)
+            if a[r1:r2+1,c1:c2+1].tolist()!=p["output"]: return None
+        a=np.array(t["test"][0]["input"]); nz=np.argwhere(a!=0)
+        if len(nz)==0: return None
+        r1,c1=nz.min(0); r2,c2=nz.max(0)
+        return a[r1:r2+1,c1:c2+1].tolist()
+    def _scale(s,t):
+        for f in [2,3,4,5]:
+            if all(np.array_equal(np.repeat(np.repeat(np.array(p["input"]),f,0),f,1),np.array(p["output"])) for p in t["train"]):
+                return np.repeat(np.repeat(np.array(t["test"][0]["input"]),f,0),f,1).tolist()
+        return None
+    def _tile(s,t):
+        for nr in range(1,6):
+            for nc in range(1,6):
+                if nr==1 and nc==1: continue
+                if all(np.array_equal(np.tile(np.array(p["input"]),(nr,nc)),np.array(p["output"])) for p in t["train"]):
+                    return np.tile(np.array(t["test"][0]["input"]),(nr,nc)).tolist()
+        return None
+    def _grav(s,t):
+        for d in ['down','up','left','right']:
+            ok=True
+            for p in t["train"]:
+                a=np.array(p["input"]); o=np.array(p["output"])
+                if a.shape!=o.shape: ok=False; break
+                bg=Counter(a.flatten().tolist()).most_common(1)[0][0]
+                r=np.full_like(a,bg); h,w=a.shape
+                for idx in range(h if d in ['down','up'] else w):
+                    if d=='down':
+                        nb=[a[rr,idx] for rr in range(h) if a[rr,idx]!=bg]
+                        for i,v in enumerate(nb): r[h-len(nb)+i,idx]=v
+                    elif d=='up':
+                        nb=[a[rr,idx] for rr in range(h) if a[rr,idx]!=bg]
+                        for i,v in enumerate(nb): r[i,idx]=v
+                    elif d=='right':
+                        nb=[a[idx,cc] for cc in range(w) if a[idx,cc]!=bg]
+                        for i,v in enumerate(nb): r[idx,w-len(nb)+i]=v
+                    elif d=='left':
+                        nb=[a[idx,cc] for cc in range(w) if a[idx,cc]!=bg]
+                        for i,v in enumerate(nb): r[idx,i]=v
+                if not np.array_equal(r,o): ok=False; break
+            if ok:
+                a=np.array(t["test"][0]["input"]); bg=Counter(a.flatten().tolist()).most_common(1)[0][0]
+                r=np.full_like(a,bg); h,w=a.shape
+                for idx in range(h if d in ['down','up'] else w):
+                    if d=='down':
+                        nb=[a[rr,idx] for rr in range(h) if a[rr,idx]!=bg]
+                        for i,v in enumerate(nb): r[h-len(nb)+i,idx]=v
+                    elif d=='up':
+                        nb=[a[rr,idx] for rr in range(h) if a[rr,idx]!=bg]
+                        for i,v in enumerate(nb): r[i,idx]=v
+                    elif d=='right':
+                        nb=[a[idx,cc] for cc in range(w) if a[idx,cc]!=bg]
+                        for i,v in enumerate(nb): r[idx,w-len(nb)+i]=v
+                    elif d=='left':
+                        nb=[a[idx,cc] for cc in range(w) if a[idx,cc]!=bg]
+                        for i,v in enumerate(nb): r[idx,i]=v
+                return r.tolist()
+        return None
+    def _fill(s,t):
+        from collections import deque
+        for p in t["train"]:
+            if len(p["input"])!=len(p["output"]) or len(p["input"][0])!=len(p["output"][0]): return None
+        for fc in range(10):
+            ok=True
+            for p in t["train"]:
+                a=np.array(p["input"]); o=np.array(p["output"]); h,w=a.shape
+                bg=Counter(a.flatten().tolist()).most_common(1)[0][0]
+                vis=np.zeros((h,w),dtype=bool); q=deque()
+                for rr in range(h):
+                    for c in [0,w-1]:
+                        if a[rr,c]==bg and not vis[rr,c]: q.append((rr,c)); vis[rr,c]=True
+                for c in range(w):
+                    for rr in [0,h-1]:
+                        if a[rr,c]==bg and not vis[rr,c]: q.append((rr,c)); vis[rr,c]=True
+                while q:
+                    rr,c=q.popleft()
+                    for dr,dc in [(-1,0),(1,0),(0,-1),(0,1)]:
+                        nr,nc=rr+dr,c+dc
+                        if 0<=nr<h and 0<=nc<w and not vis[nr,nc] and a[nr,nc]==bg: vis[nr,nc]=True; q.append((nr,nc))
+                e=a.copy()
+                for rr in range(h):
+                    for c in range(w):
+                        if a[rr,c]==bg and not vis[rr,c]: e[rr,c]=fc
+                if not np.array_equal(e,o): ok=False; break
+            if ok:
+                a=np.array(t["test"][0]["input"]); h,w=a.shape; bg=Counter(a.flatten().tolist()).most_common(1)[0][0]
+                vis=np.zeros((h,w),dtype=bool); q=deque()
+                for rr in range(h):
+                    for c in [0,w-1]:
+                        if a[rr,c]==bg and not vis[rr,c]: q.append((rr,c)); vis[rr,c]=True
+                for c in range(w):
+                    for rr in [0,h-1]:
+                        if a[rr,c]==bg and not vis[rr,c]: q.append((rr,c)); vis[rr,c]=True
+                while q:
+                    rr,c=q.popleft()
+                    for dr,dc in [(-1,0),(1,0),(0,-1),(0,1)]:
+                        nr,nc=rr+dr,c+dc
+                        if 0<=nr<h and 0<=nc<w and not vis[nr,nc] and a[nr,nc]==bg: vis[nr,nc]=True; q.append((nr,nc))
+                r=a.copy()
+                for rr in range(h):
+                    for c in range(w):
+                        if a[rr,c]==bg and not vis[rr,c]: r[rr,c]=fc
+                return r.tolist()
+        return None
+    def _overlay(s,t):
+        for sp in ['h','v']:
+            for op in ['or','and']:
+                ok=True
+                for p in t["train"]:
+                    a=np.array(p["input"]); o=np.array(p["output"]); h,w=a.shape
+                    if sp=='h' and h%2==0: t1,t2=a[:h//2],a[h//2:]
+                    elif sp=='v' and w%2==0: t1,t2=a[:,:w//2],a[:,w//2:]
+                    else: ok=False; break
+                    if o.shape!=t1.shape: ok=False; break
+                    e=np.where(t1!=0,t1,t2) if op=='or' else np.where((t1!=0)&(t2!=0),t1,0)
+                    if not np.array_equal(e,o): ok=False; break
+                if ok:
+                    a=np.array(t["test"][0]["input"]); h,w=a.shape
+                    if sp=='h': t1,t2=a[:h//2],a[h//2:]
+                    else: t1,t2=a[:,:w//2],a[:,w//2:]
+                    return (np.where(t1!=0,t1,t2) if op=='or' else np.where((t1!=0)&(t2!=0),t1,0)).tolist()
+        return None
+    def _rmcol(s,t):
+        for rc in range(1,10):
+            ok=True
+            for p in t["train"]:
+                if len(p["input"])!=len(p["output"]) or len(p["input"][0])!=len(p["output"][0]): ok=False; break
+                for r in range(len(p["input"])):
+                    for c in range(len(p["input"][0])):
+                        ic,oc=p["input"][r][c],p["output"][r][c]
+                        if ic==rc:
+                            if oc!=0: ok=False; break
+                        elif ic!=oc: ok=False; break
+                    if not ok: break
+                if not ok: break
+            if ok: return [[0 if c==rc else c for c in row] for row in t["test"][0]["input"]]
+        return None
+    def _mirror(s,t):
+        for ax in ['h','v']:
+            ok=True
+            for p in t["train"]:
+                a=np.array(p["input"]); o=np.array(p["output"])
+                if a.shape!=o.shape: ok=False; break
+                m=np.fliplr(a) if ax=='h' else np.flipud(a)
+                e=a.copy(); mask=a==0; e[mask]=m[mask]
+                if not np.array_equal(e,o): ok=False; break
+            if ok:
+                a=np.array(t["test"][0]["input"])
+                m=np.fliplr(a) if ax=='h' else np.flipud(a)
+                r=a.copy(); mask=a==0; r[mask]=m[mask]; return r.tolist()
+        return None
+
+# ======================== vLLM INFERENCE ENGINE ========================
+
+class VLLMEngine:
+    """Wraps vLLM for batched SOAR inference."""
+    def __init__(self, model_path, tp=4):
+        from vllm import LLM
+        print(f"Loading vLLM: {model_path} (TP={tp})")
+        self.llm = LLM(
+            model=model_path,
+            tensor_parallel_size=tp,
+            dtype="bfloat16",
+            trust_remote_code=True,
+            max_model_len=8192,
+            gpu_memory_utilization=0.92,
+        )
+        print("  ✓ vLLM loaded")
+
+    def generate(self, prompts, temp=0.9, max_tokens=2048):
+        from vllm import SamplingParams
+        params = SamplingParams(
+            temperature=temp, top_p=0.95, max_tokens=max_tokens,
+            repetition_penalty=1.05, stop=["```\n\n", "\n\n\n\n"],
+        )
+        outputs = self.llm.generate(prompts, params, use_tqdm=False)
+        return [o.outputs[0].text for o in outputs]
+
+    def generate_chat(self, messages_list, temp=0.9, max_tokens=2048):
+        """Generate from list of chat message dicts."""
+        from vllm import SamplingParams
+        params = SamplingParams(
+            temperature=temp, top_p=0.95, max_tokens=max_tokens,
+            repetition_penalty=1.05,
+        )
+        # Apply chat template
+        tokenizer = self.llm.get_tokenizer()
+        prompts = []
+        for msgs in messages_list:
+            text = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+            prompts.append(text)
+        outputs = self.llm.generate(prompts, params, use_tqdm=False)
+        return [o.outputs[0].text for o in outputs]
+
+
+class TransformersEngine:
+    """Fallback: plain transformers (slower but always works)."""
+    def __init__(self, model_path):
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        print(f"Loading transformers: {model_path}")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_path, dtype=torch.bfloat16, device_map="auto", trust_remote_code=True)
+        self.model.eval()
+        print("  ✓ transformers loaded")
+
+    def generate_chat(self, messages_list, temp=0.9, max_tokens=2048):
+        import torch
+        results = []
+        for msgs in messages_list:
+            text = self.tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+            inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=8192)
+            inputs = {k:v.to(self.model.device) for k,v in inputs.items()}
+            with torch.no_grad():
+                out = self.model.generate(**inputs, max_new_tokens=max_tokens, temperature=temp,
+                                          top_p=0.95, do_sample=True, pad_token_id=self.tokenizer.eos_token_id,
+                                          repetition_penalty=1.05)
+            results.append(self.tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))
+        return results
+
+
+# ======================== SOAR SOLVER ========================
+
+def soar_solve_batch(engine, tasks_dict, n_samples=200, n_refine=100):
+    """
+    Solve ALL tasks using batched vLLM inference.
+    Key insight: vLLM is fastest when you submit big batches, not one-at-a-time.
+    """
+    task_ids = list(tasks_dict.keys())
+    all_programs = {tid: [] for tid in task_ids}  # tid -> [(code, acc, test_out, exec_results)]
+
+    # ---- Phase 1: Batched sampling ----
+    print(f"\n  SOAR Phase 1: Sampling {n_samples} programs per task across {len(task_ids)} tasks")
+
+    for wave in range(0, n_samples, 10):  # Sample in waves of 10
+        if tleft() < TOTAL_HOURS*3600*0.25: break  # Reserve 25% for refine + TTT
+
+        # Build batch: 10 prompts per unsolved task
+        batch_prompts = []
+        batch_tids = []
+        for tid in task_ids:
+            # Skip if already solved perfectly
+            if any(p[1] == 1.0 for p in all_programs[tid]): continue
+            prompt = sample_prompt(tasks_dict[tid])
+            msgs = [{"role": "user", "content": prompt}]
+            for _ in range(min(10, n_samples - wave)):
+                batch_prompts.append(msgs)
+                batch_tids.append(tid)
+
+        if not batch_prompts: break
+
+        # Generate entire batch at once (vLLM handles this efficiently)
+        try:
+            responses = engine.generate_chat(batch_prompts, temp=TEMP_SAMPLE, max_tokens=MAX_TOKENS)
+        except Exception as e:
+            print(f"    Wave {wave} error: {e}")
+            continue
+
+        # Evaluate all generated programs
+        for tid, text in zip(batch_tids, responses):
+            code = extract_code(text)
+            if code:
+                acc, er, to = eval_code(code, tasks_dict[tid])
+                all_programs[tid].append((code, acc, to, er))
+
+        # Progress
+        solved = sum(1 for tid in task_ids if any(p[1]==1.0 for p in all_programs[tid]))
+        partial = sum(1 for tid in task_ids if all_programs[tid] and not any(p[1]==1.0 for p in all_programs[tid]))
+        print(f"    Wave {wave+10}: solved={solved}, partial={partial}, "
+              f"unsolved={len(task_ids)-solved-partial}, time_left={tleft()/3600:.2f}h")
+
+    # ---- Phase 2: Batched refinement ----
+    print(f"\n  SOAR Phase 2: Refining top programs")
+
+    # Collect tasks that need refinement (not perfectly solved)
+    needs_refine = [tid for tid in task_ids
+                    if all_programs[tid] and not any(p[1]==1.0 for p in all_programs[tid])]
+
+    for wave in range(0, n_refine, 5):
+        if tleft() < TOTAL_HOURS*3600*0.15: break  # Reserve 15% for TTT
+
+        batch_prompts = []
+        batch_tids = []
+        for tid in needs_refine:
+            if any(p[1]==1.0 for p in all_programs[tid]): continue
+            # Pick best program to refine
+            best = max(all_programs[tid], key=lambda x: x[1])
+            rp = refine_prompt(tasks_dict[tid], best[0], best[3])
+            msgs = [{"role": "user", "content": rp}]
+            for _ in range(min(5, n_refine - wave)):
+                batch_prompts.append(msgs)
+                batch_tids.append(tid)
+
+        if not batch_prompts: break
+
+        try:
+            responses = engine.generate_chat(batch_prompts, temp=TEMP_REFINE, max_tokens=MAX_TOKENS)
+        except Exception as e:
+            print(f"    Refine wave {wave} error: {e}")
+            continue
+
+        for tid, text in zip(batch_tids, responses):
+            code = extract_code(text)
+            if code:
+                acc, er, to = eval_code(code, tasks_dict[tid])
+                all_programs[tid].append((code, acc, to, er))
+
+        solved = sum(1 for tid in task_ids if any(p[1]==1.0 for p in all_programs[tid]))
+        print(f"    Refine wave {wave+5}: solved={solved}, time_left={tleft()/3600:.2f}h")
+
+    # ---- Phase 3: Vote ----
+    results = {}
+    for tid in task_ids:
+        scores = defaultdict(float)
+        for code, acc, to, _ in all_programs[tid]:
+            if to is None: continue
+            k = ghash(to)
+            scores[k] += 1 + 1000*acc  # SOAR scoring
+        ranked = sorted(scores.items(), key=lambda x: -x[1])
+        preds = [[list(row) for row in k] for k,_ in ranked[:2]]
+        results[tid] = preds
+
+    return results
+
+
+# ======================== DATA LOADING ========================
+
+def load_tasks():
+    tasks = {}
+    for base in DATA_PATHS:
+        if not os.path.exists(base): continue
+        # Single JSON file
+        for fname in ["arc-agi-2_test_challenges.json","test_challenges.json"]:
+            p = os.path.join(base, fname)
+            if os.path.exists(p):
+                with open(p) as f: tasks = json.load(f)
+                print(f"Loaded {len(tasks)} tasks from {fname}")
+                return tasks
+        # Directory of JSONs
+        for f in sorted(os.listdir(base)):
+            if f.endswith(".json") and "solution" not in f and "sample" not in f:
+                fp = os.path.join(base, f)
+                with open(fp) as fh: d = json.load(fh)
+                if isinstance(d, dict) and "train" in d:
+                    tasks[f.replace(".json","")] = d
+                elif isinstance(d, dict):
+                    tasks.update(d)
+        if tasks:
+            print(f"Loaded {len(tasks)} tasks from {base}")
+            return tasks
+    # HuggingFace fallback (needs internet)
+    print("Falling back to HuggingFace...")
+    from datasets import load_dataset
+    ds = load_dataset("arc-agi-community/arc-agi-2", split="train")
+    for i,row in enumerate(ds):
+        tasks[f"task_{i:04d}"] = {"train":row["fewshots"],"test":row["question"]}
+    print(f"Loaded {len(tasks)} tasks")
+    return tasks
+
+
+# ======================== MAIN ========================
+
+def main():
+    global START
+    START = time.time()
+    print("="*70)
+    print("ARC-AGI-2 SOLVER — TARGET: 33.47%+")
+    print("="*70)
+
+    # Load tasks
+    tasks = load_tasks()
+    ids = sorted(tasks.keys())
+    print(f"Total tasks: {len(ids)}")
+
+    submission = {}
+    stats = defaultdict(int)
+    heur = Heur()
+
+    # ---- Phase 0: Heuristics (instant) ----
+    print("\n--- Phase 0: Heuristics ---")
+    heur_solved = {}
+    remaining = {}
+    for tid in ids:
+        pred = heur.solve(tasks[tid])
+        if pred:
+            submission[tid] = {"attempt_1": pred, "attempt_2": pred}
+            heur_solved[tid] = True
+            stats["heuristic"] += 1
+        else:
+            remaining[tid] = tasks[tid]
+    print(f"Heuristic: {len(heur_solved)}/{len(ids)}")
+    print(f"Remaining: {len(remaining)}")
+
+    # ---- Phase 1: SOAR Program Synthesis (vLLM) ----
+    print("\n--- Phase 1: SOAR Program Synthesis ---")
+    model_path = find(SOAR_PATHS)
+    print(f"Model: {model_path}")
+
+    engine = None
+    try:
+        engine = VLLMEngine(model_path, tp=TP_SIZE)
+    except Exception as e:
+        print(f"vLLM failed: {e}")
+        try:
+            engine = TransformersEngine(model_path)
+        except Exception as e2:
+            print(f"Transformers also failed: {e2}")
+
+    soar_results = {}
+    if engine and remaining:
+        soar_results = soar_solve_batch(engine, remaining, N_SAMPLES, N_REFINE)
+
+    # Build submission from SOAR results
+    for tid, preds in soar_results.items():
+        if preds:
+            while len(preds) < 2: preds.append(preds[0])
+            submission[tid] = {"attempt_1": preds[0], "attempt_2": preds[1]}
+            stats["soar"] += 1
+        else:
+            # Fallback: return input unchanged
+            submission[tid] = {
+                "attempt_1": copy.deepcopy(tasks[tid]["test"][0]["input"]),
+                "attempt_2": copy.deepcopy(tasks[tid]["test"][0]["input"]),
+            }
+            stats["unsolved"] += 1
+
+    # Fill any missing tasks
+    for tid in ids:
+        if tid not in submission:
+            submission[tid] = {
+                "attempt_1": copy.deepcopy(tasks[tid]["test"][0]["input"]),
+                "attempt_2": copy.deepcopy(tasks[tid]["test"][0]["input"]),
+            }
+            stats["unsolved"] += 1
+
+    # ---- Save ----
+    os.makedirs(os.path.dirname(OUTPUT) or ".", exist_ok=True)
+    with open(OUTPUT, "w") as f:
+        json.dump(submission, f)
+
+    elapsed = time.time() - START
+    print(f"\n{'='*70}")
+    print(f"DONE in {elapsed/3600:.2f}h")
+    print(f"  heuristic={stats['heuristic']} | soar={stats['soar']} | unsolved={stats['unsolved']}")
+    print(f"  Total: {len(submission)}/{len(ids)}")
+    print(f"  Output: {OUTPUT}")
+    print(f"{'='*70}")
+
+
+if __name__ == "__main__":
+    main()