File size: 1,121 Bytes
0a9174e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | 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()
|