#!/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])