Instantiates an EquivariantStructureDenoisingModule model for protein-ligand complex structure denoising. :param protein_model_cfg: Protein model configuration. :param score_cfg: Score configuration. :param task_cfg: Task configuration. :param state_dict: Optional (potentially-p
(
protein_model_cfg: DictConfig,
score_cfg: DictConfig,
task_cfg: DictConfig,
state_dict: Optional[STATE_DICT] = None,
)
| 706 | |
| 707 | |
| 708 | def resolve_score_head( |
| 709 | protein_model_cfg: DictConfig, |
| 710 | score_cfg: DictConfig, |
| 711 | task_cfg: DictConfig, |
| 712 | state_dict: Optional[STATE_DICT] = None, |
| 713 | ) -> torch.nn.Module: |
| 714 | """Instantiates an EquivariantStructureDenoisingModule model for protein-ligand complex |
| 715 | structure denoising. |
| 716 | |
| 717 | :param protein_model_cfg: Protein model configuration. |
| 718 | :param score_cfg: Score configuration. |
| 719 | :param task_cfg: Task configuration. |
| 720 | :param state_dict: Optional (potentially-pretrained) state dictionary. |
| 721 | :return: EquivariantStructureDenoisingModule model. |
| 722 | """ |
| 723 | model = EquivariantStructureDenoisingModule( |
| 724 | score_cfg.fiber_dim, |
| 725 | input_dim=protein_model_cfg.residue_dim, |
| 726 | input_pair_dim=protein_model_cfg.pair_dim, |
| 727 | hidden_dim=score_cfg.hidden_dim, |
| 728 | n_stacks=score_cfg.n_stacks, |
| 729 | n_heads=protein_model_cfg.n_heads, |
| 730 | dropout=task_cfg.dropout, |
| 731 | ) |
| 732 | if score_cfg.from_pretrained and state_dict is not None: |
| 733 | try: |
| 734 | model.load_state_dict( |
| 735 | { |
| 736 | ".".join(k.split(".")[1:]): v |
| 737 | for k, v in state_dict.items() |
| 738 | if k.startswith("score_head") |
| 739 | } |
| 740 | ) |
| 741 | log.info("Successfully loaded pretrained score weights.") |
| 742 | except Exception as e: |
| 743 | log.warning(f"Skipping loading of pretrained score weights due to: {e}.") |
| 744 | return model |
| 745 | |
| 746 | |
| 747 | def resolve_confidence_head( |
no test coverage detected