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

---
library_name: transformers
pipeline_tag: text-generation
base_model: Qwen/Qwen3-4B-Base
base_model_relation: finetune
language:
- en
tags:
- lean
- lean4
- theorem-proving
- qwen3
- lean-world-model
- grpo
datasets:
- 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:

```bash
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`:

```lean4
import Mathlib

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

The complete Transformers example is:

```python
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:

```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
```

<details>
<summary>Full generated response</summary>

````text
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
```
````

</details>

[Full response](https://huggingface.co/EPFL-LARA/lean-world-model-qwen3-4b-grpo/blob/main/example_response.txt) ·
[Input/output JSON](https://huggingface.co/EPFL-LARA/lean-world-model-qwen3-4b-grpo/blob/main/example_output.json) ·
[Lean proof](https://huggingface.co/EPFL-LARA/lean-world-model-qwen3-4b-grpo/blob/main/example_proof.lean)

## vLLM

For multiple samples and longer output budgets:

```bash
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:

````text
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](https://huggingface.co/collections/EPFL-LARA/formal-theorem-proving-with-lean-world-models-6ad2b4aa4400b3e2933f380d)