| 20 | |
| 21 | |
| 22 | class ProteinDataProcessor: |
| 23 | def __init__( |
| 24 | self, |
| 25 | device, |
| 26 | scale=16.0, |
| 27 | ref_scale=5.0, |
| 28 | multiplicity=1, |
| 29 | inference_multiplicity=1, |
| 30 | backend="torch", |
| 31 | ): |
| 32 | self.device = device |
| 33 | self.scale = scale |
| 34 | self.ref_scale = ref_scale |
| 35 | # if multiplicity > 1, effective batch size is multiplicity * batch_size |
| 36 | self.multiplicity = multiplicity |
| 37 | self.inference_multiplicity = inference_multiplicity |
| 38 | self.backend = backend |
| 39 | if self.backend == "mlx": |
| 40 | self.center_random_fn = mlx_center_random |
| 41 | elif self.backend == "torch": |
| 42 | self.center_random_fn = torch_center_random |
| 43 | else: |
| 44 | raise ValueError(f"Unsupported backend: {self.backend}. Choose 'torch' or 'mlx'.") |
| 45 | |
| 46 | def process_esm( |
| 47 | self, |
| 48 | batch, |
| 49 | esm_model=None, |
| 50 | esm_dict=None, |
| 51 | af2_to_esm=None, |
| 52 | inference=False, |
| 53 | ): |
| 54 | sequence = batch["aa_seq"] |
| 55 | B = len(sequence) |
| 56 | L = batch["res_type"].shape[1] |
| 57 | num_tokens = batch["cropped_num_tokens"] |
| 58 | |
| 59 | aatype, mask, residx, linker_mask, _ = batch_encode_sequences( |
| 60 | sequence, residue_index_offset=512, chain_linker="G" * 25, |
| 61 | ) |
| 62 | |
| 63 | aatype, mask, residx, linker_mask = map( |
| 64 | lambda x: x.to(self.device), (aatype, mask, residx, linker_mask) |
| 65 | ) |
| 66 | |
| 67 | if residx is None: |
| 68 | residx = torch.arange(L, device=self.device).expand_as(aatype) |
| 69 | |
| 70 | esmaa = af2_idx_to_esm_idx(aatype, mask, af2_to_esm) |
| 71 | |
| 72 | multiplicity = self.multiplicity if not inference else self.inference_multiplicity |
| 73 | |
| 74 | esm_s_, _ = compute_language_model_representations( |
| 75 | esmaa, esm_model, esm_dict, backend=self.backend |
| 76 | ) |
| 77 | |
| 78 | esm_s_ = esm_s_.detach() |
| 79 | mask, linker_mask = mask.detach().bool(), linker_mask.detach().bool() |
no outgoing calls
no test coverage detected