@@ -346,50 +346,57 @@ def _log_iteration(
346346 def _save_checkpoint (self , iteration : int ) -> None :
347347 """
348348 Save a checkpoint
349-
349+
350350 Args:
351351 iteration: Current iteration number
352352 """
353353 checkpoint_dir = os .path .join (self .output_dir , "checkpoints" )
354354 os .makedirs (checkpoint_dir , exist_ok = True )
355-
355+
356356 # Create specific checkpoint directory
357357 checkpoint_path = os .path .join (checkpoint_dir , f"checkpoint_{ iteration } " )
358358 os .makedirs (checkpoint_path , exist_ok = True )
359-
359+
360360 # Save the database
361361 self .database .save (checkpoint_path , iteration )
362-
362+
363363 # Save the best program found so far
364364 best_program = None
365365 if self .database .best_program_id :
366366 best_program = self .database .get (self .database .best_program_id )
367367 else :
368368 best_program = self .database .get_best_program ()
369-
369+
370370 if best_program :
371371 # Save the best program at this checkpoint
372372 best_program_path = os .path .join (checkpoint_path , f"best_program{ self .file_extension } " )
373373 with open (best_program_path , "w" ) as f :
374374 f .write (best_program .code )
375-
375+
376376 # Save metrics
377377 best_program_info_path = os .path .join (checkpoint_path , "best_program_info.json" )
378378 with open (best_program_info_path , "w" ) as f :
379379 import json
380- json .dump ({
381- "id" : best_program .id ,
382- "generation" : best_program .generation ,
383- "iteration" : iteration ,
384- "metrics" : best_program .metrics ,
385- "language" : best_program .language ,
386- "timestamp" : best_program .timestamp ,
387- "saved_at" : time .time ()
388- }, f , indent = 2 )
389-
390- logger .info (f"Saved best program at checkpoint { iteration } with metrics: "
391- f"{ ', ' .join (f'{ name } ={ value :.4f} ' for name , value in best_program .metrics .items ())} " )
392-
380+
381+ json .dump (
382+ {
383+ "id" : best_program .id ,
384+ "generation" : best_program .generation ,
385+ "iteration" : iteration ,
386+ "metrics" : best_program .metrics ,
387+ "language" : best_program .language ,
388+ "timestamp" : best_program .timestamp ,
389+ "saved_at" : time .time (),
390+ },
391+ f ,
392+ indent = 2 ,
393+ )
394+
395+ logger .info (
396+ f"Saved best program at checkpoint { iteration } with metrics: "
397+ f"{ ', ' .join (f'{ name } ={ value :.4f} ' for name , value in best_program .metrics .items ())} "
398+ )
399+
393400 logger .info (f"Saved checkpoint at iteration { iteration } to { checkpoint_path } " )
394401
395402 def _save_best_program (self , program : Optional [Program ] = None ) -> None :
0 commit comments