Add multi-provider LLM solver: Gemini, DeepSeek, GLM, Ollama

scripts/llm_solver_cloud.py

@@ -0,0 +1,441 @@
+"""
+PEMF ARC-AGI — LLM Program Synthesis (Multi-Provider)
+=====================================================
+
+Supports:
+  - Google Gemini (free tier: 15 RPM, generous limits)
+  - DeepSeek V4 (very cheap: $0.07/M input tokens)  
+  - GLM-4 / ChatGLM (free tier available)
+  - Ollama local (any model)
+  - Any OpenAI-compatible API
+
+Usage:
+  # Gemini (free, recommended starting point)
+  export LLM_PROVIDER=gemini
+  export GEMINI_API_KEY=your_key_here
+  python llm_solver_cloud.py
+
+  # DeepSeek (cheapest cloud option)
+  export LLM_PROVIDER=deepseek
+  export DEEPSEEK_API_KEY=your_key_here
+  python llm_solver_cloud.py
+
+  # GLM
+  export LLM_PROVIDER=glm
+  export GLM_API_KEY=your_key_here
+  python llm_solver_cloud.py
+
+  # Ollama local
+  export LLM_PROVIDER=ollama
+  export OLLAMA_MODEL=qwen2.5-coder:32b
+  python llm_solver_cloud.py
+"""
+
+import os
+import sys
+import json
+import time
+import re
+import glob
+import numpy as np
+from typing import Dict, List, Optional, Tuple
+from collections import Counter
+import urllib.request
+
+
+# =============================================================================
+# PROVIDER CONFIGS
+# =============================================================================
+
+PROVIDERS = {
+    "gemini": {
+        "name": "Google Gemini",
+        "base_url": "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent",
+        "default_model": "gemini-2.0-flash",
+        "env_key": "GEMINI_API_KEY",
+        "free_tier": "15 RPM, 1M tokens/day",
+        "get_key_url": "https://aistudio.google.com/apikey",
+    },
+    "deepseek": {
+        "name": "DeepSeek",
+        "base_url": "https://api.deepseek.com/v1/chat/completions",
+        "default_model": "deepseek-chat",
+        "env_key": "DEEPSEEK_API_KEY",
+        "free_tier": "$0.07/M input, $0.27/M output",
+        "get_key_url": "https://platform.deepseek.com/api_keys",
+    },
+    "glm": {
+        "name": "GLM (Zhipu AI)",
+        "base_url": "https://open.bigmodel.cn/api/paas/v4/chat/completions",
+        "default_model": "glm-4-flash",
+        "env_key": "GLM_API_KEY",
+        "free_tier": "glm-4-flash is free",
+        "get_key_url": "https://open.bigmodel.cn/usercenter/apikeys",
+    },
+    "ollama": {
+        "name": "Ollama (local)",
+        "base_url": "http://localhost:11434/api/generate",
+        "default_model": "qwen2.5-coder:32b",
+        "env_key": None,
+    },
+}
+
+
+# =============================================================================
+# API CALLERS
+# =============================================================================
+
+def call_gemini(prompt: str, api_key: str, model: str = "gemini-2.0-flash",
+                temperature: float = 0.7) -> str:
+    """Call Google Gemini API."""
+    url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}"
+    payload = {
+        "contents": [{"parts": [{"text": prompt}]}],
+        "generationConfig": {
+            "temperature": temperature,
+            "maxOutputTokens": 2048,
+        }
+    }
+    data = json.dumps(payload).encode('utf-8')
+    req = urllib.request.Request(url, data=data,
+                                headers={"Content-Type": "application/json"},
+                                method='POST')
+    try:
+        with urllib.request.urlopen(req, timeout=120) as resp:
+            result = json.loads(resp.read().decode())
+            candidates = result.get('candidates', [])
+            if candidates:
+                parts = candidates[0].get('content', {}).get('parts', [])
+                if parts:
+                    return parts[0].get('text', '')
+        return "ERROR: No response content"
+    except Exception as e:
+        return f"ERROR: {e}"
+
+
+def call_deepseek(prompt: str, api_key: str, model: str = "deepseek-chat",
+                  temperature: float = 0.7) -> str:
+    """Call DeepSeek API (OpenAI-compatible)."""
+    url = "https://api.deepseek.com/v1/chat/completions"
+    payload = {
+        "model": model,
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": 2048,
+        "temperature": temperature,
+    }
+    data = json.dumps(payload).encode('utf-8')
+    req = urllib.request.Request(url, data=data,
+                                headers={"Content-Type": "application/json",
+                                         "Authorization": f"Bearer {api_key}"},
+                                method='POST')
+    try:
+        with urllib.request.urlopen(req, timeout=120) as resp:
+            result = json.loads(resp.read().decode())
+            return result['choices'][0]['message']['content']
+    except Exception as e:
+        return f"ERROR: {e}"
+
+
+def call_glm(prompt: str, api_key: str, model: str = "glm-4-flash",
+             temperature: float = 0.7) -> str:
+    """Call GLM/Zhipu API (OpenAI-compatible)."""
+    url = "https://open.bigmodel.cn/api/paas/v4/chat/completions"
+    payload = {
+        "model": model,
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": 2048,
+        "temperature": temperature,
+    }
+    data = json.dumps(payload).encode('utf-8')
+    req = urllib.request.Request(url, data=data,
+                                headers={"Content-Type": "application/json",
+                                         "Authorization": f"Bearer {api_key}"},
+                                method='POST')
+    try:
+        with urllib.request.urlopen(req, timeout=120) as resp:
+            result = json.loads(resp.read().decode())
+            return result['choices'][0]['message']['content']
+    except Exception as e:
+        return f"ERROR: {e}"
+
+
+def call_ollama(prompt: str, model: str = "qwen2.5-coder:32b",
+                temperature: float = 0.7) -> str:
+    """Call local Ollama."""
+    url = "http://localhost:11434/api/generate"
+    payload = {
+        "model": model,
+        "prompt": prompt,
+        "stream": False,
+        "options": {"temperature": temperature, "num_predict": 2048},
+    }
+    data = json.dumps(payload).encode('utf-8')
+    req = urllib.request.Request(url, data=data,
+                                headers={"Content-Type": "application/json"},
+                                method='POST')
+    try:
+        with urllib.request.urlopen(req, timeout=180) as resp:
+            result = json.loads(resp.read().decode())
+            return result.get('response', '')
+    except Exception as e:
+        return f"ERROR: {e}"
+
+
+def call_llm(prompt: str, provider: str, api_key: str = "",
+             model: str = "", temperature: float = 0.7) -> str:
+    """Unified LLM caller."""
+    if provider == "gemini":
+        return call_gemini(prompt, api_key, model or "gemini-2.0-flash", temperature)
+    elif provider == "deepseek":
+        return call_deepseek(prompt, api_key, model or "deepseek-chat", temperature)
+    elif provider == "glm":
+        return call_glm(prompt, api_key, model or "glm-4-flash", temperature)
+    elif provider == "ollama":
+        return call_ollama(prompt, model or "qwen2.5-coder:32b", temperature)
+    else:
+        return f"ERROR: Unknown provider {provider}"
+
+
+# =============================================================================
+# PROMPT, EXTRACTION, VERIFICATION (same as before)
+# =============================================================================
+
+def build_prompt(task: Dict) -> str:
+    train_pairs = task.get('train', [])
+    examples = []
+    for i, pair in enumerate(train_pairs):
+        examples.append(
+            f"Example {i+1}:\n"
+            f"  Input:  {json.dumps(pair['input'])}\n"
+            f"  Output: {json.dumps(pair['output'])}"
+        )
+    examples_str = "\n".join(examples)
+
+    inputs = [np.array(p['input']) for p in train_pairs]
+    outputs = [np.array(p['output']) for p in train_pairs]
+    same_shape = all(i.shape == o.shape for i, o in zip(inputs, outputs))
+    in_colors = sorted(set(c for i in inputs for c in np.unique(i).tolist()))
+    out_colors = sorted(set(c for o in outputs for c in np.unique(o).tolist()))
+
+    analysis = f"  Same input/output shape: {same_shape}\n"
+    analysis += f"  Input colors: {in_colors}, Output colors: {out_colors}\n"
+    if not same_shape:
+        for i, o in zip(inputs[:1], outputs[:1]):
+            analysis += f"  Shape: {i.shape} -> {o.shape}\n"
+
+    return f"""Solve this ARC-AGI puzzle. Write ONLY a Python function, no explanations.
+
+{examples_str}
+
+Analysis:
+{analysis}
+```python
+import numpy as np
+from collections import Counter, deque
+from scipy.ndimage import label
+
+def transform(grid: list[list[int]]) -> list[list[int]]:
+    grid = np.array(grid)
+"""
+
+
+def extract_code(response: str) -> Optional[str]:
+    for pattern in [r'```python\s*(.*?)```', r'```\s*(.*?)```']:
+        matches = re.findall(pattern, response, re.DOTALL)
+        for match in matches:
+            if 'def transform' in match:
+                return match.strip()
+    idx = response.find('def transform')
+    if idx >= 0:
+        before = response[:idx]
+        import_start = max(before.rfind('import '), before.rfind('from '))
+        start = import_start if import_start >= 0 else idx
+        code = response[start:]
+        end = code.find('```')
+        if end > 0:
+            code = code[:end]
+        return code.strip()
+    stripped = response.strip()
+    if stripped.startswith(('import', 'def transform', 'from')):
+        return stripped
+    return None
+
+
+def verify_program(code: str, train_pairs: List[Dict]) -> bool:
+    namespace = {'np': np, 'numpy': np, 'Counter': Counter,
+                 'deque': __import__('collections').deque}
+    try:
+        # Allow scipy import in generated code
+        try:
+            import scipy.ndimage
+            namespace['scipy'] = __import__('scipy')
+        except ImportError:
+            pass
+        exec(code, namespace)
+    except Exception:
+        return False
+    if 'transform' not in namespace:
+        return False
+    fn = namespace['transform']
+    for pair in train_pairs:
+        try:
+            result = fn([row[:] for row in pair['input']])
+            if result is None:
+                return False
+            r = np.array(result, dtype=int)
+            e = np.array(pair['output'], dtype=int)
+            if r.shape != e.shape or not np.array_equal(r, e):
+                return False
+        except Exception:
+            return False
+    return True
+
+
+def apply_program(code: str, test_input):
+    namespace = {'np': np, 'numpy': np, 'Counter': Counter,
+                 'deque': __import__('collections').deque}
+    try:
+        import scipy.ndimage
+        namespace['scipy'] = __import__('scipy')
+    except ImportError:
+        pass
+    try:
+        exec(code, namespace)
+        result = namespace['transform']([row[:] for row in test_input])
+        if result is not None:
+            return np.array(result, dtype=int).tolist()
+    except Exception:
+        pass
+    return None
+
+
+# =============================================================================
+# SYNTHESIS + MAIN
+# =============================================================================
+
+def synthesize_task(task, provider, api_key, model, n_candidates=8, verbose=False):
+    prompt = build_prompt(task)
+    for i in range(n_candidates):
+        temp = 0.1 if i == 0 else min(0.4 + 0.15 * i, 1.2)
+        response = call_llm(prompt, provider, api_key, model, temp)
+        if response.startswith("ERROR:"):
+            if verbose: print(f"    C{i+1}: {response[:60]}")
+            # Rate limit — wait and retry
+            if "429" in response or "rate" in response.lower():
+                time.sleep(5)
+            continue
+        code = extract_code(response)
+        if code is None:
+            if verbose: print(f"    C{i+1}: no code")
+            continue
+        if verbose: print(f"    C{i+1}: {len(code)}ch", end="")
+        if verify_program(code, task['train']):
+            if verbose: print(" ✅")
+            return (f"llm_c{i+1}", code)
+        else:
+            if verbose: print(" ❌")
+    return None
+
+
+def main():
+    PROVIDER = os.environ.get("LLM_PROVIDER", "gemini")
+    config = PROVIDERS.get(PROVIDER, {})
+    API_KEY = os.environ.get(config.get("env_key", ""), "") if config.get("env_key") else ""
+    MODEL = os.environ.get("LLM_MODEL", config.get("default_model", ""))
+    N_CANDIDATES = int(os.environ.get("N_CANDIDATES", "8"))
+    ARC_DIR = os.environ.get("ARC_DIR", "arc_data/training")
+    ALREADY_SOLVED = os.environ.get("ALREADY_SOLVED", "already_solved.json")
+    OUTPUT = os.environ.get("OUTPUT_FILE", "llm_results.json")
+
+    print("=" * 60)
+    print(f"PEMF ARC-AGI — LLM Synthesis ({config.get('name', PROVIDER)})")
+    print("=" * 60)
+    print(f"Provider: {PROVIDER}")
+    print(f"Model: {MODEL}")
+    print(f"Candidates/task: {N_CANDIDATES}")
+    if not API_KEY and PROVIDER != "ollama":
+        print(f"\n⚠️  No API key! Set {config.get('env_key', '???')}")
+        print(f"   Get key: {config.get('get_key_url', '?')}")
+        return
+    print()
+
+    # Load already solved
+    already_solved = set()
+    if os.path.exists(ALREADY_SOLVED):
+        with open(ALREADY_SOLVED) as f:
+            already_solved = set(json.load(f))
+        print(f"Symbolic solved: {len(already_solved)}")
+
+    # Load tasks
+    task_files = sorted(glob.glob(os.path.join(ARC_DIR, "*.json")))
+    unsolved = [(os.path.basename(tf).replace('.json',''), tf)
+                for tf in task_files
+                if os.path.basename(tf).replace('.json','') not in already_solved]
+    print(f"Total tasks: {len(task_files)}, unsolved: {len(unsolved)}")
+    print()
+
+    # Run
+    results = {}
+    solved = 0
+    total_time = 0
+
+    for idx, (tid, tf) in enumerate(unsolved):
+        with open(tf) as f:
+            task = json.load(f)
+        print(f"[{idx+1:3d}/{len(unsolved)}] {tid}:", end=" ", flush=True)
+        start = time.time()
+        result = synthesize_task(task, PROVIDER, API_KEY, MODEL, N_CANDIDATES, verbose=False)
+        elapsed = time.time() - start
+        total_time += elapsed
+
+        if result:
+            rule, code = result
+            solved += 1
+            test_outputs = [apply_program(code, t['input']) for t in task.get('test', [])]
+            results[tid] = {'status': 'solved', 'rule': rule, 'code': code,
+                           'test_outputs': test_outputs, 'time_s': round(elapsed, 2)}
+            print(f"✅ ({elapsed:.1f}s)")
+        else:
+            results[tid] = {'status': 'failed', 'time_s': round(elapsed, 2)}
+            print(f"❌ ({elapsed:.1f}s)")
+
+        # Rate limit respect
+        if PROVIDER == "gemini":
+            time.sleep(4)  # 15 RPM = 1 every 4s
+        elif PROVIDER in ("deepseek", "glm"):
+            time.sleep(1)
+
+        # Save every 10
+        if (idx + 1) % 10 == 0:
+            _save(OUTPUT, PROVIDER, MODEL, N_CANDIDATES, solved, idx+1,
+                  total_time, already_solved, len(task_files), results)
+            print(f"  [Saved: {solved}/{idx+1}, total {len(already_solved)+solved}/{len(task_files)}]")
+
+    # Final save
+    _save(OUTPUT, PROVIDER, MODEL, N_CANDIDATES, solved, len(unsolved),
+          total_time, already_solved, len(task_files), results)
+
+    print(f"\n{'='*60}")
+    print(f"LLM solved:      {solved}/{len(unsolved)}")
+    print(f"Symbolic:        {len(already_solved)}")
+    print(f"TOTAL:           {len(already_solved)+solved}/{len(task_files)} ({100*(len(already_solved)+solved)/len(task_files):.1f}%)")
+    print(f"Saved: {OUTPUT}")
+
+
+def _save(path, provider, model, n_cand, solved, attempted, total_time,
+          already_solved, total_tasks, results):
+    with open(path, 'w') as f:
+        json.dump({
+            'provider': provider, 'model': model, 'n_candidates': n_cand,
+            'llm_solved': solved, 'attempted': attempted,
+            'total_time_s': round(total_time, 1),
+            'symbolic_solved': len(already_solved),
+            'total_solved': len(already_solved) + solved,
+            'total_tasks': total_tasks,
+            'solve_rate': round(100*(len(already_solved)+solved)/total_tasks, 2),
+            'results': results,
+        }, f, indent=2)
+
+
+if __name__ == "__main__":
+    main()