one more tiny change needed for trainer orchestrator

Author: lucidrains

alphafold3_pytorch/configs.py

@@ -29,11 +29,14 @@ def safe_deep_get(
     dotpath: str | List[str],  # dotpath notation, so accessing {'a': {'b'': {'c': 1}}} would be "a.b.c"
     default = None
 ):
-    if isinstance(dotpath ,str):
+    if isinstance(dotpath, str):
         dotpath = dotpath.split('.')
 
     for key in dotpath:
-        if key not in d:
+        if (
+            not isinstance(d, dict) or \
+            key not in d
+        ):
             return default
 
         d = d[key]
@@ -48,7 +51,7 @@ def yaml_config_path_to_dict(
     if isinstance(path, str):
         path = Path(path)
 
-    assert path.is_file()
+    assert path.is_file(), f'cannot find {str(path)}'
 
     with open(str(path), 'r') as f:
         maybe_config_dict = yaml.safe_load(f)
@@ -113,7 +116,7 @@ def create_instance_from_yaml_file(
         return af3_config.create_instance()
 
 class TrainerConfig(BaseModelWithExtra):
-    model: Alphafold3Config
+    model: Alphafold3Config | None = None
     num_train_steps: int
     batch_size: int
     grad_accum_every: int
@@ -142,6 +145,7 @@ def from_yaml_file(
     def create_instance(
         self,
         dataset: Dataset,
+        model: Alphafold3 | None = None,
         fabric: Fabric | None = None,
         test_dataset: Dataset | None = None,
         optimizer: Optimizer | None = None,
@@ -152,7 +156,12 @@ def create_instance(
 
         trainer_kwargs = self.model_dump()
 
-        alphafold3 = self.model.create_instance()
+        assert exists(self.model) ^ exists(model), 'either model is available on the trainer config, or passed in when creating the instance, but not both or neither'
+
+        if exists(self.model):
+            alphafold3 = self.model.create_instance()
+        else:
+            alphafold3 = model
 
         trainer_kwargs.update(dict(
             model = alphafold3,

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.1.28"
+version = "0.1.29"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

tests/test_trainer.py

@@ -14,7 +14,8 @@
     DataLoader,
     Trainer,
     TrainerConfig,
-    create_trainer_from_yaml
+    create_trainer_from_yaml,
+    create_alphafold3_from_yaml
 )
 
 # mock dataset
@@ -188,3 +189,21 @@ def test_trainer_config():
     # take a single training step
 
     trainer()
+
+# test creating trainer without model, given when creating instance
+
+def test_trainer_config_without_model():
+    curr_dir = Path(__file__).parents[0]
+
+    af3_yaml_path = curr_dir / 'alphafold3.yaml'
+    trainer_yaml_path = curr_dir / 'trainer_without_model.yaml'
+
+    alphafold3 = create_alphafold3_from_yaml(af3_yaml_path)
+
+    trainer = create_trainer_from_yaml(
+        trainer_yaml_path,
+        model = alphafold3,
+        dataset = MockAtomDataset(16)
+    )
+
+    assert isinstance(trainer, Trainer)