File size: 1,543 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 | #!/usr/bin/env python3
"""用 create_uma_finetune_dataset.py 生成的 data yaml 更新 demo 配置文件。"""
import argparse
from pathlib import Path
import yaml
def load_yaml(path: Path):
with open(path) as f:
return yaml.safe_load(f)
def save_yaml(path: Path, data):
with open(path, "w") as f:
yaml.dump(data, f, default_flow_style=False, sort_keys=False, allow_unicode=True)
def update_demo_config(demo_config_path: Path, generated_data_path: Path):
demo_cfg = load_yaml(demo_config_path)
gen_data = load_yaml(generated_data_path)
# 生成的 yaml 直接就是 data 部分的内容（没有顶层 data 键）
demo_cfg["data"] = gen_data
# 保持 train/val 路径使用 ONESCIENCE_DATASETS_DIR 变量（数据放在 UMA/data/ 下）
demo_cfg["data"]["train_dataset"]["splits"]["train"][
"src"
] = "${ONESCIENCE_DATASETS_DIR}/data/oc20_finetune/train"
demo_cfg["data"]["val_dataset"]["splits"]["val"][
"src"
] = "${ONESCIENCE_DATASETS_DIR}/data/oc20_finetune/val"
save_yaml(demo_config_path, demo_cfg)
print(f"[OK] 已更新 {demo_config_path}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--demo-config", type=Path, required=True)
parser.add_argument(
"--generated-data",
type=Path,
default=Path("data/oc20_finetune/data/uma_conserving_data_task_energy_force.yaml"),
)
args = parser.parse_args()
update_demo_config(args.demo_config, args.generated_data)
|