83 lines
3.8 KiB
Python
83 lines
3.8 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Generate Lean proofs with the paper's prompt and last-complete-block parser."""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from prompt_utils import (
|
||
|
|
LEAN4_DEFAULT_HEADER, DEFAULT_GENERATION_TOP_P,
|
||
|
|
build_prover_prompt_for_example, default_generation_temperature,
|
||
|
|
default_generation_max_tokens, default_lean_stop_sequences,
|
||
|
|
parse_generated_lean, split_context_and_command,
|
||
|
|
)
|
||
|
|
|
||
|
|
DEFAULT_MODEL = 'EPFL-LARA/lean-world-model-qwen3-4b-grpo'
|
||
|
|
FEEDBACK = True
|
||
|
|
REASONING = False
|
||
|
|
|
||
|
|
|
||
|
|
def build_prompt(source: str) -> str:
|
||
|
|
context, declaration, found = split_context_and_command(source)
|
||
|
|
if not found:
|
||
|
|
raise ValueError("Supply a Lean theorem, lemma, or example ending in 'by sorry'.")
|
||
|
|
example = {"formal_statement": declaration, "header": context or LEAN4_DEFAULT_HEADER}
|
||
|
|
_, prompt = build_prover_prompt_for_example(
|
||
|
|
example, use_informal_statement=False, is_deepseek=False,
|
||
|
|
tgat_prompt=FEEDBACK, reasoning=REASONING,
|
||
|
|
)
|
||
|
|
return prompt
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
||
|
|
parser.add_argument("files", nargs="+", type=Path, help="Lean problem files, including imports and a theorem to complete")
|
||
|
|
parser.add_argument("--model", default=DEFAULT_MODEL)
|
||
|
|
parser.add_argument("--samples", type=int, default=1)
|
||
|
|
parser.add_argument("--temperature", type=float)
|
||
|
|
parser.add_argument("--top-p", type=float, default=DEFAULT_GENERATION_TOP_P)
|
||
|
|
parser.add_argument("--max-new-tokens", type=int, default=default_generation_max_tokens(False))
|
||
|
|
parser.add_argument("--tensor-parallel-size", type=int, default=1)
|
||
|
|
parser.add_argument("--seed", type=int, default=0)
|
||
|
|
parser.add_argument("--print-prompt", action="store_true", help="Show the exact raw prompt without loading the model")
|
||
|
|
args = parser.parse_args()
|
||
|
|
if args.samples < 1 or args.max_new_tokens < 1:
|
||
|
|
parser.error("--samples and --max-new-tokens must be positive")
|
||
|
|
prompts = [build_prompt(path.read_text(encoding="utf-8")) for path in args.files]
|
||
|
|
if args.print_prompt:
|
||
|
|
for prompt in prompts:
|
||
|
|
print(prompt, end="")
|
||
|
|
return
|
||
|
|
|
||
|
|
from vllm import LLM, SamplingParams
|
||
|
|
|
||
|
|
llm = LLM(model=args.model, tokenizer=args.model,
|
||
|
|
tensor_parallel_size=args.tensor_parallel_size, dtype="auto",
|
||
|
|
max_model_len=32768, trust_remote_code=False, seed=args.seed,
|
||
|
|
gpu_memory_utilization=0.90)
|
||
|
|
tokenizer = llm.get_tokenizer()
|
||
|
|
tokenized = [tokenizer.encode(prompt, add_special_tokens=False) for prompt in prompts]
|
||
|
|
if any(len(ids) >= 32768 for ids in tokenized):
|
||
|
|
raise ValueError("The problem prompt must leave space in the 32,768-token context for a proof.")
|
||
|
|
temperature = args.temperature if args.temperature is not None else default_generation_temperature(args.samples)
|
||
|
|
params = [SamplingParams(n=args.samples, temperature=temperature,
|
||
|
|
top_p=args.top_p,
|
||
|
|
max_tokens=min(args.max_new_tokens, 32768-len(ids)),
|
||
|
|
stop=default_lean_stop_sequences(False),
|
||
|
|
include_stop_str_in_output=True, seed=args.seed)
|
||
|
|
for ids in tokenized]
|
||
|
|
outputs = llm.generate([{"prompt_token_ids": ids} for ids in tokenized], params)
|
||
|
|
for path, output in zip(args.files, outputs, strict=True):
|
||
|
|
for i, completion in enumerate(output.outputs):
|
||
|
|
print(json.dumps({
|
||
|
|
"source": str(path), "sample_id": i,
|
||
|
|
"response": completion.text,
|
||
|
|
"lean_code": parse_generated_lean(completion.text, is_deepseek=False),
|
||
|
|
"finish_reason": completion.finish_reason,
|
||
|
|
}, ensure_ascii=False), flush=True)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|