| 73 | return pred, target |
| 74 | |
| 75 | def inference(self, batch): |
| 76 | graph = batch["graph"] |
| 77 | |
| 78 | if self.graph_construction_model: |
| 79 | graph = self.graph_construction_model(graph) |
| 80 | input = graph.residue_feature.float() |
| 81 | |
| 82 | if isinstance(self.model.structure_model, gvp.SurfGVP): |
| 83 | output = self.model(graph, input, batch["surf_graph"]) |
| 84 | else: |
| 85 | output = self.model(graph, input) |
| 86 | node_feature = output["node_feature"] |
| 87 | pred = self.linear(node_feature) |
| 88 | |
| 89 | if self.plddt_threshold: |
| 90 | # for AlphaFold2-predicted structure pdb files, plddt is saved as b_factor |
| 91 | plddt_mask = batch["graph"].b_factor < self.plddt_threshold |
| 92 | pred[plddt_mask] = output["sequence_logits"][plddt_mask] |
| 93 | |
| 94 | return pred, graph.num_residues |