Update alphafold3.py (#218)

Author: amorehead

alphafold3_pytorch/alphafold3.py

@@ -2806,6 +2806,7 @@ def forward(
         ligand_loss_weight = 10.,
         return_loss_breakdown = False,
         single_structure_input=False,
+        filepaths: List[str] | None = None,
     ) -> ElucidatedAtomDiffusionReturn:
 
         # diffusion loss
@@ -2875,7 +2876,7 @@ def forward(
                 )
             except Exception as e:
                 # NOTE: For many (random) unit test inputs, permutation alignment can be unstable
-                logger.warning(f"Skipping multi-chain permutation alignment due to: {e}")
+                logger.warning(f"Skipping multi-chain permutation alignment {f'for {filepaths}' if exists(filepaths) else ''} due to: {e}")
 
         # main diffusion mse loss
 
@@ -6366,7 +6367,8 @@ def forward(
         detach_when_recycling: bool = None,
         min_conf_resolution: float = 0.1,
         max_conf_resolution: float = 4.0,
-        hard_validate: bool = False
+        hard_validate: bool = False,
+        filepaths: List[str] | None = None
     ) -> (
         Float['b m 3'] |
         Float['ts b m 3'] |
@@ -6849,6 +6851,7 @@ def forward(
                 nucleotide_loss_weight = self.nucleotide_loss_weight,
                 ligand_loss_weight = self.ligand_loss_weight,
                 single_structure_input = single_structure_input,
+                filepaths = filepaths,
             )
 
         # confidence head
@@ -6920,7 +6923,7 @@ def forward(
                         )
                     except Exception as e:
                         # NOTE: For many (random) unit test inputs, permutation alignment can be unstable
-                        logger.warning(f"Skipping multi-chain permutation alignment due to: {e}")
+                        logger.warning(f"Skipping multi-chain permutation alignment {f'for {filepaths}' if exists(filepaths) else ''} due to: {e}")
 
                 assert exists(
                     distogram_atom_indices