Method forward

flowdock/models/components/flowdock.py:1973–2025

Perform a forward pass through the model. :param batch: A batch dictionary. :param training: Whether the model is in training mode. :param iter_id: The current iteration ID. :param observed_block_contacts: Observed block contacts. :param contact_prediction: W

(
        self,
        batch: MODEL_BATCH,
        iter_id: Union[int, str] = 0,
        observed_block_contacts: Optional[torch.Tensor] = None,
        contact_prediction: bool = True,
        infer_geometry_prior: bool = False,
        score: bool = False,
        use_template: bool = False,
        **kwargs: Dict[str, Any],
    )

        return ret

    def forward(
        self,
        batch: MODEL_BATCH,
        iter_id: Union[int, str] = 0,
        observed_block_contacts: Optional[torch.Tensor] = None,
        contact_prediction: bool = True,
        infer_geometry_prior: bool = False,
        score: bool = False,
        use_template: bool = False,
        **kwargs: Dict[str, Any],
    ) -> Union[MODEL_BATCH, torch.Tensor]:
        """Perform a forward pass through the model.

        :param batch: A batch dictionary.
        :param training: Whether the model is in training mode.
        :param iter_id: The current iteration ID.
        :param observed_block_contacts: Observed block contacts.
        :param contact_prediction: Whether to predict contacts.
        :param infer_geometry_prior: Whether to predict using a geometry prior.
        :param score: Whether to predict a denoised complex structure.
        :param use_template: Whether to use a template protein structure.
        :param kwargs: Additional keyword arguments.
        :return: Batch dictionary with outputs or ligand binding affinity.
        """
        prepare_batch(batch)

        batch = self.run_encoder_stack(
            batch,
            use_template=use_template,
            use_plddt=self.global_cfg.use_plddt,
            **kwargs,
        )

        if contact_prediction:
            self.run_contact_map_stack(
                batch,
                iter_id,
                observed_block_contacts=observed_block_contacts,
                **kwargs,
            )

        if infer_geometry_prior:
            assert (
                batch["misc"]["protein_only"] is False
            ), "Only protein-ligand complexes are supported for a geometry prior."
            self.infer_geometry_prior(batch, **kwargs)

        if score:
            batch["outputs"]["denoised_prediction"] = self.run_score_head(
                batch, embedding_iter_id=iter_id, **kwargs
            )

        return batch


if __name__ == "__main__":