Default header handling costs about half the compiling formalizations (two-line fix in ensure_mathlib_import)

[EfrainGaray]

Hi! 👋 First of all, thank you for releasing the model, the dataset and the evaluation code all together 🙏 It is rare to get all three, and it is the only reason I could check this end to end instead of guessing.

I have been running MathForm-8B on my own machine (a single 16 GB card) and I think I found something that is costing you a lot of valid formalizations, with a fix that turned out to be two lines in a function you already wrote.

What I measured

MathForm-8B in Q8_0, first 25 problems of FormalMATH-Lite, 8 samples each, default formalizer prompt, everything compiled against Lean 4.19 with full Mathlib (6305 of 6305 modules built).

	samples compiling	pass@8
as-is	95/200 (47.5%)	21/25 (84%)
header replaced with import Mathlib	188/200 (94%)	25/25 (100%)

Same generations, same everything. Only the header changed 😅

Why it happens

98 of the 105 failures are imports of modules that do not exist in Mathlib:

module	times
Mathlib.Algebra.BigOperators.Basic	22
Mathlib.Data.Nat.Prime	17
Mathlib.Data.Nat.Pow	8
Mathlib.Data.Nat.Factorial	6
Mathlib.Analysis.SpecialFunctions.Trigonometric	6

They look like paths from older Mathlib versions, so I suspect the model picked them up from training data rather than getting confused.

I did want to be sure it was my report and not my setup, so I checked before blaming anything: the real modules it cites, like Mathlib.Data.Real.Sqrt and Mathlib.Order.Basic, are present and compiled fine. The failing ones have no source file at all.

After the header fix, the 12 remaining failures out of 200 are genuine: ten type errors, one syntax error. That part I would not touch, it is the model doing real work and occasionally missing 🙂

The fix

You already have ensure_mathlib_import() in evaluation/utils.py. It only adds the import when there is none, so a wrong import Mathlib.Data.Nat.Prime sails right through:

def ensure_mathlib_import(code: str) -> str:
    if not code or not code.strip():
        return code
    if re.search(r"^\s*import\s+", code, flags=re.M):
        return code
    return f"import Mathlib\n\n{code.lstrip()}"

Replacing whatever imports were generated keeps the current behaviour for import-less output and recovers the rest:

def ensure_mathlib_import(code: str) -> str:
    if not code or not code.strip():
        return code
    body = re.sub(r"^\s*import\s+\S+[ \t]*\n?", "", code, flags=re.M)
    return f"import Mathlib\n\n{body.lstrip()}"

I tested it against empty input, no imports, one bad import, several imports, an already-correct header, and a file with open lines in between. All behave as expected. I would be glad to open a PR if that helps 🚀

Two smaller things

stepfun in prompts.py already pins the header ("Your code should start with: import Mathlib"), but DEFAULT_INFER_PROMPT_TEMPLATE is formalizer, which does not. Switching the default gets most of the benefit at the prompt level, without touching any code.

And whenever you build the next version of FormalVerse, normalizing headers there would stop the model from learning module names that are no longer valid.

Happy to share the raw generations and the per-file compile results if they are useful to you 📎 And thanks again for open-sourcing the whole thing, it made all of this a pleasure to dig into.