simple function for extending input transform config

Author: lucidrains

alphafold3_pytorch/__init__.py

@@ -32,6 +32,10 @@
     Alphafold3
 )
 
+from alphafold3_pytorch.inputs import (
+    register_input_transform
+)
+
 from alphafold3_pytorch.trainer import (
     Trainer,
     DataLoader,

alphafold3_pytorch/inputs.py

@@ -1,4 +1,4 @@
-from typing import TypedDict, Literal
+from typing import Type, TypedDict, Literal, Callable
 
 from alphafold3_pytorch.typing import (
     typecheck,
@@ -83,3 +83,12 @@ def single_protein_input_and_single_nucleic_acid_to_atom_input(
     SingleProteinInput: single_protein_input_to_atom_input,
     SingleProteinSingleNucleicAcidInput: single_protein_input_and_single_nucleic_acid_to_atom_input
 }
+
+# function for extending the config
+
+@typecheck
+def register_input_transform(
+    input_type: Type,
+    fn: Callable[[TypedDict], AtomInput]
+):
+    INPUT_TO_ATOM_TRANSFORM[input_type] = fn

alphafold3_pytorch/trainer.py

@@ -93,7 +93,7 @@ def collate_af3_inputs(
         if not exists(maybe_to_atom_fn):
             raise TypeError(f'invalid input type {type(i)} being passed into Trainer that is not converted to AtomInput correctly')
 
-        atom_inputs = maybe_to_atom_fn(i)
+        atom_inputs.append(maybe_to_atom_fn(i))
 
     # separate input dictionary into keys and values
 

pyproject.toml

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