初始化项目，由ModelHub XC社区提供模型
Model: EPFL-LARA/lean-world-model-qwen3-4b-grpo Source: Original Platform
This commit is contained in:
ModelHub XC
303
prompt_utils.py
Normal file
303
prompt_utils.py
Normal file
@@ -0,0 +1,303 @@
|
||||
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
|
||||
Reference in New Issue
Block a user