Ensure no null BatchedAtomInputs are passed to the Trainer (#214)

* Update trainer.py

* Update inputs.py

* Update trainer.py

* Update trainer.py

* Update inputs.py

* Update inputs.py

* Update inputs.py

* Update utils.py

* Update pyproject.toml

Author: amorehead

alphafold3_pytorch/inputs.py

@@ -13,6 +13,7 @@
 from io import StringIO
 from itertools import groupby
 from pathlib import Path
+from retrying import retry
 from typing import Any, Callable, Dict, List, Literal, Set, Tuple, Type
 
 import einx
@@ -74,7 +75,7 @@
     remove_consecutive_duplicate,
 )
 from alphafold3_pytorch.tensor_typing import Bool, Float, Int, typecheck
-from alphafold3_pytorch.utils.utils import default, exists, first
+from alphafold3_pytorch.utils.utils import default, exists, first, not_exists
 
 # silence RDKit's warnings
 
@@ -3303,7 +3304,7 @@ def pdb_input_to_molecule_input(
 
 # datasets
 
-# PDB dataset that returns a PDBInput based on folder
+# PDB dataset that returns either a PDBInput or AtomInput based on folder
 
 
 class PDBDataset(Dataset):
@@ -3321,6 +3322,7 @@ def __init__(
         crop_size: int = 384,
         training: bool | None = None,  # extra training flag placed by Alex on PDBInput
         sample_only_pdb_ids: Set[str] | None = None,
+        return_atom_inputs: bool = False,
         **pdb_input_kwargs,
     ):
         if isinstance(folder, str):
@@ -3333,6 +3335,7 @@ def __init__(
         self.sample_type = sample_type
         self.training = training
         self.sample_only_pdb_ids = sample_only_pdb_ids
+        self.return_atom_inputs = return_atom_inputs
         self.pdb_input_kwargs = pdb_input_kwargs
 
         self.cropping_config = {
@@ -3369,8 +3372,8 @@ def __len__(self):
         """Return the number of PDB mmCIF files in the dataset."""
         return len(self.files)
 
-    def __getitem__(self, idx: int | str) -> PDBInput:
-        """Return a PDBInput object for the specified index."""
+    def get_item(self, idx: int | str) -> PDBInput | AtomInput | None:
+        """Return either a PDBInput or an AtomInput object for the specified index."""
         sampled_id = None
 
         if exists(self.sampler):
@@ -3412,15 +3415,24 @@ def __getitem__(self, idx: int | str) -> PDBInput:
         if self.training:
             cropping_config = self.cropping_config
 
-        pdb_input = PDBInput(
+        i = PDBInput(
             mmcif_filepath=str(mmcif_filepath),
             chains=(chain_id_1, chain_id_2),
             cropping_config=cropping_config,
             training=self.training,
             **self.pdb_input_kwargs,
         )
 
-        return pdb_input
+        if self.return_atom_inputs:
+            i = maybe_transform_to_atom_input(i)
+        
+        return i
+
+    def __getitem__(self, idx: int | str, max_attempts: int = 5) -> PDBInput | AtomInput:
+        """Return either a PDBInput or an AtomInput object for the specified index."""
+        retry_decorator = retry(retry_on_result=not_exists, stop_max_attempt_number=max_attempts)
+        i = retry_decorator(self.get_item)(idx)
+        return i
 
 
 # the config used for keeping track of all the disparate inputs and their transforms down to AtomInput
@@ -3461,7 +3473,7 @@ def maybe_transform_to_atom_input(i: Any, raise_exception: bool = False) -> Atom
 
     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"
+            f"Invalid input type {type(i)} being passed into Trainer that is not converted to AtomInput correctly"
         )
 
     try:

alphafold3_pytorch/trainer.py

@@ -111,8 +111,8 @@ def collate_inputs_to_batched_atom_input(
     inputs: List,
     int_pad_value = -1,
     atoms_per_window: int | None = None,
-    map_input_fn: Callable | None = None
-
+    map_input_fn: Callable | None = None,
+    transform_to_atom_inputs: bool = True,
 ) -> BatchedAtomInput:
 
     if exists(map_input_fn):
@@ -121,7 +121,25 @@ def collate_inputs_to_batched_atom_input(
     # go through all the inputs
     # and for any that is not AtomInput, try to transform it with the registered input type to corresponding registered function
 
-    atom_inputs = maybe_transform_to_atom_inputs(inputs)
+    if transform_to_atom_inputs:
+        atom_inputs = maybe_transform_to_atom_inputs(inputs)
+
+        if len(atom_inputs) < len(inputs):
+            # if some of the `inputs` could not be converted into `atom_inputs`,
+            # randomly select a subset of the `atom_inputs` to duplicate to match
+            # the expected number of `atom_inputs`
+            assert (
+                len(atom_inputs) > 0
+            ), "No `AtomInput` objects could be created for the current batch."
+            atom_inputs = random.choices(atom_inputs, k=len(inputs))  # nosec
+    else:
+        atom_inputs = inputs
+
+    assert all(isinstance(i, AtomInput) for i in atom_inputs), (
+        "All inputs must be of type `AtomInput`. "
+        "If you want to transform the inputs to `AtomInput`, "
+        "set `transform_to_atom_inputs=True`."
+    )
 
     # take care of windowing the atompair_inputs and atompair_ids if they are not windowed already
 
@@ -248,9 +266,14 @@ def DataLoader(
     *args,
     atoms_per_window: int | None = None,
     map_input_fn: Callable | None = None,
+    transform_to_atom_inputs: bool = True,
     **kwargs
 ):
-    collate_fn = partial(collate_inputs_to_batched_atom_input, atoms_per_window = atoms_per_window)
+    collate_fn = partial(
+        collate_inputs_to_batched_atom_input,
+        atoms_per_window = atoms_per_window,
+        transform_to_atom_inputs = transform_to_atom_inputs,
+    )
 
     if exists(map_input_fn):
         collate_fn = partial(collate_fn, map_input_fn = map_input_fn)

alphafold3_pytorch/utils/utils.py

@@ -1,15 +1,6 @@
 import numpy as np
 
-from typing import Any, List
-
-def first(arr: List) -> Any:
-    """
-    Returns first element of list
-
-    :param arr: the list
-    :return: the element
-    """
-    return arr[0]
+from typing import Any, Iterable, List
 
 
 def exists(val: Any) -> bool:
@@ -21,6 +12,15 @@ def exists(val: Any) -> bool:
     return val is not None
 
 
+def not_exists(val: Any) -> bool:
+    """Check if a value does not exist.
+
+    :param val: The value to check.
+    :return: `True` if the value does not exist, otherwise `False`.
+    """
+    return val is None
+
+
 def default(v: Any, d: Any) -> Any:
     """Return default value `d` if `v` does not exist (i.e., is `None`).
 
@@ -31,6 +31,15 @@ def default(v: Any, d: Any) -> Any:
     return v if exists(v) else d
 
 
+def first(arr: Iterable[Any]) -> Any:
+    """Return the first element of an iterable object such as a list.
+
+    :param arr: An iterable object.
+    :return: The first element of the iterable object.
+    """
+    return arr[0]
+
+
 def always(value):
     """Always return a value."""
 

pyproject.toml

@@ -47,6 +47,7 @@ dependencies = [
     "pydantic>=2.8.2",
     "pyyaml",
     "rdkit>=2023.9.6",
+    "retrying",
     "scikit-learn>=1.5.0",
     "sh>=2.0.7",
     "shortuuid",