304 lines
11 KiB
Python
304 lines
11 KiB
Python
import re
|
|
from typing import Any
|
|
|
|
|
|
LEAN4_DEFAULT_HEADER = (
|
|
"import Mathlib\nimport Aesop\n\n"
|
|
"set_option maxHeartbeats 0\n\n"
|
|
"open BigOperators Real Nat Topology Rat\n\n"
|
|
)
|
|
DEFAULT_LEAN_VERSION = "v4.15.0"
|
|
DEFAULT_TIMEOUT_PER_CMD = 60
|
|
DEFAULT_MISSING_LEAN_CODE = "sorry, no code found in output"
|
|
DEFAULT_DEEPSEEK_GENERATION_MAX_TOKENS = 32768
|
|
DEFAULT_STANDARD_GENERATION_MAX_TOKENS = 32768
|
|
DEFAULT_GENERATION_TOP_P = 0.95
|
|
DEEPSEEK_STYLE_MODEL_MARKERS = (
|
|
"deepseek-ai",
|
|
"goedel-lm",
|
|
)
|
|
DECLARATION_START_RE = re.compile(
|
|
r"(?m)^\s*(?:@\[[^\n]*\]\s*)*"
|
|
r"(?:(?:private|protected|noncomputable|unsafe|local)\s+)*"
|
|
r"(theorem|lemma|example)\b"
|
|
)
|
|
|
|
def remove_imports(text: str) -> str:
|
|
"""Remove all lines starting with 'import' (ignoring leading spaces)."""
|
|
cleaned_lines = [
|
|
line for line in text.splitlines()
|
|
if not re.match(r'^\s*import\s+', line)
|
|
]
|
|
return "\n".join(cleaned_lines).strip()
|
|
|
|
def strip_trailing_backticks(text: str) -> str:
|
|
return re.sub(r"\s*```(?:\s*)$", "", text).strip()
|
|
|
|
_LEAN_FENCE_OPEN_RE = re.compile(r"```lean4?(?![A-Za-z0-9_])")
|
|
_LEAN_FENCE_RE = re.compile(r"```")
|
|
|
|
|
|
def closed_lean_blocks(text: str) -> list[str]:
|
|
"""Bodies of every *closed* ```lean4 block, left to right.
|
|
|
|
A fence that itself opens a new lean block does not close the current one --
|
|
the earlier block is an abandoned draft and is skipped. A plain regex would
|
|
read ```` ```lean4 A ```lean4 B ``` ```` as a single block bodied "A ```lean4 B".
|
|
"""
|
|
blocks: list[str] = []
|
|
pos = 0
|
|
while True:
|
|
opener = _LEAN_FENCE_OPEN_RE.search(text, pos)
|
|
if opener is None:
|
|
return blocks
|
|
body_start = opener.end()
|
|
closer = _LEAN_FENCE_RE.search(text, body_start)
|
|
if closer is None:
|
|
return blocks
|
|
if _LEAN_FENCE_OPEN_RE.match(text, closer.start()) is not None:
|
|
pos = closer.start()
|
|
continue
|
|
blocks.append(text[body_start : closer.start()])
|
|
pos = closer.end()
|
|
|
|
|
|
def extract_lean_code(text: str, idx: int = 0) -> str | None:
|
|
matches = closed_lean_blocks(text or "")
|
|
if not matches:
|
|
return None
|
|
lean_code = matches[idx].strip()
|
|
return strip_trailing_backticks(lean_code)
|
|
|
|
def is_deepseek_style_model(
|
|
model_id: str | None,
|
|
loadable_model_id: str | None = None,
|
|
) -> bool:
|
|
candidates = [candidate for candidate in (model_id, loadable_model_id) if candidate]
|
|
return any(
|
|
any(marker in candidate.lower() for marker in DEEPSEEK_STYLE_MODEL_MARKERS)
|
|
for candidate in candidates
|
|
)
|
|
|
|
def split_context_and_command(code: str) -> tuple[str, str, bool]:
|
|
text = (code or "").strip()
|
|
if not text:
|
|
return "", "", False
|
|
|
|
match = DECLARATION_START_RE.search(text)
|
|
if match is None:
|
|
return "", text, False
|
|
|
|
return text[: match.start()].strip(), text[match.start() :].strip(), True
|
|
|
|
def build_formal_statement(formal_statement: str, header: str = LEAN4_DEFAULT_HEADER, informal_statement: str = "") -> str:
|
|
"""
|
|
Prepare the prompt for the model by combining the header, informal statement, and formal statement.
|
|
The formal statement is wrapped in a lean4 code block for clarity.
|
|
"""
|
|
header = (header or LEAN4_DEFAULT_HEADER).strip()
|
|
|
|
return f"""
|
|
{header.strip()}
|
|
|
|
{informal_statement.strip()}
|
|
{formal_statement.strip()}
|
|
""".strip()
|
|
|
|
|
|
def preprocess_formal_statement(formal_statement: str) -> str:
|
|
"""
|
|
Preprocess the formal statement to ensure it ends with 'sorry' if it ends with 'by' or ':='.
|
|
This is a safety measure to ensure the prompt is well-formed for the model.
|
|
"""
|
|
formal_statement = formal_statement.strip()
|
|
if formal_statement.endswith("by"):
|
|
formal_statement += " sorry"
|
|
elif formal_statement.endswith(":="):
|
|
formal_statement += " by sorry"
|
|
elif not formal_statement.endswith("sorry"):
|
|
raise ValueError(f"Formal statement must end with 'sorry' or 'by' for safety. Got:\n{formal_statement}")
|
|
|
|
return formal_statement
|
|
|
|
PROMPT_INTRO = "Solve the following problem with Lean 4 code and explanatory comments:"
|
|
# Phase-1 instructions: formalmathatepfl/sft_classic and formalmathatepfl/feedback_data_training.
|
|
PLAIN_INSTRUCTION = (
|
|
"Replace every sorry statement with an appropriate proof. "
|
|
"Provide a complete solution in the lean4 code block."
|
|
)
|
|
FEEDBACK_CLAUSE = (
|
|
", annotated with Lean 4 compiler feedback blocks in the form "
|
|
"`/- <feedback> ... </feedback> -/` at the appropriate locations"
|
|
)
|
|
FEEDBACK_INSTRUCTION = PLAIN_INSTRUCTION[:-1] + FEEDBACK_CLAUSE + "."
|
|
# Phase-2 (reasoning) instructions: `question` / `question_feedback` of formalmathatepfl/reasoning_data
|
|
# and formalmathatepfl/rl_data. Byte-identical to LEARNING_PLAIN_TAIL / LEARNING_FEEDBACK_TAIL in
|
|
# annotate_numina/unify_reasoning_data.py (tests/test_prompt_utils_reasoning.py pins the text).
|
|
REASONING_SECTIONS = (
|
|
"[Problem Restatement], [Useful Observations], [Reasoning Path], [Key Derivation], "
|
|
"[Proof Type], [Formal Anchors] and [Final Solution]"
|
|
)
|
|
REPAIR_SECTIONS = "[Error Analysis], [Revised Reasoning] and [Updated Proof Plan]"
|
|
REASONING_INSTRUCTION = (
|
|
"Replace every sorry statement with an appropriate proof. Before writing any Lean code, give an "
|
|
f"informal solution with the sections {REASONING_SECTIONS}. Then provide the complete solution in "
|
|
f"the lean4 code block. If a previous attempt failed, first give {REPAIR_SECTIONS}, then provide "
|
|
"the corrected solution."
|
|
)
|
|
REASONING_FEEDBACK_INSTRUCTION = (
|
|
"Replace every sorry statement with an appropriate proof. Before writing any Lean code, give an "
|
|
f"informal solution with the sections {REASONING_SECTIONS}. Then provide the complete solution in "
|
|
f"the lean4 code block{FEEDBACK_CLAUSE}. If a previous attempt failed, first give {REPAIR_SECTIONS}, "
|
|
f"then provide the corrected solution{FEEDBACK_CLAUSE}."
|
|
)
|
|
INSTRUCTIONS = (REASONING_FEEDBACK_INSTRUCTION, REASONING_INSTRUCTION, FEEDBACK_INSTRUCTION, PLAIN_INSTRUCTION)
|
|
|
|
|
|
def instruction_for(tgat_prompt: bool = False, reasoning: bool = False) -> str:
|
|
if reasoning:
|
|
return REASONING_FEEDBACK_INSTRUCTION if tgat_prompt else REASONING_INSTRUCTION
|
|
return FEEDBACK_INSTRUCTION if tgat_prompt else PLAIN_INSTRUCTION
|
|
|
|
|
|
def prompt_with_instruction(formal_statement: str, instruction: str) -> str:
|
|
"""The SFT prompt layout: intro, the statement in a lean4 block, the instruction, and the
|
|
trailing newline every training prompt ends with (the model starts its answer at column 0)."""
|
|
return f"{PROMPT_INTRO}\n\n```lean4\n{formal_statement}\n```\n\n{instruction}\n"
|
|
|
|
|
|
def with_instruction(prompt: str, tgat_prompt: bool = False, reasoning: bool = False) -> str:
|
|
"""Re-target a prebuilt SFT-style prompt (e.g. a dataset `question`) to the requested
|
|
instruction. Prompts without a known instruction sentence are returned unchanged."""
|
|
wanted = instruction_for(tgat_prompt=tgat_prompt, reasoning=reasoning)
|
|
for present in INSTRUCTIONS: # longest first: the plain sentence is a prefix of the others
|
|
if present in prompt:
|
|
return prompt.replace(present, wanted, 1)
|
|
return prompt
|
|
|
|
|
|
def prompt_template(formal_statement: str) -> str:
|
|
return prompt_with_instruction(formal_statement, PLAIN_INSTRUCTION)
|
|
|
|
|
|
def prompt_template_feedback(formal_statement: str) -> str:
|
|
return prompt_with_instruction(formal_statement, FEEDBACK_INSTRUCTION)
|
|
|
|
|
|
def prompt_template_reasoning(formal_statement: str) -> str:
|
|
return prompt_with_instruction(formal_statement, REASONING_INSTRUCTION)
|
|
|
|
|
|
def prompt_template_reasoning_feedback(formal_statement: str) -> str:
|
|
return prompt_with_instruction(formal_statement, REASONING_FEEDBACK_INSTRUCTION)
|
|
|
|
|
|
def deepseek_prompt_template(formal_statement: str) -> str:
|
|
return f"""Complete the following Lean 4 code:
|
|
|
|
```lean4
|
|
{formal_statement}
|
|
```
|
|
|
|
Before producing the Lean 4 code to formally prove the given theorem, provide a detailed proof plan outlining the main proof steps and strategies.
|
|
The plan should highlight key ideas, intermediate lemmas, and proof structures that will guide the construction of the final formal proof.
|
|
""".strip()
|
|
|
|
|
|
def informal_statement_from_example(
|
|
example: dict[str, Any],
|
|
use_informal_statement: bool,
|
|
) -> str:
|
|
if not use_informal_statement:
|
|
return ""
|
|
for key in ("informal_prefix", "informal_statement"):
|
|
value = example.get(key)
|
|
if isinstance(value, str) and value.strip():
|
|
return value.strip()
|
|
return ""
|
|
|
|
|
|
def prover_prompt_template(
|
|
formal_statement: str,
|
|
is_deepseek: bool = False,
|
|
tgat_prompt: bool = False,
|
|
reasoning: bool = False,
|
|
) -> str:
|
|
"""`reasoning` asks for the phase-2 informal solution sections before the proof and
|
|
`tgat_prompt` for the feedback-annotated proof; the two combine. The DeepSeek layout is
|
|
used only when neither is requested."""
|
|
if reasoning or tgat_prompt:
|
|
return prompt_with_instruction(formal_statement, instruction_for(tgat_prompt=tgat_prompt, reasoning=reasoning))
|
|
if is_deepseek:
|
|
return deepseek_prompt_template(formal_statement=formal_statement)
|
|
return prompt_template(formal_statement=formal_statement)
|
|
|
|
|
|
def build_prover_prompt_for_example(
|
|
example: dict[str, Any],
|
|
use_informal_statement: bool,
|
|
is_deepseek: bool,
|
|
tgat_prompt: bool = False,
|
|
reasoning: bool = False,
|
|
) -> tuple[str, str]:
|
|
formal_statement = build_formal_statement(
|
|
formal_statement=preprocess_formal_statement(example["formal_statement"]),
|
|
header=example.get("header", None),
|
|
informal_statement=informal_statement_from_example(
|
|
example,
|
|
use_informal_statement=use_informal_statement,
|
|
),
|
|
)
|
|
prompt = prover_prompt_template(
|
|
formal_statement=formal_statement,
|
|
is_deepseek=is_deepseek,
|
|
tgat_prompt=tgat_prompt,
|
|
reasoning=reasoning,
|
|
)
|
|
return formal_statement, prompt
|
|
|
|
|
|
def parse_generated_lean(
|
|
raw_text: str,
|
|
example: dict[str, Any] | None = None,
|
|
is_deepseek: bool = False,
|
|
missing_code: str | None = None,
|
|
) -> str | None:
|
|
lean_code = extract_lean_code(raw_text, idx=-1 if is_deepseek else -1)
|
|
if lean_code is None:
|
|
if missing_code is None:
|
|
return None
|
|
lean_code = missing_code
|
|
|
|
if is_deepseek and isinstance(example, dict):
|
|
header = example.get("header")
|
|
if isinstance(header, str) and header.strip():
|
|
return f"{header.strip()}\n\n{lean_code.strip()}".strip()
|
|
|
|
return lean_code.strip()
|
|
|
|
|
|
def default_lean_stop_sequences(is_deepseek: bool) -> list[str] | None:
|
|
return None if is_deepseek else None #["``` ", "```\n"]
|
|
|
|
|
|
def default_generation_max_tokens(
|
|
is_deepseek: bool,
|
|
deepseek_max_tokens: int = DEFAULT_DEEPSEEK_GENERATION_MAX_TOKENS,
|
|
standard_max_tokens: int = DEFAULT_STANDARD_GENERATION_MAX_TOKENS,
|
|
) -> int:
|
|
return deepseek_max_tokens if is_deepseek else standard_max_tokens
|
|
|
|
|
|
def default_generation_temperature(max_attempts: int) -> float:
|
|
if max_attempts <= 0:
|
|
raise ValueError(f"max_attempts must be positive, got {max_attempts}")
|
|
if max_attempts == 1:
|
|
return 0.3
|
|
if max_attempts <= 16:
|
|
return 0.5
|
|
if max_attempts <= 32:
|
|
return 0.7
|
|
if max_attempts <= 64:
|
|
return 0.8
|
|
return 1.0
|