| import time | |
| import torch | |
| import numpy as np | |
| from inference import AlphaFoldEmbedInference, AlphaFoldEmbedTokenizer | |
| def benchmark(): | |
| model = AlphaFoldEmbedInference() | |
| tokenizer = AlphaFoldEmbedTokenizer() | |
| model.eval() | |
| # Dummy data | |
| seqs = ["".join(np.random.choice(list("ACDEFGHIKLMNPQRSTVWY"), size=150)) for _ in range(100)] | |
| # true plddt | |
| true_plddts = [np.random.uniform(50, 100, size=150) for _ in range(100)] | |
| start_time = time.time() | |
| maes = [] | |
| with torch.no_grad(): | |
| for seq, true_plddt in zip(seqs, true_plddts): | |
| tokens = tokenizer.tokenize(seq) | |
| input_ids = torch.tensor([tokens]) | |
| out = model(input_ids) | |
| pred_plddt = out["plddt"].squeeze(0).numpy() | |
| mae = np.mean(np.abs(pred_plddt - true_plddt)) | |
| maes.append(mae) | |
| end_time = time.time() | |
| throughput = len(seqs) / (end_time - start_time) | |
| avg_mae = np.mean(maes) | |
| print(f"pLDDT MAE: {avg_mae:.4f}") | |
| print(f"Sequence Processing Throughput: {throughput:.2f} seq/s") | |
| if __name__ == "__main__": | |
| benchmark() | |