mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-31 03:05:21 +00:00
allaun
475f6319ea
This squashes all local history (768 commits) onto the scrubbed PR #90 baseline. Individual commits were lost during filter-repo corruption; the working tree content is preserved intact. Build: N/A (working tree state only)
337 lines
11 KiB
Python
337 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
DeepSeek-Prover-V2-7B integration for Lean proof generation.
|
|
|
|
Two backends:
|
|
1. "ollama" (default) — talks to a local Ollama instance serving
|
|
deepseek-prover-v2:7b or a quantised variant. URL from
|
|
OLLAMA_URL env var (default http://localhost:11434).
|
|
2. "huggingface" — loads the model directly via `transformers`.
|
|
Requires: pip install transformers torch accelerate.
|
|
|
|
Environment variables:
|
|
DEEPSEEK_PROVER_BACKEND — "ollama" | "huggingface"
|
|
OLLAMA_URL — Ollama endpoint (default http://localhost:11434)
|
|
DEEPSEEK_API_KEY — API key for DeepSeek cloud (alternative to local)
|
|
PROOF_SERVER_URLS — comma-separated fallback proof server URLs
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
DEFAULT_OLLAMA_URL = "http://localhost:11434"
|
|
DEFAULT_MODEL = "deepseek-prover-v2:7b"
|
|
FALLBACK_MODELS = ["deepseek-coder-v2:16b", "codellama:13b"]
|
|
PROMPT_TEMPLATE = (
|
|
"You are a Lean 4 theorem prover. Complete the following proof.\n"
|
|
"Output ONLY valid Lean 4 code — no markdown fences, no explanations.\n\n"
|
|
"{context}\n\n"
|
|
"-- Prove:\n{theorem_statement}\n\n"
|
|
"{error_feedback}"
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class ProverConfig:
|
|
model: str = DEFAULT_MODEL
|
|
temperature: float = 0.6
|
|
top_p: float = 0.95
|
|
max_tokens: int = 4096
|
|
num_candidates: int = 4
|
|
stop_on_first: bool = True
|
|
|
|
|
|
@dataclass
|
|
class CandidateResult:
|
|
code: str
|
|
model: str
|
|
compile_log: str = ""
|
|
passed: bool = False
|
|
iteration: int = 0
|
|
latency_ms: float = 0.0
|
|
|
|
|
|
class DeepSeekProver:
|
|
"""Generate Lean proof candidates using DeepSeek-Prover-V2-7B."""
|
|
|
|
def __init__(
|
|
self,
|
|
backend: Optional[str] = None,
|
|
config: Optional[ProverConfig] = None,
|
|
api_key: Optional[str] = None,
|
|
):
|
|
self.config = config or ProverConfig()
|
|
self.backend = backend or os.environ.get(
|
|
"DEEPSEEK_PROVER_BACKEND", "ollama"
|
|
)
|
|
self.api_key = api_key or os.environ.get("DEEPSEEK_API_KEY", "")
|
|
self._hf_pipeline = None
|
|
|
|
# ------------------------------------------------------------------
|
|
# Public API
|
|
# ------------------------------------------------------------------
|
|
|
|
def generate(
|
|
self,
|
|
theorem_statement: str,
|
|
context: str = "",
|
|
error_feedback: str = "",
|
|
) -> list[CandidateResult]:
|
|
"""Generate N candidate proofs (N = config.num_candidates).
|
|
|
|
Returns a list of CandidateResult objects, each containing the
|
|
generated Lean code and metadata.
|
|
"""
|
|
prompt = PROMPT_TEMPLATE.format(
|
|
context=context,
|
|
theorem_statement=theorem_statement,
|
|
error_feedback=error_feedback,
|
|
)
|
|
|
|
if self.backend == "huggingface":
|
|
return self._generate_hf(prompt)
|
|
elif self.backend == "deepseek_api":
|
|
return self._generate_api(prompt)
|
|
else:
|
|
return self._generate_ollama(prompt)
|
|
|
|
def prove(
|
|
self,
|
|
theorem_statement: str,
|
|
context: str = "",
|
|
max_iterations: int = 5,
|
|
lean_file: str = "/tmp/proof_target.lean",
|
|
lake_workdir: Optional[str] = None,
|
|
) -> CandidateResult:
|
|
"""Iterative proof search: generate → compile → feedback → retry.
|
|
|
|
Args:
|
|
theorem_statement: The Lean theorem statement to prove.
|
|
context: Surrounding Lean module code (imports, definitions).
|
|
max_iterations: Max generate-compile cycles.
|
|
lean_file: Where to write the candidate for compilation.
|
|
lake_workdir: Working directory for `lake build`.
|
|
|
|
Returns:
|
|
CandidateResult with passed=True on success, or last failure.
|
|
"""
|
|
error_feedback = ""
|
|
for iteration in range(1, max_iterations + 1):
|
|
candidates = self.generate(
|
|
theorem_statement, context, error_feedback
|
|
)
|
|
best = candidates[0] if candidates else CandidateResult("", self.config.model)
|
|
best.iteration = iteration
|
|
|
|
if self._compile_check(best, lean_file, lake_workdir):
|
|
best.passed = True
|
|
return best
|
|
|
|
error_feedback = self._extract_errors(best.compile_log)
|
|
if not error_feedback:
|
|
error_feedback = "Compilation failed. Check syntax."
|
|
|
|
return best
|
|
|
|
# ------------------------------------------------------------------
|
|
# Ollama backend
|
|
# ------------------------------------------------------------------
|
|
|
|
def _generate_ollama(self, prompt: str) -> list[CandidateResult]:
|
|
models = [self.config.model] + FALLBACK_MODELS
|
|
last_error = ""
|
|
|
|
for model in models:
|
|
try:
|
|
return self._ollama_request(prompt, model)
|
|
except Exception as exc:
|
|
last_error = f"{exc}"
|
|
continue
|
|
|
|
raise RuntimeError(
|
|
f"All Ollama models failed. Last error: {last_error}"
|
|
)
|
|
|
|
def _ollama_request(
|
|
self, prompt: str, model: str
|
|
) -> list[CandidateResult]:
|
|
import requests
|
|
|
|
url = f"{os.environ.get('OLLAMA_URL', DEFAULT_OLLAMA_URL)}/api/generate"
|
|
results = []
|
|
|
|
for _ in range(self.config.num_candidates):
|
|
t0 = time.perf_counter()
|
|
resp = requests.post(
|
|
url,
|
|
json={
|
|
"model": model,
|
|
"prompt": prompt,
|
|
"stream": False,
|
|
"options": {
|
|
"temperature": self.config.temperature,
|
|
"top_p": self.config.top_p,
|
|
"num_predict": self.config.max_tokens,
|
|
},
|
|
},
|
|
timeout=120,
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
code = self._clean_lean_output(data.get("response", ""))
|
|
|
|
results.append(
|
|
CandidateResult(
|
|
code=code,
|
|
model=model,
|
|
latency_ms=(time.perf_counter() - t0) * 1000,
|
|
)
|
|
)
|
|
|
|
return results
|
|
|
|
# ------------------------------------------------------------------
|
|
# HuggingFace backend
|
|
# ------------------------------------------------------------------
|
|
|
|
def _generate_hf(self, prompt: str) -> list[CandidateResult]:
|
|
if self._hf_pipeline is None:
|
|
self._hf_pipeline = self._load_hf_model()
|
|
|
|
results = []
|
|
for _ in range(self.config.num_candidates):
|
|
t0 = time.perf_counter()
|
|
outputs = self._hf_pipeline(
|
|
prompt,
|
|
max_new_tokens=self.config.max_tokens,
|
|
temperature=self.config.temperature,
|
|
top_p=self.config.top_p,
|
|
do_sample=True,
|
|
)
|
|
code = self._clean_lean_output(outputs[0]["generated_text"])
|
|
results.append(
|
|
CandidateResult(
|
|
code=code,
|
|
model="deepseek-ai/DeepSeek-Prover-V2-7B",
|
|
latency_ms=(time.perf_counter() - t0) * 1000,
|
|
)
|
|
)
|
|
|
|
return results
|
|
|
|
def _load_hf_model(self):
|
|
"""Lazy-load the HuggingFace pipeline."""
|
|
from transformers import pipeline
|
|
|
|
return pipeline(
|
|
"text-generation",
|
|
model="deepseek-ai/DeepSeek-Prover-V2-7B",
|
|
device_map="auto",
|
|
torch_dtype="auto",
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# DeepSeek API backend
|
|
# ------------------------------------------------------------------
|
|
|
|
def _generate_api(self, prompt: str) -> list[CandidateResult]:
|
|
import requests
|
|
|
|
if not self.api_key:
|
|
raise RuntimeError(
|
|
"DEEPSEEK_API_KEY not set for deepseek_api backend"
|
|
)
|
|
|
|
results = []
|
|
for _ in range(self.config.num_candidates):
|
|
t0 = time.perf_counter()
|
|
resp = requests.post(
|
|
"https://api.deepseek.com/chat/completions",
|
|
headers={
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
json={
|
|
"model": "deepseek-prover-v2",
|
|
"messages": [
|
|
{
|
|
"role": "system",
|
|
"content": "You are a Lean 4 theorem prover. "
|
|
"Output ONLY valid Lean code.",
|
|
},
|
|
{"role": "user", "content": prompt},
|
|
],
|
|
"temperature": self.config.temperature,
|
|
"max_tokens": self.config.max_tokens,
|
|
},
|
|
timeout=120,
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
choice = data["choices"][0]["message"]["content"]
|
|
code = self._clean_lean_output(choice)
|
|
results.append(
|
|
CandidateResult(
|
|
code=code,
|
|
model="deepseek-prover-v2",
|
|
latency_ms=(time.perf_counter() - t0) * 1000,
|
|
)
|
|
)
|
|
|
|
return results
|
|
|
|
# ------------------------------------------------------------------
|
|
# Helpers
|
|
# ------------------------------------------------------------------
|
|
|
|
@staticmethod
|
|
def _clean_lean_output(text: str) -> str:
|
|
"""Strip markdown fences and trim."""
|
|
text = re.sub(r"^```(?:lean)?\s*\n?", "", text, flags=re.MULTILINE)
|
|
text = re.sub(r"\n```\s*$", "", text, flags=re.MULTILINE)
|
|
return text.strip()
|
|
|
|
@staticmethod
|
|
def _compile_check(
|
|
candidate: CandidateResult,
|
|
lean_file: str = "/tmp/proof_target.lean",
|
|
lake_workdir: Optional[str] = None,
|
|
) -> bool:
|
|
"""Write candidate to file and attempt lake build.
|
|
|
|
Returns True if the module compiles cleanly.
|
|
"""
|
|
Path(lean_file).write_text(candidate.code)
|
|
cmd = ["lake", "build"]
|
|
if lake_workdir:
|
|
cmd = ["lake", "-d", lake_workdir, "build"]
|
|
|
|
try:
|
|
import subprocess
|
|
|
|
result = subprocess.run(
|
|
cmd,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=120,
|
|
cwd=lake_workdir or Path(lean_file).parent,
|
|
)
|
|
candidate.compile_log = result.stdout + result.stderr
|
|
return result.returncode == 0
|
|
except subprocess.TimeoutExpired as exc:
|
|
candidate.compile_log = f"TIMEOUT: {exc}"
|
|
return False
|
|
|
|
@staticmethod
|
|
def _extract_errors(compile_log: str) -> str:
|
|
"""Extract Lean error messages for feedback."""
|
|
lines = compile_log.split("\n")
|
|
errors = [l for l in lines if "error:" in l.lower() or "sorry" in l]
|
|
return "\n".join(errors[:10])
|