exclude filepath and chains from alphafold3 forward

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -4832,8 +4832,6 @@ def forward(
         pde_labels: Int['b n n'] | Int['b m m'] | None = None,
         plddt_labels: Int['b n'] | Int['b m'] | None = None,
         resolved_labels: Int['b n'] | Int['b m'] | None = None,
-        chains: Int['b 2'] | None = None,
-        filepath: List[str] | None = None,
         return_loss_breakdown = False,
         return_loss: bool = None,
         return_present_sampled_atoms: bool = False,

alphafold3_pytorch/inputs.py

@@ -125,6 +125,9 @@
 def flatten(arr):
     return [el for sub_arr in arr for el in sub_arr]
 
+def without_keys(d: dict, exclude: set):
+    return {k: v for k, v in d.items() if k not in exclude}
+
 def pad_to_len(t, length, value = 0, dim = -1):
     assert dim < 0
     zeros = (0, 0) * (-dim - 1)
@@ -149,6 +152,9 @@ def inner(x, *args, **kwargs):
 
 # atom level, what Alphafold3 accepts
 
+UNCOLLATABLE_ATOM_INPUT_FIELDS = {'filepath'}
+ATOM_INPUT_EXCLUDE_MODEL_FIELDS = {'filepath', 'chains'}
+
 @typecheck
 @dataclass
 class AtomInput:
@@ -217,6 +223,9 @@ class BatchedAtomInput:
     def dict(self):
         return asdict(self)
 
+    def model_forward_dict(self):
+        return without_keys(self.dict(), ATOM_INPUT_EXCLUDE_MODEL_FIELDS)
+
 # functions for saving an AtomInput to disk or loading from disk to AtomInput
 
 @typecheck

alphafold3_pytorch/trainer.py

@@ -26,7 +26,8 @@
     BatchedAtomInput,
     Alphafold3Input,
     PDBInput,
-    maybe_transform_to_atom_inputs
+    maybe_transform_to_atom_inputs,
+    UNCOLLATABLE_ATOM_INPUT_FIELDS,
 )
 
 from alphafold3_pytorch.data import (
@@ -147,7 +148,7 @@ def collate_inputs_to_batched_atom_input(
 
     outputs = []
 
-    for group_index, grouped in enumerate(zip(*atom_inputs)):
+    for key, grouped in zip(keys, zip(*atom_inputs)):
         # if all None, just return None
 
         not_none_grouped = [*filter(exists, grouped)]
@@ -158,8 +159,8 @@ def collate_inputs_to_batched_atom_input(
 
         # collate list of input filepath strings
 
-        if keys[group_index] == "filepath":
-            outputs.append(not_none_grouped)
+        if key in UNCOLLATABLE_ATOM_INPUT_FIELDS:
+            outputs.append(grouped)
             continue
 
         # default to empty tensor for any Nones
@@ -640,7 +641,7 @@ def __call__(
                     # model forwards
 
                     loss, loss_breakdown = self.model(
-                        **inputs.dict(),
+                        **inputs.model_forward_dict(),
                         return_loss_breakdown = True
                     )
 
@@ -702,7 +703,7 @@ def __call__(
 
                     for valid_batch in self.valid_dataloader:
                         valid_loss, loss_breakdown = eval_model(
-                            **valid_batch.dict(),
+                            **valid_batch.model_forward_dict(),
                             return_loss_breakdown = True
                         )
 
@@ -742,7 +743,7 @@ def __call__(
 
                 for test_batch in self.test_dataloader:
                     test_loss, loss_breakdown = eval_model(
-                        **test_batch.dict(),
+                        **test_batch.model_forward_dict(),
                         return_loss_breakdown = True
                     )
 

pyproject.toml

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

tests/test_input.py

@@ -115,7 +115,7 @@ def test_alphafold3_input(directed_bonds):
         )
     )
 
-    alphafold3(**batched_atom_input.dict(), num_sample_steps = 1)
+    alphafold3(**batched_atom_input.model_forward_dict(), num_sample_steps = 1)
 
 def test_atompos_input():
 
@@ -168,7 +168,7 @@ def test_atompos_input():
         )
     )
 
-    loss = alphafold3(**batched_atom_input.dict())
+    loss = alphafold3(**batched_atom_input.model_forward_dict())
     loss.backward()
 
     # sampling
@@ -244,7 +244,7 @@ def test_pdbinput_input():
         ),
     )
 
-    loss = alphafold3(**batched_atom_input.dict())
+    loss = alphafold3(**batched_atom_input.model_forward_dict())
     loss.backward()
 
     # sampling