able to instantiate a config from a deeply nested key in the config yml

Author: lucidrains

alphafold3_pytorch/configs.py

@@ -23,10 +23,27 @@
 def exists(v):
     return v is not None
 
+@typecheck
+def safe_deep_get(
+    d: dict,
+    dotpath: str | List[str],  # dotpath notation, so accessing {'a': {'b'': {'c': 1}}} would be "a.b.c"
+    default = None
+):
+    if isinstance(dotpath ,str):
+        dotpath = dotpath.split('.')
+
+    for key in dotpath:
+        if key not in d:
+            return default
+
+        d = d[key]
+
+    return d
+
 @typecheck
 def yaml_config_path_to_dict(
     path: str | Path
-) -> dict | None:
+) -> dict:
 
     if isinstance(path, str):
         path = Path(path)
@@ -73,16 +90,26 @@ class Alphafold3Config(BaseModelWithExtra):
 
     @staticmethod
     @typecheck
-    def from_yaml_file(path: str | Path):
+    def from_yaml_file(
+        path: str | Path,
+        dotpath: str | List[str] = []
+    ):
         config_dict = yaml_config_path_to_dict(path)
+        config_dict = safe_deep_get(config_dict, dotpath)
+        assert exists(config_dict), f'config not found at path {".".join(dotpath)}'
+
         return Alphafold3Config(**config_dict)
 
     def create_instance(self) -> Alphafold3:
         alphafold3 = Alphafold3(**self.model_dump())
         return alphafold3
 
-    def create_instance_from_yaml_file(path: str | Path) -> Alphafold3:
-        af3_config = Alphafold3Config.from_yaml_file(path)
+    def create_instance_from_yaml_file(
+        path: str | Path,
+        dotpath: str | List[str] = []
+    ) -> Alphafold3:
+
+        af3_config = Alphafold3Config.from_yaml_file(path, dotpath)
         return af3_config.create_instance()
 
 class TrainerConfig(BaseModelWithExtra):
@@ -102,8 +129,14 @@ class TrainerConfig(BaseModelWithExtra):
 
     @staticmethod
     @typecheck
-    def from_yaml_file(path: str | Path):
+    def from_yaml_file(
+        path: str | Path,
+        dotpath: str | List[str] = []
+    ):
         config_dict = yaml_config_path_to_dict(path)
+        config_dict = safe_deep_get(config_dict, dotpath)
+        assert exists(config_dict), f'config not found at path {".".join(dotpath)}'
+
         return TrainerConfig(**config_dict)
 
     def create_instance(
@@ -137,10 +170,11 @@ def create_instance(
 
     def create_instance_from_yaml_file(
         path: str | Path,
+        dotpath: str | List[str] = [],
         **kwargs
     ) -> Trainer:
 
-        trainer_config = TrainerConfig.from_yaml_file(path)
+        trainer_config = TrainerConfig.from_yaml_file(path, dotpath)
         return trainer_config.create_instance(**kwargs)
 
 # convenience functions

tests/test_af3.py

@@ -577,6 +577,10 @@ def test_alphafold3_without_msa_and_templates():
 def test_alphafold3_config():
     curr_dir = Path(__file__).parents[0]
     af3_yaml = curr_dir / 'alphafold3.yaml'
+    trainer_yml = curr_dir / 'trainer.yaml'
 
     alphafold3 = create_alphafold3_from_yaml(af3_yaml)
     assert isinstance(alphafold3, Alphafold3)
+
+    alphafold3_from_trainer_yml = create_alphafold3_from_yaml(trainer_yml, 'model')
+    assert isinstance(alphafold3_from_trainer_yml, Alphafold3)