File size: 4,774 Bytes
bf314e8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | from __future__ import annotations
import argparse
import logging
import os
from pathlib import Path
import yaml
from onescience.utils.uma.scripts.create_finetune_dataset import (
compute_normalizer_and_linear_reference,
launch_processing,
)
from onescience.utils.uma.units.mlip_unit.api.inference import UMATask
logging.basicConfig(level=logging.INFO)
TEMPLATE_DIR = Path("configs")
DATA_YAML_DIR = Path("data")
REGRESSION_LABEL_TO_TASK_YAML = {
"e": DATA_YAML_DIR / Path("uma_conserving_data_task_energy.yaml"),
"ef": DATA_YAML_DIR / Path("uma_conserving_data_task_energy_force.yaml"),
"efs": DATA_YAML_DIR / Path("uma_conserving_data_task_energy_force_stress.yaml"),
}
UMA_SM_FINETUNE_YAML = Path("uma_sm_finetune_template.yaml")
def create_yaml(
train_path: str,
val_path: str,
force_rms: float,
linref_coeff: list,
output_dir: str,
dataset_name: str,
regression_tasks: str,
# base_model_name: str,
):
data_task_yaml = TEMPLATE_DIR / REGRESSION_LABEL_TO_TASK_YAML[regression_tasks]
with open(data_task_yaml) as file:
template = yaml.safe_load(file)
template["dataset_name"] = dataset_name
template["normalizer_rmsd"] = force_rms
template["elem_refs"] = linref_coeff
template["train_dataset"]["splits"]["train"]["src"] = train_path
template["val_dataset"]["splits"]["val"]["src"] = val_path
# add extra large vaccum box for molecules
# if dataset_name == str(UMATask.OMOL):
# template["train_dataset"]["a2g_args"]["molecule_cell_size"] = 1000.0
os.makedirs(output_dir / DATA_YAML_DIR, exist_ok=True)
with open(
output_dir / REGRESSION_LABEL_TO_TASK_YAML[regression_tasks], "w"
) as yaml_file:
yaml.dump(template, yaml_file, default_flow_style=False, sort_keys=False)
uma_finetune_yaml = TEMPLATE_DIR / UMA_SM_FINETUNE_YAML
with open(uma_finetune_yaml) as file:
template_ft = yaml.safe_load(file)
#template_ft["base_model_name"] = base_model_name
template_ft["defaults"][0]["data"] = REGRESSION_LABEL_TO_TASK_YAML[
regression_tasks
].stem
with open(output_dir / UMA_SM_FINETUNE_YAML, "w") as yaml_file:
yaml.dump(template_ft, yaml_file, default_flow_style=False, sort_keys=False)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--train-dir",
type=str,
required=True,
help="Directory of ASE atoms objects to convert for training.",
)
parser.add_argument(
"--val-dir",
type=str,
required=True,
help="Directory of ASE atoms objects to convert for validation.",
)
parser.add_argument(
"--uma-task",
type=str,
required=True,
choices=[t.value for t in UMATask],
help="choose a uma task to finetune",
)
parser.add_argument(
"--regression-tasks",
type=str,
choices=["e", "ef", "efs"],
required=True,
help="Choose to finetune based on regression task set (you must have the corresponding labels in your dataset), can be energy (e), energy+force (ef) or energy+force+stress(efs)",
)
#parser.add_argument(
# "--base-model",
# type=str,
# default="uma-s-1",
# help="Name of base uma model",
#)
parser.add_argument(
"--output-dir",
type=Path,
required=True,
help="Output directory to save required finetuning artifacts.",
)
parser.add_argument(
"--num-workers",
type=int,
default=8,
help="Number of parallel workers for processing files.",
)
args = parser.parse_args()
assert not Path(
args.output_dir
).exists(), f"{args.output_dir} can't already exist, please choose a different dir"
# Launch processing for training data
train_path = args.output_dir / "train"
launch_processing(args.train_dir, train_path, args.num_workers)
force_rms, linref_coeff = compute_normalizer_and_linear_reference(
train_path, args.num_workers
)
val_path = args.output_dir / "val"
launch_processing(args.val_dir, val_path, args.num_workers)
create_yaml(
train_path=str(train_path),
val_path=str(val_path),
force_rms=float(force_rms),
linref_coeff=linref_coeff,
output_dir=args.output_dir,
dataset_name=args.uma_task,
regression_tasks=args.regression_tasks,
# base_model_name=args.base_model,
)
logging.info(f"Generated dataset and data config yaml in {args.output_dir}")
logging.info(
f"To run finetuning, run fairchem -c {args.output_dir}/{UMA_SM_FINETUNE_YAML}"
)
|