LeanHammer Training

This is the training script for a transformer-based premise selector, developed as part of LeanHammer.

Training data preparation

To extract data, use the hammer branch of ntp-toolkit. Please follow the instructions at the start of README.md.

After data extraction, training data will be in some ntp-toolkit/Examples/mathlib, in which there should exist the following files:

Examples/mathlib/Imports/*.jsonl                   # imports of each module
Examples/mathlib/Declarations/*.jsonl              # declarations in each module
Examples/mathlib/TrainingDataWithPremises/*.jsonl  # pairs of (proof state, set of premises) in each module
Examples/mathlib/WithImports/*.lean                # source Lean code for each module, modified with a `import Hammer` line inserted at the beginning (for benchmarking purposes)
Examples/mathlib/HammerBlacklist.jsonl             # blacklist of very basic logic theorems not included in training
Examples/mathlib/revision                          # revision (commit or tag) of Mathlib extracted

Training

Training is done in train.py; run python train.py -h for CLI arguments. See train.sh for an example training run.

Evaluation & benchmark

Exporting the model, premise corpus, and pre-computed embeddings

The premise selection deployment would need to access the model, Mathlib premise corpus, and pre-computed embeddings of these premises. We store these on Hugging Face.

Run python scripts/upload.py (you need access to l3lab/lean-premises; alternatively you may specify your own Hugging Face dataset repo in both scripts/upload.py and the premise selection server environment). This uploads the model weights to model_push_repo at Mathlib's revision (e.g. v4.18.0), and the premise corpus and its pre-computed embeddings for premises in Mathlib, Batteries, and Lean core to l3lab/lean-premises at Mathlib's revision (e.g. v4.18.0).

For deployment to a premise selection service, see lean-premise-server.

Retrieving premises for tactic benchmarking

For benchmarking the proof rate of LeanHammer, run the following data preparation scripts first:

Run python scripts/export_decls.py to generate retrieved_premises/(valid|test)_decls.json, which includes the valid (resp. test) set of Mathlib theorems. Each entry in the JSON file includes the theorem name, module name, and the ground-truth set of premises in the human-written proof. This file is only used for benchmarking.

After training, train.py also evaluates the model and statically retrieves premises for theorems in the evaluation sets to retrieved_premises/dot_(valid|test)-MODEL_NAME.json. These statically retrieved premises are solely used for purposes of making LeanHammer benchmarking faster, since the hammer tactic won't have to call the cloud API service for every theorem and every setting tested. During deployment, retrieval is dynamically done by calling the server.

(If for some reason you want to re-generate these retrieved premises, you may run scripts/retrieve_premises.py, but usually this script does not have to be run because it is part of train.py.) Analogously, scripts/leandojo_retrieve.py retrieves premises for a trained LeanDojo ReProver retriever.

LeanHammer benchmark

TODO: the benchmarking repository is undergoing refactor.

Set the appropriate paths in scripts/tactic_benchmark.sh (which requires a path to ntp-toolkit checked out at the tacticBenchmark branch), and run scripts/tactic_benchmark.sh. This tries the LeanHammer on the theorems in the evaluation sets, and saves results to results_DATE-(valid|test)/MODEL_NAME/*.json. Each JSON file is a run of a model in a specific setting. One can use python scripts/print_results.py [json files] to pretty-print the evaluation results, including a breakdown of failure modes and cumulative proof rates. Corresponding to the setting names in the paper, the names used in the json filenames are: aesop_with_premises = aesop, hammerCore_nosimp = auto, aesop_hammerCore_nosimp = aesop+auto, and aesop_hammerCore_nosimp_with_premises = full.

Overview of all files

models.py contains multiple dataclasses that represent Premises, PremiseSets, goal States, etc. These dataclasses are shared between training and the premise selection server code.

data.py contains a training/evaluation dataset class RetrievalDataset, and a data collator.

evaluator.py contains PremiseRetrievalEvaluator, a subclass of SentenceEvaluator from sentence-transformers, for evaluating metrics like recall@k.

loss.py contains MaskedCachedMultipleNegativesRankingLoss, a version of in-batch contrastive loss with a retrieval mask (see our paper).

train.py is the main training entrypoint.

Scripts under scripts/ are (currently) testing scripts for after training, including evaluation, LeanHammer benchmarking, exporting pre-computed embeddings, and Hugging Face upload.