77import typing
88import uuid
99from itp_interface .tools .simple_lean4_sync_executor import SimpleLean4SyncExecutor
10- from itp_interface .tools .lean4_context_helper import Lean4ContextHelper
1110from itp_interface .tools .coq_training_data_generator import GenericTrainingDataGenerationTransform , TrainingDataGenerationType
12- from itp_interface .tools .training_data_format import MergableCollection , TrainingDataMetadataFormat , TheoremProvingTrainingDataCollection , TheoremProvingTrainingDataFormat
13- from itp_interface .tools .training_data import TrainingData
11+ from itp_interface .tools .training_data_format import MergableCollection , TrainingDataMetadataFormat , ExtractionDataCollection , TheoremProvingTrainingDataFormat
12+ from itp_interface .tools .training_data import TrainingData , DataLayoutFormat
1413
1514class Local4DataExtractionTransform (GenericTrainingDataGenerationTransform ):
1615 def __init__ (self ,
@@ -24,70 +23,45 @@ def __init__(self,
2423 self .max_search_results = max_search_results
2524 self .max_parallelism = max_parallelism
2625
27- def get_meta_object (self ) -> MergableCollection :
28- return TrainingDataMetadataFormat (training_data_buffer_size = self .buffer_size )
26+ def get_meta_object (self ) -> TrainingDataMetadataFormat :
27+ return TrainingDataMetadataFormat (
28+ training_data_buffer_size = self .buffer_size ,
29+ data_filename_prefix = "extraction_data_" ,
30+ lemma_ref_filename_prefix = "extraction_lemma_refs_" )
2931
3032 def get_data_collection_object (self ) -> MergableCollection :
31- return TheoremProvingTrainingDataCollection ()
33+ return ExtractionDataCollection ()
3234
3335 def load_meta_from_file (self , file_path ) -> MergableCollection :
3436 return TrainingDataMetadataFormat .load_from_file (file_path )
3537
3638 def load_data_from_file (self , file_path ) -> MergableCollection :
37- return TheoremProvingTrainingDataCollection .load_from_file (file_path , self .logger )
39+ return ExtractionDataCollection .load_from_file (file_path , self .logger )
3840
39- def __call__ (self , training_data : TrainingData , project_id : str , lean_executor : SimpleLean4SyncExecutor , print_coq_executor_callback : typing .Callable [[], SimpleLean4SyncExecutor ], theorems : typing .List [str ] = None , other_args : dict = {}) -> TrainingData :
40- print_lean_executor = print_coq_executor_callback ()
41- lean_context_helper = Lean4ContextHelper (print_lean_executor , self .depth , self .logger )
42- lean_context_helper .__enter__ ()
41+ def __call__ (self ,
42+ training_data : TrainingData ,
43+ project_id : str ,
44+ lean_executor : SimpleLean4SyncExecutor ,
45+ print_coq_executor_callback : typing .Callable [[], SimpleLean4SyncExecutor ],
46+ theorems : typing .List [str ] = None ,
47+ other_args : dict = {}) -> TrainingData :
4348 file_namespace = lean_executor .main_file .replace ('/' , '.' )
4449 self .logger .info (f"=========================Processing { file_namespace } )
4550 theorem_id = str (uuid .uuid4 ())
4651 theorems = set (theorems ) if theorems is not None else None
47- assert len ( theorems ) == 1 , "Only one theorem can be processed at a time"
52+ cnt = 0
4853 with lean_executor :
4954 lean_executor .set_run_exactly ()
50- lean_executor ._skip_to_theorem (theorems .pop ())
51- start_goals = lean_context_helper .get_focussed_goals_from_proof_context (lean_executor .proof_context )
52- ran_next = lean_executor .run_next ()
53- cmd_ran = lean_executor .current_stmt
54- try :
55- theorem_id = theorem_id + "/" + lean_executor .get_current_lemma_name ()
56- except :
57- pass
58- proof_id = theorem_id
59- while not lean_executor .execution_complete and ran_next :
60- if lean_executor .is_in_proof_mode ():
61- end_goals = lean_context_helper .get_focussed_goals_from_proof_context (lean_executor .proof_context )
62- else :
63- end_goals = []
64- if len (start_goals ) > 0 and \
65- (len (start_goals ) != len (end_goals ) or not all (s_g == e_g for s_g , e_g in zip (start_goals , end_goals ))):
66- tdf = TheoremProvingTrainingDataFormat (
67- proof_id = proof_id ,
68- all_useful_defns_theorems = [],
69- start_goals = start_goals ,
70- end_goals = end_goals ,
71- proof_steps = [cmd_ran ],
72- simplified_goals = [],
73- addition_state_info = {},
74- file_path = lean_executor .main_file ,
75- theorem_name = proof_id ,
76- project_id = project_id )
77- training_data .merge (tdf )
78- if lean_executor .is_in_proof_mode ():
79- start_goals = end_goals
80- ran_next = lean_executor .run_next ()
81- cmd_ran = lean_executor .current_stmt
82- else :
83- ran_next = False
84- training_data .meta .num_theorems += 1
55+ line_infos = lean_executor .extract_all_theorems_and_definitions ()
56+ for line_info in line_infos :
57+ if theorems is not None and line_info .name not in theorems :
58+ continue
59+ training_data .merge (line_info )
60+ cnt += 1
61+ training_data .meta .last_proof_id = theorem_id
8562 self .logger .info (f"===============Finished processing { file_namespace } )
86- self .logger .info (f"Total theorems processed in this transform: 1" )
87- try :
88- lean_context_helper .__exit__ (None , None , None )
89- except :
90- pass
63+ self .logger .info (f"Total declarations processed in this transform: { cnt } )
64+ return training_data
9165
9266
9367if __name__ == "__main__" :
@@ -110,13 +84,14 @@ def _print_lean_executor_callback():
11084 search_lean_exec = SimpleLean4SyncExecutor (main_file = file_name , project_root = project_dir )
11185 search_lean_exec .__enter__ ()
11286 return search_lean_exec
113- transform = Local4DataGenerationTransform (0 , buffer_size = 1000 )
87+ transform = Local4DataExtractionTransform (0 , buffer_size = 1000 )
11488 training_data = TrainingData (
11589 output_path ,
11690 "training_metadata.json" ,
11791 training_meta = transform .get_meta_object (),
118- logger = logger )
92+ logger = logger ,
93+ layout = DataLayoutFormat .DECLARATION_EXTRACTION )
11994 with SimpleLean4SyncExecutor (project_root = project_dir , main_file = file_name , use_human_readable_proof_context = True , suppress_error_log = True ) as coq_exec :
120- transform (training_data , project_id , coq_exec , _print_lean_executor_callback , theorems = ['{"namespace": "Lean4Proj1 ", "name": "test2"}'
95+ transform (training_data , project_id , coq_exec , _print_lean_executor_callback , theorems = ["test2" , "test ""test1" ])
12196 save_info = training_data .save ()
12297 logger .info (f"Saved training data to { save_info } )
0 commit comments