update

README.md

@@ -306,8 +306,8 @@ After the tactic generator is trained, we combine it with best-first search to p
 
 For models without retrieval, run:
 ```bash
-python prover/evaluate.py --data-path data/leandojo_benchmark_4/random/ --ckpt_path PATH_TO_MODEL_CHECKPOINT --split test --num-workers 8 --num-gpus 1
-python prover/evaluate.py --data-path data/leandojo_benchmark_4/novel_premises/ --ckpt_path PATH_TO_MODEL_CHECKPOINT --split test --num-workers 8 --num-gpus 1
+python prover/evaluate.py --data-path data/leandojo_benchmark_4/random/ --ckpt_path PATH_TO_MODEL_CHECKPOINT --split test --num-workers 12 --num-gpus 1
+python prover/evaluate.py --data-path data/leandojo_benchmark_4/novel_premises/ --ckpt_path PATH_TO_MODEL_CHECKPOINT --split test --num-workers 12 --num-gpus 1
 ```
 
 For models with retrieval, first use the retriever to index the corpus (pre-computing the embeddings of all premises):

generator/confs/cli_lean4_novel_premises.yaml

@@ -39,9 +39,9 @@ model:
   length_penalty: 0.0
   ret_ckpt_path: null
   eval_num_retrieved: 100
-  eval_num_workers: 6  # Lower this number if you don't have 80GB GPU memory.
+  eval_num_workers: 5  # Lower this number if you don't have 80GB GPU memory.
   eval_num_gpus: 1
-  eval_num_theorems: 300  # Lower this number will make validation faster (but noiser).
+  eval_num_theorems: 250  # Lower this number will make validation faster (but noiser).
 
 data:
   data_path: data/leandojo_benchmark_4/novel_premises/

generator/confs/cli_lean4_random.yaml

@@ -39,9 +39,9 @@ model:
   length_penalty: 0.0
   ret_ckpt_path: null
   eval_num_retrieved: 100
-  eval_num_workers: 6  # Lower this number if you don't have 80GB GPU memory.
+  eval_num_workers: 5  # Lower this number if you don't have 80GB GPU memory.
   eval_num_gpus: 1
-  eval_num_theorems: 300  # Lower this number will make validation faster (but noiser).
+  eval_num_theorems: 250  # Lower this number will make validation faster (but noiser).
 
 data:
   data_path: data/leandojo_benchmark_4/random/

generator/model.py

@@ -1,5 +1,6 @@
 """Lightning module for the tactic generator."""
 
+import os
 import torch
 import openai
 import pickle
@@ -243,7 +244,7 @@ def on_validation_epoch_end(self) -> None:
 
         from prover.evaluate import evaluate  # Avoid circular import.
 
-        ckpt_path = f"{self.trainer.log_dir}/checkpoints/last.ckpt"
+        ckpt_path = f"{self.trainer.log_dir}/checkpoints/last-tmp.ckpt"
         self.trainer.save_checkpoint(ckpt_path)
         logger.info(f"Saved checkpoint to {ckpt_path}. Evaluating...")
         torch.cuda.empty_cache()
@@ -278,6 +279,9 @@ def on_validation_epoch_end(self) -> None:
         self.log("Pass@1_val", acc, on_step=False, on_epoch=True, sync_dist=True)
         logger.info(f"Pass@1: {acc}")
 
+        if os.path.exists(ckpt_path):
+            os.remove(ckpt_path)
+
     ##############
     # Prediction #
     ##############

retrieval/datamodule.py

@@ -232,16 +232,15 @@ def prepare_data(self) -> None:
         pass
 
     def setup(self, stage: Optional[str] = None) -> None:
-        if stage in (None, "fit"):
-            self.ds_train = RetrievalDataset(
-                [os.path.join(self.data_path, "train.json")],
-                self.corpus,
-                self.num_negatives,
-                self.num_in_file_negatives,
-                self.max_seq_len,
-                self.tokenizer,
-                is_train=True,
-            )
+        self.ds_train = RetrievalDataset(
+            [os.path.join(self.data_path, "train.json")],
+            self.corpus,
+            self.num_negatives,
+            self.num_in_file_negatives,
+            self.max_seq_len,
+            self.tokenizer,
+            is_train=True,
+        )
 
         if stage in (None, "fit", "validate"):
             self.ds_val = RetrievalDataset(