lean-world-model-qwen3-4b-grpo/README.md

library_name, pipeline_tag, base_model, base_model_relation, language, tags, datasets

library_name	pipeline_tag	base_model	base_model_relation	language	tags	datasets
transformers	text-generation	Qwen/Qwen3-4B-Base	finetune	en	lean lean4 theorem-proving qwen3 lean-world-model grpo	EPFL-LARA/lean-world-model-cpt EPFL-LARA/lean-world-model-sft Goedel-LM/RL_dataset_V2

Lean World-Model Prover — Qwen3-4B

Paper model: WM4 · Stage: After GRPO

Generates Lean proofs with predicted goal states and compiler diagnostics in inline comments, with support for repair attempts within one response.

Hugging Face Transformers

Load the published weights from the Hub and generate on your GPU:

pip install 'transformers==4.57.3' accelerate torch
hf download EPFL-LARA/lean-world-model-qwen3-4b-grpo example_transformers.py inference.py prompt_utils.py Example.lean --local-dir lean-world-model-qwen3-4b-grpo
cd lean-world-model-qwen3-4b-grpo
python example_transformers.py

The included Example.lean is the input; replace it with your own imports, context, and theorem ending in by sorry:

import Mathlib

theorem demo (n : ℕ) : n + 0 = n := by sorry

The complete Transformers example is:

from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed

from inference import build_prompt

MODEL_ID = 'EPFL-LARA/lean-world-model-qwen3-4b-grpo'

source = Path("Example.lean").read_text(encoding="utf-8")
prompt = build_prompt(source)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, dtype=torch.bfloat16, device_map="auto", attn_implementation="sdpa",
).eval()
inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(model.device)
prompt_length = inputs["input_ids"].shape[1]
budget = min(2048, 32768 - prompt_length)
if budget <= 0:
    raise ValueError("The prompt must leave room for a proof in the 32,768-token context.")

set_seed(0)
with torch.inference_mode():
    output = model.generate(
        **inputs, max_new_tokens=budget, do_sample=True,
        temperature=0.3, top_p=0.95, top_k=0, min_p=0.0,
        repetition_penalty=1.0, num_beams=1,
        pad_token_id=tokenizer.pad_token_id,
    )
response = tokenizer.decode(output[0, prompt_length:], skip_special_tokens=True)
print(response, end="")

build_prompt supplies this checkpoint's exact raw completion prompt. Tokenization uses add_special_tokens=False; the context budget is 32,768 tokens. This compact example allows up to 2,048 generated tokens and prints only the generated response.

Example output

Recorded with the script above on an H100 in BF16, using seed 0, temperature 0.3, and top-p 0.95. This is the extracted final proof, checked with Lean 4.15 and Mathlib:

import Mathlib

-- Theorem: For any natural number n, n + 0 = n.
theorem demo (n : ℕ) : n + 0 = n := by
  /- <feedback>
  -- proof state:
  n : ℕ
  ⊢ n + 0 = n
  </feedback> -/
  induction n with
  /- <feedback>
  -- proof state:
  case zero
  ⊢ 0 + 0 = 0
  </feedback> -/
  | zero => rfl
  /- <feedback>
  -- proof state:
  ih : n + 0 = n
  case succ
  ⊢ n + 1 + 0 = n + 1
  </feedback> -/
  | succ n ih => simp_all

Full generated response

Here is the complete Lean 4 proof annotated with Lean 4 compiler feedback blocks:

```lean4
import Mathlib

-- Theorem: For any natural number n, n + 0 = n.
theorem demo (n : ℕ) : n + 0 = n := by
  /- <feedback>
  -- proof state:
  n : ℕ
  ⊢ n + 0 = n
  </feedback> -/
  induction n with
  /- <feedback>
  -- proof state:
  case zero
  ⊢ 0 + 0 = 0
  </feedback> -/
  | zero => rfl
  /- <feedback>
  -- proof state:
  ih : n + 0 = n
  case succ
  ⊢ n + 1 + 0 = n + 1
  </feedback> -/
  | succ n ih => simp_all
```

Full response · Input/output JSON · Lean proof

vLLM

For multiple samples and longer output budgets:

pip install 'vllm==0.12.0'
hf download EPFL-LARA/lean-world-model-qwen3-4b-grpo inference.py prompt_utils.py Example.lean --local-dir lean-world-model-qwen3-4b-grpo
python lean-world-model-qwen3-4b-grpo/inference.py lean-world-model-qwen3-4b-grpo/Example.lean --samples 1

The vLLM runner uses the same raw prompt. Its default temperature is 0.3 for one sample and 0.7 for 32 samples; top-p is 0.95. Output tokens fit within the context remaining after the prompt. Use --print-prompt to inspect the exact input without loading weights.

Prompt template

The feedback template follows evaluate_prover_model_run.py. Replace the placeholder with the imports, context, and theorem ending in by sorry, and keep the final newline:

Solve the following problem with Lean 4 code and explanatory comments:

```lean4
{formal_statement_with_imports}
```

Replace every sorry statement with an appropriate proof. Provide a complete solution in the lean4 code block, annotated with Lean 4 compiler feedback blocks in the form `/- <feedback> ... </feedback> -/` at the appropriate locations.

Expected response

Proofs appear in lean4 code fences with /- <feedback> ... </feedback> -/ comments containing predicted goal states and diagnostics. A response can contain successive corrected proof attempts.

Use the last complete Lean code block as the final proof candidate for Lean verification. The vLLM runner returns the full response, extracted lean_code, and finish_reason as JSONL.

Formal Theorem Proving with Lean World Models