Update README.md

README.md

@@ -1 +1,598 @@
-Pythagoras-Prover-Diffusion-4B
+---
+license: apache-2.0
+base_model:
+- Qwen/Qwen3-4B
+pipeline_tag: text-generation
+library_name: transformers
+datasets:
+- Pythagoras-LM/SFT_Dataset
+---
+<div align="center">
+  <img src="https://github.com/Pythagoras-LM/Pythagoras-LM.github.io/blob/main/figures/pythagoras-without-background.png?raw=true" width="240" alt="Pythagoras-Prover logo"><br>
+  <h1>Pythagoras-Prover</h1>
+</div>
+
+<div align="center">
+  
+[![Website](https://img.shields.io/badge/%F0%9F%A4%96%20Homepage-Pythagoras-536af5?color=536af5&logoColor=white)](https://pythagoras-lm.github.io/)
+[![Paper](https://img.shields.io/badge/Paper-arXiv-b31b1b?logo=arxiv&logoColor=white)](https://arxiv.org/abs/XXXX.XXXXX)
+[![GitHub](https://img.shields.io/badge/GitHub-Code-black.svg?logo=github)](https://github.com/Pythagoras-LM/Pythagoras-Prover)
+[![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+
+</div>
+
+
+## 1. Introduction
+
+We introduce **Pythagoras-Prover**, a compute-efficient family of open-source large language models for formal theorem proving in Lean 4. The family comprises two autoregressive provers at 4B and 32B parameters, together with **Pythagoras-Prover-Diffusion**, the first diffusion-based theorem prover, which iteratively refines Lean proofs at inference time. All three models are artefacts of a single methodological approach: a scalable, Lean-verified synthetic data pipeline. At its centre is **Augmented Lean Formalisation** (ALF), a structured mutation scheme that expands a verified seed corpus into formal variants without per-instance Lean compilation, then re-uses them as a self-distillation signal during training. This design lets careful data construction stand in for raw scale, closing much of the gap between small open provers and their largest counterparts — without relying on inference-time self-correction.
+
+<div align="center">
+  <img src="https://github.com/Pythagoras-LM/Pythagoras-LM.github.io/blob/main/figures/prover_fig1_v2_hi.png?raw=true" alt="Pythagoras-Prover benchmark overview" width="90%">
+</div>
+
+
+## 2. Model Summary
+
+---
+
+**A Lean-Verified Synthetic Data Pipeline**
+
+- Natural-language problems from general-math and competition sources are autoformalised into Lean and gated on the type-checker using predominantly sub-30B open models, with an auto-informalisation and alignment step discarding faithful-but-wrong formalisations to yield a verified seed corpus partitioned into easy, medium, and hard tiers.
+- A rubric-guided distillation stage re-prompts on each rejected instance to target the specific Lean type-checker error responsible for its failure, lifting autoformalisation success and roughly doubling the verified training set.
+
+---
+
+**Model Training**
+
+- LoRA-only supervised fine-tuning of Qwen3-4B and Qwen3-32B under an 8K context, paired with a dynamic proof-reasoning filter and a difficulty-ordered easy→medium→hard curriculum, followed by reinforcement learning with a Lean-compilation reward and a final continued-SFT stage on the ALF corpus.
+
+---
+
+**Augmented Lean Formalisation**
+- ALF emits one structured variant per category — simplification, generalisation, lemma proposal, proof-step decomposition, and reformulation — for every seed statement, replacing per-instance Lean verification with a cheap statement-alignment check and expanding the seed corpus into roughly 2M formal variants.
+- The post-RL prover proves the mutations, and these self-distilled proofs form a corpus that trains both the autoregressive and diffusion provers from a single recipe.
+
+---
+
+**The Smallest Efficient Open-Source Lean Theorem Prover**
+
+- We train Pythagoras-Prover-4B, one of the smallest and most compute-efficient open-source Lean theorem provers to date, reaching 86.07% on MiniF2F-Test at Pass@32 and surpassing the prior state of the art at a fraction of its parameter count.
+- Pythagoras-Prover-Diffusion adapts a block-diffusion formulation with a tactic-based masking objective aligned to the discrete reasoning steps of Lean — to our knowledge the first demonstration that a diffusion language model can verifiably solve Lean theorems at non-trivial rates.
+
+---
+
+The resulting models set a new bar for compute-efficient formal proving. **Pythagoras-Prover-32B** achieves state-of-the-art performance among open-source provers, reaching **93.03%** on MiniF2F-Test and solving **93 of 672** problems on PutnamBench, while **Pythagoras-Prover-4B** outperforms DeepSeek-Prover-V2-671B on MiniF2F-Test despite being roughly **167× smaller** — with no self-correction or test-time reinforcement learning. We additionally release **MiniF2F-ALF**, an ALF-mutated companion benchmark on which every evaluated prover degrades.
+
+## 3. Benchmark Performance
+
+We evaluate Pythagoras-Prover on three Lean 4 benchmarks — MiniF2F-Test, PutnamBench, and the MiniF2F-ALF benchmark we introduce — under a single unified protocol (Lean 4.9.0-rc1, a 30,000-token generation limit, and a verbatim-statement pass criterion). Across all three, Pythagoras-Prover matches or exceeds open-source provers an order of magnitude larger, and does so **without** inference-time self-correction or test-time reinforcement learning.
+
+
+<div align="center">
+  <table style="margin: 0 auto;">
+    <thead>
+      <tr>
+        <th>Method</th>
+        <th>#Params</th>
+        <th>Pass@32</th>
+        <th>Pass@1024</th>
+        <th>Best (N)</th>
+      </tr>
+    </thead>
+    <tbody>
+      <tr><td>Goedel-Prover-SFT</td><td>7B</td><td>57.6</td><td>–</td><td>62.7 (3200)</td></tr>
+      <tr><td>STP</td><td>7B</td><td>–</td><td>–</td><td>67.6 (25600)</td></tr>
+      <tr><td>Kimina-Prover-Preview-72B</td><td>72B</td><td>68.85</td><td>–</td><td>80.74 (8192)</td></tr>
+      <tr><td>DeepSeek-Prover-V2-7B</td><td>7B</td><td>75.6</td><td>–</td><td>82.0 (8192)</td></tr>
+      <tr><td>DeepSeek-Prover-V2-671B</td><td>671B</td><td>82.4</td><td>–</td><td>88.9 (8192)</td></tr>
+      <tr><td>Kimina-Prover-8B-Distill</td><td>8B</td><td>77.86</td><td>–</td><td>–</td></tr>
+      <tr><td>Kimina-Prover-70B</td><td>70B</td><td>84.0</td><td>87.7</td><td>92.2 (TTRL)</td></tr>
+      <tr><td>Goedel-Prover-V2-8B</td><td>8B</td><td>84.6</td><td>87.9</td><td>90.2 (8192)</td></tr>
+      <tr><td>&nbsp;&nbsp;+ Self-Correction</td><td>8B</td><td>86.7</td><td>89.3</td><td>–</td></tr>
+      <tr><td>Goedel-Prover-V2-32B</td><td>32B</td><td>88.1</td><td>91.8</td><td>92.2 (8192)</td></tr>
+      <tr><td>&nbsp;&nbsp;+ Self-Correction</td><td>32B</td><td>90.4</td><td>92.6</td><td>–</td></tr>
+      <tr><td><strong>Pythagoras-Prover-4B</strong></td><td>4B</td><td><strong>86.07</strong></td><td><strong>88.11</strong></td><td><strong>89.75 (2048)</strong></td></tr>
+      <tr><td><strong>Pythagoras-Prover-32B</strong></td><td>32B</td><td><strong>89.75</strong></td><td><strong>92.62</strong></td><td><strong>93.03 (2048)</strong></td></tr>
+    </tbody>
+  </table>
+  <!-- table caption -->
+  <caption align="bottom"><strong>Table 1</strong>: <em>Pythagoras-Prover-4B exceeds DeepSeek-Prover-V2-671B's pass@8192 result (88.9%) at pass@2048 — a quarter of the budget and ~167× fewer parameters. Pythagoras-Prover-32B sets the strongest reported MiniF2F-Test pass rate without self-correction or test-time RL.</em></caption>
+</div>
+
+
+<br>
+
+<div align="center">
+  <table style="margin: 0 auto;">
+    <thead>
+      <tr>
+        <th>#</th>
+        <th>Model</th>
+        <th>num-solved</th>
+        <th>compute</th>
+      </tr>
+    </thead>
+    <tbody>
+      <tr><td>1</td><td><strong>Pythagoras-Prover-32B</strong></td><td><strong>93</strong></td><td><strong>Pass@2048</strong></td></tr>
+      <tr><td>1</td><td><strong>Pythagoras-Prover-32B</strong></td><td><strong>59</strong></td><td><strong>Pass@64</strong></td></tr>
+      <tr><td>1</td><td><strong>Pythagoras-Prover-32B</strong></td><td><strong>48</strong></td><td><strong>Pass@32</strong></td></tr>
+      <tr><td>2</td><td>Goedel-Prover-V2-32B (self-correction mode)</td><td>86</td><td>Pass@184</td></tr>
+      <tr><td>2</td><td>Goedel-Prover-V2-32B (self-correction mode)</td><td>57</td><td>Pass@32</td></tr>
+      <tr><td>2</td><td>Goedel-Prover-V2-32B</td><td>43</td><td>Pass@32</td></tr>
+      <tr><td>3</td><td>DeepSeek-Prover-V2-671B</td><td>47</td><td>Pass@1024</td></tr>
+      <tr><td>3</td><td>DeepSeek-Prover-V2-671B</td><td>22</td><td>Pass@32</td></tr>
+      <tr><td>4</td><td>DSP+</td><td>23</td><td>Pass@128</td></tr>
+      <tr><td>5</td><td>Bourbaki</td><td>14</td><td>Pass@512</td></tr>
+      <tr><td>6</td><td>Kimina-Prover-7B-Distill</td><td>10</td><td>Pass@192</td></tr>
+      <tr><td>7</td><td>Self-play Theorem Prover</td><td>8</td><td>Pass@3200</td></tr>
+      <tr><td>8</td><td>Goedel-Prover-SFT</td><td>7</td><td>Pass@512</td></tr>
+      <tr><td>9</td><td>ABEL (closed-source)</td><td>7</td><td>Pass@596</td></tr>
+    </tbody>
+  </table>
+  <!-- table caption -->
+  <caption align="bottom"><strong>Table 2</strong>: <em>PutnamBench leaderboard (problems solved out of 657). Pythagoras-Prover-32B takes the top rank, solving 93 problems at Pass@2048 — 7 more than the previous best (Goedel-Prover-V2-32B, 86 at Pass@184 in self-correction mode) and nearly double DeepSeek-Prover-V2-671B's 47 at Pass@1024, despite being roughly 20× smaller. Seed-Prover (331 solved) is omitted from the ranked rows as it is closed-source with undisclosed test-time compute.</em></caption>
+</div>
+
+<br>
+
+<div align="center">
+  <table style="margin: 0 auto;">
+    <thead>
+      <tr>
+        <th>Model</th>
+        <th>Pass@32</th>
+      </tr>
+    </thead>
+    <tbody>
+      <tr><td>DeepSeek-Prover-V2-671B</td><td>79.71</td></tr>
+      <tr><td>Goedel-Prover-V2-8B</td><td>82.58</td></tr>
+      <tr><td>Goedel-Prover-V2-32B</td><td>83.61</td></tr>
+      <tr><td><strong>Pythagoras-Prover-4B</strong></td><td><strong>83.19</strong></td></tr>
+      <tr><td><strong>Pythagoras-Prover-32B</strong></td><td><strong>85.04</strong></td></tr>
+    </tbody>
+  </table>
+  <!-- table caption -->
+  <caption align="bottom"><strong>Table 3</strong>: <em>Performance of current state-of-the-art provers on MiniF2F-ALF (Pass@32, %). As MiniF2F-ALF is introduced in this work, all results are evaluated by us under a unified setup.</em></caption>
+</div>
+
+<br>
+
+<div align="center">
+  <table style="margin: 0 auto;">
+    <thead>
+      <tr>
+        <th>Model</th>
+        <th>Performance</th>
+      </tr>
+    </thead>
+    <tbody>
+      <tr><td>Pythagoras-Prover</td><td>86.1%</td></tr>
+      <tr><td>Pythagoras-Prover∗</td><td>74.6%</td></tr>
+      <tr><td><strong>Diffusion Pythagoras-Prover</strong></td><td><strong>63.3%</strong></td></tr>
+    </tbody>
+  </table>
+  <caption align="bottom"><strong>Table 4</strong>: <em>Diffusion Theorem Proving Performance (Pass@32). ∗ denotes the setting where training tokens are restricted to 4096 and evaluation is performed solely at 8192 tokens.</em></caption>
+</div>
+
+## 4. Model & Dataset Downloads
+
+We release our Pythagoras-Prover models, the training dataset and the new MiniF2F-ALF benchmark for future future research.
+
+<div align="center">
+  
+| Model | Download |
+| -------- | -------- |
+|    Pythagoras-Prover-32B    |   Coming Soon   |
+|    Pythagoras-Prover-4B    |   [🤗Download](https://huggingface.co/Pythagoras-LM/Pythagoras-Prover-4B)    |
+|    Pythagoras-Prover-Diffusion-4B    |   [🤗Download](https://huggingface.co/Pythagoras-LM/Pythagoras-Prover-Diffusion-4B)    |
+
+</div>
+
+
+<div align="center">
+
+| Dataset | Download |
+| -------- | -------- |
+|   Pythagoras-Prover-SFT    |   [🤗Download](https://huggingface.co/datasets/Pythagoras-LM/SFT_Dataset)   |
+|   Pythagoras-Prover-Distill-4B    |   Coming Soon    |
+|    Pythagoras-Prover-Distill-32B    |   Coming Soon    |
+</div>
+
+## 5. Quick Start
+
+This checkpoint is an A2D/BD3LM diffusion prover. It should be decoded with a diffusion sampler rather than `AutoModelForCausalLM.generate`. The snippet below is self-contained: it registers the A2D Qwen3 model class, implements the BD3LM block-diffusion sampler used for MiniF2F evaluation, loads the checkpoint with Transformers, and runs one Lean theorem-proving prompt. The decoding settings match the MiniF2F evaluation setting used for the reported Pass@32 result: bf16, SDPA attention, 8192 generated tokens, 8192 diffusion steps, and block size 8.
+
+````python
+import copy
+import math
+import torch
+import torch.nn.functional as F
+import transformers
+from transformers import AutoConfig, AutoModel, AutoModelForMaskedLM, AutoTokenizer
+from transformers.cache_utils import DynamicCache
+from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
+from transformers.modeling_outputs import BaseModelOutputWithPast
+
+
+class A2DQwen3Config(transformers.Qwen3Config):
+    model_type = "a2d-qwen3"
+
+
+class A2DQwen3Model(transformers.Qwen3Model):
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        position_ids=None,
+        past_key_values=None,
+        inputs_embeds=None,
+        use_cache=None,
+        cache_position=None,
+        **kwargs,
+    ):
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            if attention_mask is None:
+                attention_mask = torch.ones(
+                    inputs_embeds.shape[:2],
+                    device=inputs_embeds.device,
+                    dtype=torch.long,
+                )
+
+            if not (isinstance(attention_mask, torch.Tensor) and attention_mask.ndim == 4):
+                attention_mask = _prepare_4d_attention_mask(attention_mask, self.dtype)
+
+            causal_mask_mapping = {"full_attention": attention_mask}
+            if self.has_sliding_layers:
+                causal_mask_mapping["sliding_attention"] = attention_mask
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask_mapping[decoder_layer.attention_type],
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+        )
+
+
+class A2DQwen3LMHeadModel(transformers.Qwen3ForCausalLM):
+    config_class = A2DQwen3Config
+
+    def __init__(self, config):
+        transformers.Qwen3PreTrainedModel.__init__(self, config)
+        self.model = A2DQwen3Model(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = torch.nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+
+AutoConfig.register("a2d-qwen3", A2DQwen3Config)
+AutoModel.register(A2DQwen3Config, A2DQwen3LMHeadModel)
+AutoModelForMaskedLM.register(A2DQwen3Config, A2DQwen3LMHeadModel)
+
+
+def prepare_for_sampling(x, block_size, pad_token_id):
+    batch_size, seq_len = x.shape
+    device = x.device
+
+    valid = x != pad_token_id
+    pos_raw = torch.cumsum(valid.to(torch.long), dim=-1)
+    logical_pos = pos_raw - 1
+    position_ids = torch.where(valid, logical_pos, torch.zeros_like(logical_pos)).to(
+        device=device,
+        dtype=torch.long,
+    )
+
+    pos = torch.arange(seq_len, device=device)
+    block_ids = torch.div(pos, block_size, rounding_mode="floor")
+    block_ids = block_ids.view(1, seq_len).expand(batch_size, -1)
+    block_ids = torch.where(valid, block_ids, torch.full_like(block_ids, -1))
+
+    bid_q = block_ids.view(batch_size, 1, seq_len, 1)
+    bid_k = block_ids.view(batch_size, 1, 1, seq_len)
+    attention_mask = (bid_k <= bid_q) & (bid_q >= 0) & (bid_k >= 0)
+    return attention_mask, position_ids
+
+
+def add_gumbel_noise(logits, temperature):
+    if temperature == 0:
+        return logits
+    logits = logits.to(torch.float64)
+    noise = torch.rand_like(logits, dtype=torch.float64)
+    return logits.exp() / ((-torch.log(noise)) ** temperature)
+
+
+def get_num_transfer_tokens(mask_index, steps):
+    mask_num = mask_index.sum(dim=1, keepdim=True)
+    num_transfer_tokens = torch.zeros(
+        mask_num.size(0),
+        steps,
+        device=mask_index.device,
+        dtype=torch.int64,
+    )
+    for i in range(mask_num.size(0)):
+        for t, s, j in zip(range(steps, 0, -1), range(steps - 1, -1, -1), range(steps)):
+            reverse_transfer_prob = 1 - (s / t)
+            transfer = torch.round(mask_num[i, 0].to(torch.float64) * reverse_transfer_prob).to(torch.int64)
+            transfer = torch.minimum(transfer, mask_num[i, 0])
+            num_transfer_tokens[i, j] = transfer
+            mask_num[i, 0] -= transfer
+            if mask_num[i, 0].item() == 0:
+                break
+
+    rows = []
+    max_len = 0
+    for row in num_transfer_tokens:
+        nonzero = row[row > 0]
+        rows.append(nonzero)
+        max_len = max(max_len, nonzero.numel())
+
+    padded_rows = []
+    for row in rows:
+        if row.numel() < max_len:
+            row = torch.cat([row, torch.zeros(max_len - row.numel(), dtype=row.dtype, device=row.device)])
+        padded_rows.append(row)
+    return torch.stack(padded_rows, dim=0)
+
+
+def diffusion_step_block(logits, x_block, mask_block, num_transfer_step, temperature):
+    batch_size, block_len, _ = logits.shape
+    logits_with_noise = add_gumbel_noise(logits, temperature=temperature)
+    x0 = torch.argmax(logits_with_noise, dim=-1)
+
+    probs = F.softmax(logits, dim=-1)
+    x0_probs = torch.gather(probs, dim=-1, index=x0.unsqueeze(-1)).squeeze(-1)
+
+    x0 = torch.where(mask_block, x0, x_block)
+    neg_inf = torch.full_like(x0_probs, -float("inf"))
+    confidence = torch.where(mask_block, x0_probs, neg_inf)
+
+    transfer = torch.zeros_like(x0, dtype=torch.bool)
+    for batch_idx in range(batch_size):
+        k = int(num_transfer_step[batch_idx].item())
+        if k <= 0:
+            continue
+        valid_count = (confidence[batch_idx] > -float("inf")).sum().item()
+        if valid_count == 0:
+            continue
+        _, selected = torch.topk(confidence[batch_idx], min(k, valid_count))
+        transfer[batch_idx, selected] = True
+
+    x_block_new = x_block.clone()
+    x_block_new[transfer] = x0[transfer]
+    return x_block_new
+
+
+@torch.no_grad()
+def generate_bd3lm(
+    model,
+    tokenizer,
+    input_ids,
+    max_new_tokens=8192,
+    steps=8192,
+    block_size=8,
+    temperature=1.0,
+):
+    if isinstance(input_ids, torch.Tensor):
+        input_ids = input_ids.tolist()
+    if input_ids and isinstance(input_ids[0], int):
+        input_ids = [input_ids]
+
+    device = model.device
+    mask_id = tokenizer.mask_token_id
+    pad_id = tokenizer.pad_token_id
+    eos_id = tokenizer.eos_token_id
+
+    inputs = [torch.tensor(ids, dtype=torch.long, device=device) for ids in input_ids]
+    prompt_lens = [ids.shape[0] for ids in inputs]
+    max_prompt_len = max(prompt_lens)
+    padded_prompt_len = math.ceil(max_prompt_len / block_size) * block_size
+
+    x = torch.full(
+        (len(inputs), padded_prompt_len),
+        pad_id,
+        dtype=torch.long,
+        device=device,
+    )
+    for batch_idx, ids in enumerate(inputs):
+        x[batch_idx, padded_prompt_len - ids.shape[0] : padded_prompt_len] = ids
+
+    num_blocks = math.ceil(max_new_tokens / block_size)
+    steps_per_block = math.ceil(steps / num_blocks)
+    done = torch.zeros((len(inputs),), dtype=torch.bool, device=device)
+    generated = 0
+
+    for _ in range(num_blocks):
+        if done.all():
+            break
+
+        prefix_len = x.shape[1]
+        cur_block_len = min(block_size, max_new_tokens - generated)
+        if cur_block_len <= 0:
+            break
+
+        prefix_attention_mask, prefix_position_ids = prepare_for_sampling(
+            x,
+            block_size=block_size,
+            pad_token_id=pad_id,
+        )
+        prefix_outputs = model(
+            x,
+            attention_mask=prefix_attention_mask,
+            position_ids=prefix_position_ids,
+            use_cache=True,
+        )
+        prefix_cache = prefix_outputs.past_key_values
+
+        new_block = torch.full(
+            (x.shape[0], cur_block_len),
+            mask_id,
+            dtype=torch.long,
+            device=device,
+        )
+        x = torch.cat([x, new_block], dim=1)
+        total_len = x.shape[1]
+
+        block_mask_index = x[:, -cur_block_len:] == mask_id
+        num_transfer_tokens = get_num_transfer_tokens(block_mask_index, steps_per_block)
+
+        full_attention_mask, full_position_ids = prepare_for_sampling(
+            x,
+            block_size=block_size,
+            pad_token_id=pad_id,
+        )
+        block_attention_mask = full_attention_mask[:, :, prefix_len:total_len, :]
+        block_position_ids = full_position_ids[:, prefix_len:total_len]
+
+        for step_idx in range(num_transfer_tokens.size(1)):
+            x_block = x[:, prefix_len:total_len]
+            mask_block = x_block == mask_id
+            if not mask_block.any():
+                break
+
+            logits = model(
+                x_block,
+                attention_mask=block_attention_mask,
+                position_ids=block_position_ids,
+                past_key_values=copy.deepcopy(prefix_cache),
+                use_cache=False,
+            ).logits
+
+            x[:, prefix_len:total_len] = diffusion_step_block(
+                logits=logits,
+                x_block=x_block,
+                mask_block=mask_block,
+                num_transfer_step=num_transfer_tokens[:, step_idx],
+                temperature=temperature,
+            )
+
+        if eos_id is not None:
+            done = done | (x[:, prefix_len:total_len] == eos_id).any(dim=1)
+        generated += cur_block_len
+
+    return x
+
+
+def trim_completion(tokenizer, sequence_ids, prompt_ids):
+    sequence_ids = list(sequence_ids)
+    while sequence_ids and sequence_ids[0] == tokenizer.pad_token_id:
+        sequence_ids.pop(0)
+
+    start = len(prompt_ids)
+    end = len(sequence_ids)
+    stop_ids = {tokenizer.eos_token_id}
+    if getattr(tokenizer, "eot_token_id", None) is not None:
+        stop_ids.add(tokenizer.eot_token_id)
+
+    for idx in range(start, len(sequence_ids)):
+        if sequence_ids[idx] in stop_ids:
+            end = idx
+            break
+
+    return tokenizer.decode(sequence_ids[start:end], skip_special_tokens=True)
+
+model_id = "Pythagoras-LM/Pythagoras-Prover-Diffusion-4B"
+
+formal_statement = """
+import Mathlib
+import Aesop
+
+set_option maxHeartbeats 0
+
+open BigOperators Real Nat Topology Rat
+
+/-- The volume of a cone is given by the formula $V = \frac{1}{3}Bh$, where $B$ is the area of the base and $h$ is the height. The area of the base of a cone is 30 square units, and its height is 6.5 units. What is the number of cubic units in its volume? Show that it is 65.-/
+theorem mathd_algebra_478 (b h v : ℝ) (h₀ : 0 < b ∧ 0 < h ∧ 0 < v) (h₁ : v = 1 / 3 * (b * h))
+(h₂ : b = 30) (h₃ : h = 13 / 2) : v = 65 := by
+  sorry
+""".strip()
+
+prompt = """
+Complete the following Lean 4 code:
+
+```lean4
+{}```
+
+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()
+
+if not torch.cuda.is_available():
+    raise RuntimeError("CUDA is required for practical inference with this 4B diffusion checkpoint.")
+
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+if tokenizer.eos_token != "<|im_end|>":
+    tokenizer.eos_token = "<|im_end|>"
+if tokenizer.pad_token_id is None:
+    tokenizer.pad_token = tokenizer.eos_token or tokenizer.bos_token
+
+model = AutoModelForMaskedLM.from_pretrained(
+    model_id,
+    dtype=torch.bfloat16,
+    device_map={"": 0},
+    attn_implementation="sdpa",
+).eval()
+
+chat = [{"role": "user", "content": prompt.format(formal_statement)}]
+input_ids = tokenizer.apply_chat_template(
+    chat,
+    tokenize=True,
+    add_generation_prompt=True,
+    enable_thinking=False,
+)
+
+outputs = generate_bd3lm(
+    model,
+    tokenizer,
+    input_ids,
+    max_new_tokens=8192,
+    steps=8192,
+    block_size=8,
+    temperature=1.0,
+)
+
+completion = trim_completion(tokenizer, outputs[0].tolist(), input_ids)
+print(completion)
+````
+
+# Cite
+```
+@article{leang2026pythagoras,
+  title={Pythagoras-Prover: Advancing Efficient Formal Proving via Augmented Lean Formalisation},
+  author={Leang, Joshua Ong Jun and Zhao, Zheng and Stoian, Mihaela Catalina and Xu, Qiyuan and Li, Haonan and Li, Wenda and Cohen, Shay B. and Giunchiglia, Eleonora},
+  journal={arXiv preprint arXiv:2606.12594},
+  year={2026}
+}
+```