export fns for yaml config -> Alphafold3 or Trainer

Author: lucidrains

alphafold3_pytorch/__init__.py

@@ -40,7 +40,9 @@
 
 from alphafold3_pytorch.configs import (
     Alphafold3Config,
-    TrainerConfig
+    TrainerConfig,
+    create_alphafold3_from_yaml,
+    create_trainer_from_yaml
 )
 
 __all__ = [
@@ -74,5 +76,7 @@
     Alphafold3Config,
     AtomInput,
     Trainer,
-    TrainerConfig
+    TrainerConfig,
+    create_alphafold3_from_yaml,
+    create_trainer_from_yaml
 ]

alphafold3_pytorch/configs.py

@@ -142,3 +142,8 @@ def create_instance_from_yaml_file(
 
         trainer_config = TrainerConfig.from_yaml_file(path)
         return trainer_config.create_instance(**kwargs)
+
+# convenience functions
+
+create_alphafold3_from_yaml = Alphafold3Config.create_instance_from_yaml_file
+create_trainer_from_yaml = TrainerConfig.create_instance_from_yaml_file

pyproject.toml

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

tests/test_af3.py

@@ -22,8 +22,12 @@
     InputFeatureEmbedder,
     ConfidenceHead,
     DistogramHead,
-    Alphafold3,
-    Alphafold3Config
+    Alphafold3
+)
+
+from alphafold3_pytorch.configs import (
+    Alphafold3Config,
+    create_alphafold3_from_yaml
 )
 
 from alphafold3_pytorch.alphafold3 import (
@@ -574,5 +578,5 @@ def test_alphafold3_config():
     curr_dir = Path(__file__).parents[0]
     af3_yaml = curr_dir / 'alphafold3.yaml'
 
-    alphafold3 = Alphafold3Config.create_instance_from_yaml_file(af3_yaml)
+    alphafold3 = create_alphafold3_from_yaml(af3_yaml)
     assert isinstance(alphafold3, Alphafold3)

tests/test_trainer.py

@@ -13,7 +13,8 @@
     AtomInput,
     DataLoader,
     Trainer,
-    TrainerConfig
+    TrainerConfig,
+    create_trainer_from_yaml
 )
 
 # mock dataset
@@ -173,7 +174,7 @@ def test_trainer_config():
     curr_dir = Path(__file__).parents[0]
     trainer_yaml_path = curr_dir / 'trainer.yaml'
 
-    trainer = TrainerConfig.create_instance_from_yaml_file(
+    trainer = create_trainer_from_yaml(
         trainer_yaml_path,
         dataset = MockAtomDataset(16)
     )