"""Set up MLIP training commandline interface."""from__future__importannotationsfrompathlibimportPathfromtypingimportAnnotatedfromtyperimportOption,Typerimportyamlapp=Typer()
[docs]@app.command()deftrain(mlip_config:Annotated[Path,Option(help="Configuration file to pass to MLIP CLI.")],fine_tune:Annotated[bool,Option(help="Whether to fine-tune a foundational model.")]=False,log:Annotated[Path,Option(help="Path to save logs to.")]=Path("./janus_results/train-log.yml"),tracker:Annotated[bool,Option(help="Whether to save carbon emissions of calculation")]=True,summary:Annotated[Path,Option(help=("Path to save summary of inputs, start/end time, and carbon emissions.")),]=Path("./janus_results/train-summary.yml"),)->None:""" Run training for MLIP by passing a configuration file to the MLIP's CLI. Parameters ---------- mlip_config Configuration file to pass to MLIP CLI. fine_tune Whether to fine-tune a foundational model. Default is False. log Path to write logs to. Default is Path("train-log.yml"). tracker Whether to save carbon emissions of calculation in log file and summary. Default is True. summary Path to save summary of inputs, start/end time, and carbon emissions. Default is Path("train-summary.yml"). """fromjanus_core.cli.utilsimportcarbon_summary,end_summary,start_summaryfromjanus_core.training.trainimporttrainasrun_trainwithopen(mlip_config,encoding="utf8")asconfig_file:config=yaml.safe_load(config_file)iffine_tune:if"foundation_model"notinconfig:raiseValueError("Please include `foundation_model` in your configuration file")if(config["foundation_model"]notin("small","medium","large","small_off","medium_off","large_off")andnotPath(config["foundation_model"]).exists()):raiseValueError(""" Invalid foundational model. Valid options are: 'small', 'medium', 'large', 'small_off', 'medium_off', 'large_off', or a path to the model """)elif"foundation_model"inconfig:raiseValueError("Please include the `--fine-tune` option for fine-tuning")config={"mlip_config":mlip_config,"fine_tune":fine_tune,"log":log,"tracker":tracker,"summary":summary,}log_kwargs={"filemode":"w"}iflog:log_kwargs["filename"]=logoutput_files={"log":log.absolute()}# Save summary information before training beginsstart_summary(command="train",summary=summary,info={},config=config,output_files=output_files,)# Run trainingrun_train(mlip_config,attach_logger=True,log_kwargs=log_kwargs,track_carbon=tracker)# Save carbon summaryiftracker:carbon_summary(summary=summary,log=log)# Save time after training has finishedend_summary(summary)