first offer way to create Alphafold3 instance from config file

Author: lucidrains

alphafold3_pytorch/__init__.py

@@ -38,6 +38,10 @@
     AtomInput
 )
 
+from alphafold3_pytorch.configs import (
+    Alphafold3Config
+)
+
 __all__ = [
     Attention,
     Attend,

alphafold3_pytorch/alphafold3.py

@@ -2694,7 +2694,7 @@ def __init__(
         self,
         *,
         dim_single_inputs,
-        atompair_dist_bins: Float[' d'],
+        atompair_dist_bins: List[float],
         dim_single = 384,
         dim_pairwise = 128,
         num_plddt_bins = 50,
@@ -2705,6 +2705,8 @@ def __init__(
     ):
         super().__init__()
 
+        atompair_dist_bins = Tensor(atompair_dist_bins)
+
         self.register_buffer('atompair_dist_bins', atompair_dist_bins)
 
         num_dist_bins = atompair_dist_bins.shape[-1]
@@ -2828,7 +2830,7 @@ def __init__(
         dim_single = 384,
         dim_pairwise = 128,
         dim_token = 768,
-        distance_bins: Float[' dist_bins'] = torch.linspace(3, 20, 38),
+        distance_bins: List[float] = torch.linspace(3, 20, 38).tolist(),
         ignore_index = -1,
         num_dist_bins: int | None = None,
         num_plddt_bins = 50,
@@ -3018,6 +3020,8 @@ def __init__(
 
         # logit heads
 
+        distance_bins = Tensor(distance_bins)
+
         self.register_buffer('distance_bins', distance_bins)
         num_dist_bins = default(num_dist_bins, len(distance_bins))
 
@@ -3065,7 +3069,7 @@ def save(self, path: str | Path, overwrite = False):
         if isinstance(path, str):
             path = Path(path)
 
-        assert not path.is_dir() and (not path.exists() or overwrite)
+        assert path.is_file() and (not path.exists() or overwrite)
 
         path.parent.mkdir(exist_ok = True, parents = True)
 
@@ -3080,7 +3084,7 @@ def load(self, path: str | Path, strict = False):
         if isinstance(path, str):
             path = Path(path)
 
-        assert path.exists() and not path.is_dir()
+        assert path.exists() and path.is_file()
 
         package = torch.load(str(path), map_location = 'cpu')
 
@@ -3098,7 +3102,7 @@ def init_and_load(path: str | Path):
         if isinstance(path, str):
             path = Path(path)
 
-        assert path.exists() and not path.is_dir()
+        assert path.is_file()
 
         package = torch.load(str(path), map_location = 'cpu')
 

alphafold3_pytorch/configs.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+from alphafold3_pytorch.typing import typecheck
+from alphafold3_pytorch.alphafold3 import Alphafold3
+
+import yaml
+from pathlib import Path
+
+from pydantic import BaseModel
+
+# functions
+
+def exists(v):
+    return v is not None
+
+@typecheck
+def yaml_config_path_to_dict(
+    path: str | Path
+) -> dict | None:
+
+    if isinstance(path, str):
+        path = Path(path)
+
+    assert path.is_file()
+
+    with open(str(path), 'r') as f:
+        maybe_config_dict = yaml.safe_load(f)
+
+    assert exists(maybe_config_dict), f'unable to parse yaml config at {str(path)}'
+    assert isinstance(maybe_config_dict, dict), f'yaml config file is not a dictionary'
+
+    return maybe_config_dict
+
+# base pydantic classes for constructing alphafold3 and trainer from config files
+
+class BaseModelWithExtra(BaseModel):
+    class Config:
+        extra = 'allow'
+        use_enum_values = True
+
+class Alphafold3Config(BaseModelWithExtra):
+    dim_atom_inputs: int
+    dim_template_feats: int
+    dim_template_model: int
+    atoms_per_window: int
+    dim_atom: int
+    dim_atompair_inputs: int
+    dim_atompair: int
+    dim_input_embedder_token: int
+    dim_single: int
+    dim_pairwise: int
+    dim_token: int
+    ignore_index: int = -1
+    num_dist_bins: int | None
+    num_plddt_bins: int
+    num_pde_bins: int
+    num_pae_bins: int
+    sigma_data: int | float
+    diffusion_num_augmentations: int
+    loss_confidence_weight: int | float
+    loss_distogram_weight: int | float
+    loss_diffusion_weight: int | float
+
+    @staticmethod
+    @typecheck
+    def from_yaml_file(path: str | Path):
+        config_dict = yaml_config_path_to_dict(path)
+        return Alphafold3Config(**config_dict)
+
+    def create_instance(self) -> Alphafold3:
+        alphafold3 = Alphafold3(**self.dict())
+        return alphafold3
+
+    def create_instance_from_yaml_file(path: str | Path) -> Alphafold3:
+        af3_config = Alphafold3Config.from_yaml_file(path)
+        return af3_config.create_instance()

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.1.20"
+version = "0.1.21"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }
@@ -30,13 +30,12 @@ dependencies = [
     "ema-pytorch>=0.4.8",
     "environs",
     "frame-averaging-pytorch>=0.0.18",
-    "hydra-core",
     "jaxtyping>=0.2.28",
     "lightning>=2.2.5",
-    "omegaconf",
     "pandas>=1.5.3",
     "pdbeccdutils>=0.8.5",
     "pydantic>=2.7.2",
+    "pyyaml",
     "taylor-series-linear-attention>=0.1.9",
     "timeout_decorator>=0.5.0",
     'torch_geometric',