allaun/Research-Stack
1
0
Fork 0
mirror of https://github.com/allaunthefox/Research-Stack.git synced 2026-07-31 03:05:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions
Research-Stack/4-Infrastructure/shim/alphaproof_loop.py
Brandon Schneider 7047a770d1 feat: AlphaProof batch mode with copy-if pre-filter 
New CLI mode: python alphaproof_loop.py --batch <lean_dir>

Pipeline:
1. Pre-filter scans Lean codebase (11,434 theorems)
2. Classifies trivial (63.5%) vs non-trivial (36.5%)
3. Prioritizes non-trivial by tactic count + sorry weight
4. Feeds hardest problems to Ollama LLM first
5. Skips trivial theorems entirely (zero deltas)

Usage:
  python alphaproof_loop.py --batch ../../0-Core-Formalism/lean/Semantics/Semantics
  python alphaproof_loop.py --batch . --max-iter 20 --model deepseek-coder-v2:16b

Same pattern as vectorized copy_if:
  - Skip zero deltas → 40x faster (blog post)
  - Skip trivial theorems → 2.7x faster (pre-filter)
  - Focus solver on residuals → fewer iterations
2026-05-29 02:42:04 -05:00

563 lines
19 KiB
Python
Raw Blame History

"""
AlphaProof-inspired proof search loop.
Generates Lean proof candidates via a local LLM (Ollama), verifies them
through lake build, and iterates with error feedback until a valid proof
is found or max iterations exhausted.
Logs all iterations to JSONL for replay analysis.
"""
import json
import os
import re
import subprocess
import tempfile
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional, List, Dict
try:
import requests
HAS_REQUESTS = True
except ImportError:
HAS_REQUESTS = False
# Import pre-filter for copy-if pattern
try:
from lean_proof_prefilter import analyze_file, classify_theorem, extract_theorems
HAS_PREFILTER = True
except ImportError:
HAS_PREFILTER = False
# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://localhost:11434")
DEFAULT_MODEL = os.environ.get("ALPHAPROOF_MODEL", "deepseek-coder-v2:16b")
LOG_DIR = Path(__file__).parent / "alphaproof_logs"
Q16_SCALE = 65536 # 2^16 for Q16.16 fixed-point
# ---------------------------------------------------------------------------
# Ollama LLM interface
# ---------------------------------------------------------------------------
def generate_candidate(prompt: str, model: str = DEFAULT_MODEL,
temperature: float = 0.7) -> str:
"""Call Ollama API to generate a candidate Lean proof.
Args:
prompt: The problem description + error feedback.
model: Ollama model name.
temperature: Sampling temperature.
Returns:
Generated Lean code string.
"""
if not HAS_REQUESTS:
raise ImportError("requests library required. pip install requests")
system_prompt = (
"You are a Lean 4 theorem prover. Given a mathematical problem, "
"generate a complete, compilable Lean 4 module with all necessary "
"imports. Use `import Mathlib` for standard library. "
"Output ONLY the Lean code, no markdown fences or explanations. "
"The module must compile with `lake build`."
)
payload = {
"model": model,
"system": system_prompt,
"prompt": prompt,
"stream": False,
"options": {
"temperature": temperature,
"num_predict": 4096,
}
}
try:
resp = requests.post(f"{OLLAMA_URL}/api/generate", json=payload,
timeout=120)
resp.raise_for_status()
result = resp.json()
code = result.get("response", "")
# Strip markdown fences if present
code = re.sub(r'^```(?:lean|lean4)?\s*\n?', '', code, flags=re.MULTILINE)
code = re.sub(r'\n?```\s*$', '', code, flags=re.MULTILINE)
return code.strip()
except requests.RequestException as e:
return f"-- ERROR: Ollama request failed: {e}"
# ---------------------------------------------------------------------------
# Lean verification
# ---------------------------------------------------------------------------
def verify_proof(lean_code: str, module_name: str = "Candidate",
lake_project: Optional[str] = None,
timeout: int = 120) -> dict:
"""Write Lean code to a temp file and verify via lake build.
Args:
lean_code: Complete Lean 4 module source.
module_name: Name for the temp module file.
lake_project: Path to existing lake project. If None, creates a temp one.
timeout: Build timeout in seconds.
Returns:
{'success': bool, 'errors': list[str], 'warnings': list[str]}
"""
errors = []
warnings = []
with tempfile.TemporaryDirectory(prefix="alphaproof_") as tmpdir:
tmpdir = Path(tmpdir)
if lake_project is None:
# Create a minimal lake project
lakefile = tmpdir / "lakefile.lean"
lakefile.write_text(
'import Lake\n'
'open Lake DSL in\n'
'package «candidate» where\n'
' leanOptions := #[\n'
' ⟨`pp.unicode.fun, true⟩\n'
' ]\n\n'
'require mathlib from git\n'
' "https://github.com/leanprover-community/mathlib4" @ "master"\n'
)
project_dir = tmpdir
else:
project_dir = Path(lake_project)
# Write candidate file
candidate_dir = project_dir / "Candidate"
candidate_dir.mkdir(exist_ok=True)
candidate_file = candidate_dir / f"{module_name}.lean"
candidate_file.write_text(lean_code)
# Run lake build
try:
result = subprocess.run(
["lake", "build", f"Candidate.{module_name}"],
cwd=str(project_dir),
capture_output=True,
text=True,
timeout=timeout,
env={**os.environ, "LAKE_PATH": str(project_dir)}
)
if result.returncode == 0:
return {'success': True, 'errors': [], 'warnings': []}
# Parse errors from stderr
stderr = result.stderr
error_lines = []
for line in stderr.split('\n'):
line = line.strip()
if not line:
continue
if 'error' in line.lower():
error_lines.append(line)
elif 'warning' in line.lower():
warnings.append(line)
elif line.startswith('error:') or 'Error' in line:
error_lines.append(line)
if not error_lines and stderr:
error_lines = [stderr[:500]]
return {
'success': False,
'errors': error_lines,
'warnings': warnings
}
except subprocess.TimeoutExpired:
return {
'success': False,
'errors': [f'Build timed out after {timeout}s'],
'warnings': []
}
except FileNotFoundError:
return {
'success': False,
'errors': ['lake command not found. Is Lean/Lake installed?'],
'warnings': []
}
# ---------------------------------------------------------------------------
# FPGA Q16 acceleration (Python placeholder)
# ---------------------------------------------------------------------------
def q16_multiply(a: int, b: int) -> int:
"""Q16.16 fixed-point multiplication.
(a * b) >> 16 with overflow clamping.
"""
result = (a * b) >> 16
# Clamp to Q16.16 range
INT32_MAX = 2147483647
INT32_MIN = -2147483648
return max(INT32_MIN, min(INT32_MAX, result))
def q16_from_float(f: float) -> int:
"""Convert float to Q16.16."""
return max(-2147483648, min(2147483647, round(f * Q16_SCALE)))
def q16_to_float(q: int) -> float:
"""Convert Q16.16 to float."""
return q / Q16_SCALE
def fpga_accelerate(candidates: list[dict]) -> list[dict]:
"""Rank candidates using Q16.16 fixed-point scoring.
This is a Python placeholder for FPGA Q16 LUT acceleration.
In production, this would route through UART to the Tang Nano 9K
or direct memory-mapped LUT access.
Scoring heuristic:
- Shorter proofs score higher (fewer tokens)
- Proofs with fewer sorry/axiom/admit score higher
- Previous success patterns boost score
Args:
candidates: list of {'code': str, 'iteration': int, ...}
Returns:
Candidates sorted by Q16 score (descending), with 'q16_score' added.
"""
scored = []
for c in candidates:
code = c.get('code', '')
# Length penalty (shorter is better)
length_score = q16_from_float(1.0 / (1.0 + len(code) / 1000.0))
# Penalty for incomplete proofs
penalty = 0
for bad_word in ['sorry', 'admit', 'axiom', 'by omega']:
count = code.count(bad_word)
penalty += q16_from_float(count * 0.1)
# Base score from length
base = length_score
# Bonus for having structure (def, theorem, proof)
if 'theorem' in code or 'lemma' in code:
base = q16_multiply(base, q16_from_float(1.2))
final_score = max(0, base - penalty)
c['q16_score'] = final_score
c['q16_score_float'] = q16_to_float(final_score)
scored.append(c)
scored.sort(key=lambda x: x.get('q16_score', 0), reverse=True)
return scored
# ---------------------------------------------------------------------------
# Proof prioritization (copy-if pre-filter)
# ---------------------------------------------------------------------------
def prioritize_problems(lean_dir: str, max_problems: int = 20) -> List[Dict]:
"""Pre-filter and prioritize proof problems using the copy-if pattern.
Scans the Lean codebase, classifies theorems as trivial/non-trivial,
and returns the highest-priority non-trivial problems for the LLM.
This is the Lean equivalent of the vectorized copy_if:
- Trivial theorems (zero deltas) = skip
- Non-trivial theorems (non-zero deltas) = feed to solver
Args:
lean_dir: Path to Lean source directory.
max_problems: Maximum number of problems to return.
Returns:
List of prioritized problem dicts with file, theorem, and priority.
"""
if not HAS_PREFILTER:
print("[alphaproof] pre-filter not available, using default order")
return []
problems = []
lean_path = Path(lean_dir)
for lean_file in sorted(lean_path.rglob('*.lean')):
if '.lake' in str(lean_file):
continue
try:
analysis = analyze_file(str(lean_file))
for thm in analysis['theorems']:
if thm['classification'] == 'non_trivial':
# Priority: more tactics = harder = higher priority
tactic_count = len(thm.get('category', '').split(','))
has_sorry = 'sorry' in thm.get('category', '')
priority = tactic_count + (10 if has_sorry else 0)
problems.append({
'file': str(lean_file),
'theorem': thm['name'],
'kind': thm['kind'],
'line': thm['line'],
'category': thm['category'],
'priority': priority,
'block': thm['block'],
})
except Exception:
continue
# Sort by priority (highest first = hardest problems first)
problems.sort(key=lambda p: p['priority'], reverse=True)
# Deduplicate by theorem name
seen = set()
unique = []
for p in problems:
if p['theorem'] not in seen:
seen.add(p['theorem'])
unique.append(p)
return unique[:max_problems]
def batch_proof_search(lean_dir: str, max_iterations: int = 10,
model: str = DEFAULT_MODEL,
lake_project: Optional[str] = None) -> Dict:
"""Run AlphaProof on prioritized non-trivial theorems.
Uses the copy-if pre-filter to skip trivial theorems and focus
the LLM solver on the hardest problems first.
"""
problems = prioritize_problems(lean_dir)
if not problems:
print("[alphaproof] no non-trivial problems found")
return {'problems': [], 'results': []}
print(f"[alphaproof] found {len(problems)} non-trivial problems")
print(f"[alphaproof] top 5: {[p['theorem'] for p in problems[:5]]}")
results = []
for i, problem in enumerate(problems):
print(f"\n[alphaproof] [{i+1}/{len(problems)}] {problem['theorem']}"
f" (priority={problem['priority']}, category={problem['category']})")
# Build prompt from the theorem block
prompt = (
f"Prove this Lean 4 theorem. The file is {problem['file']}, "
f"line {problem['line']}.\n\n"
f"```lean\n{problem['block']}\n```\n\n"
f"Replace the sorry with a complete proof. "
f"Use omega, simp, decide, norm_num, ring, or linarith as appropriate."
)
result = search_loop(
problem=prompt,
max_iterations=max_iterations,
model=model,
lake_project=lake_project,
)
results.append({
'theorem': problem['theorem'],
'file': problem['file'],
'priority': problem['priority'],
'success': result['success'],
'iterations': result['iterations'],
'solution': result.get('solution', ''),
})
if result['success']:
print(f"[alphaproof] SOLVED {problem['theorem']} "
f"in {result['iterations']} iterations")
else:
print(f"[alphaproof] FAILED {problem['theorem']} "
f"after {result['iterations']} iterations")
# Summary
solved = sum(1 for r in results if r['success'])
print(f"\n[alphaproof] summary: {solved}/{len(results)} solved")
return {
'problems': problems,
'results': results,
'solved': solved,
'total': len(results),
}
# Main search loop
# ---------------------------------------------------------------------------
def search_loop(problem: str, max_iterations: int = 50,
model: str = DEFAULT_MODEL,
lake_project: Optional[str] = None,
temperature: float = 0.7,
build_timeout: int = 120) -> dict:
"""Main AlphaProof search loop: generate → verify → feedback → retry.
Args:
problem: Natural language description of the problem to prove.
max_iterations: Maximum number of generate-verify cycles.
model: Ollama model to use.
lake_project: Optional path to existing lake project.
temperature: LLM temperature.
build_timeout: Lean build timeout per iteration.
Returns:
{'solution': str, 'iterations': int, 'success': bool,
'candidates': list[dict], 'log_path': str}
"""
LOG_DIR.mkdir(parents=True, exist_ok=True)
log_path = LOG_DIR / f"search_{datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')}.jsonl"
candidates = []
best_solution = ""
feedback = ""
with open(log_path, 'w') as log_f:
for iteration in range(1, max_iterations + 1):
t0 = time.time()
# Build prompt with feedback from previous iterations
if feedback:
prompt = (
f"Problem:\n{problem}\n\n"
f"Previous attempt failed with these errors:\n{feedback}\n\n"
f"Please fix the Lean code and try again. "
f"Iteration {iteration}/{max_iterations}."
)
else:
prompt = f"Problem:\n{problem}\n\nPlease write a Lean 4 proof."
# Generate candidate
code = generate_candidate(prompt, model=model,
temperature=temperature)
gen_time = time.time() - t0
# Verify
t1 = time.time()
result = verify_proof(code, module_name=f"Iter{iteration}",
lake_project=lake_project,
timeout=build_timeout)
ver_time = time.time() - t1
candidate = {
'iteration': iteration,
'code': code,
'success': result['success'],
'errors': result['errors'],
'warnings': result['warnings'],
'gen_time': round(gen_time, 2),
'verify_time': round(ver_time, 2),
'timestamp': datetime.now(timezone.utc).isoformat(),
}
candidates.append(candidate)
# Log to JSONL
log_entry = {**candidate, 'code_length': len(code)}
log_f.write(json.dumps(log_entry) + '\n')
log_f.flush()
if result['success']:
best_solution = code
# Rank remaining candidates via FPGA Q16 path
candidates = fpga_accelerate(candidates)
return {
'solution': best_solution,
'iterations': iteration,
'success': True,
'candidates': candidates,
'log_path': str(log_path),
}
# Build feedback for next iteration
feedback = '\n'.join(result['errors'][:5]) # Top 5 errors
print(f"[alphaproof] iter {iteration}/{max_iterations}: "
f"FAIL ({len(result['errors'])} errors, "
f"gen={gen_time:.1f}s, ver={ver_time:.1f}s)")
# Exhausted iterations
candidates = fpga_accelerate(candidates)
return {
'solution': best_solution,
'iterations': max_iterations,
'success': False,
'candidates': candidates,
'log_path': str(log_path),
}
# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------
if __name__ == '__main__':
import sys
if len(sys.argv) < 2:
print("Usage: python alphaproof_loop.py <problem_file_or_text> "
"[--max-iter N] [--model MODEL]")
print(" python alphaproof_loop.py --batch <lean_dir> "
"[--max-iter N] [--model MODEL]")
sys.exit(1)
problem_arg = sys.argv[1]
# Batch mode: run pre-filter on Lean codebase
if problem_arg == '--batch':
lean_dir = sys.argv[2] if len(sys.argv) > 2 else '.'
max_iter = 10
model_name = DEFAULT_MODEL
for i, arg in enumerate(sys.argv[3:], 3):
if arg == '--max-iter' and i + 1 < len(sys.argv):
max_iter = int(sys.argv[i + 1])
elif arg == '--model' and i + 1 < len(sys.argv):
model_name = sys.argv[i + 1]
result = batch_proof_search(lean_dir, max_iterations=max_iter,
model=model_name)
print(f"\n{'='*60}")
print(f"Solved: {result['solved']}/{result['total']}")
print(f"{'='*60}")
sys.exit(0 if result['solved'] > 0 else 1)
# Single problem mode
if os.path.isfile(problem_arg):
problem_text = Path(problem_arg).read_text()
else:
problem_text = problem_arg
max_iter = 50
model_name = DEFAULT_MODEL
for i, arg in enumerate(sys.argv[2:], 2):
if arg == '--max-iter' and i + 1 < len(sys.argv):
max_iter = int(sys.argv[i + 1])
elif arg == '--model' and i + 1 < len(sys.argv):
model_name = sys.argv[i + 1]
result = search_loop(problem_text, max_iterations=max_iter,
model=model_name)
print(f"\n{'='*60}")
print(f"Success: {result['success']}")
print(f"Iterations: {result['iterations']}")
print(f"Log: {result['log_path']}")
if result['success']:
print(f"\nSolution:\n{result['solution']}")
Reference in a new issue View git blame Copy permalink