adapt for esmfold

tests/benchmark.py

@@ -8,7 +8,7 @@
 from transformers import EsmForMaskedLM, EsmTokenizer
 
 from faesm.esm import FAEsmForMaskedLM
-from tests.utils import generate_random_esm2_inputs
+# from tests.utils import generate_random_esm2_inputs
 
 # Set Seaborn theme and professional settings
 sns.set_theme(style="white")  # Remove grid by using "white"
@@ -27,6 +27,27 @@
     }
 )
 
+def generate_random_esm2_inputs(
+    tokenizer, batch_size=3, min_seq_length=5, max_seq_length=10, device="cuda"
+):
+    """Generate random ESM2 model inputs."""
+    random_lengths = torch.randint(
+        min_seq_length, max_seq_length + 1, (batch_size,), device=device
+    )
+    random_tokens = [
+        torch.randint(low=4, high=29, size=(length,), device=device).tolist()
+        for length in random_lengths
+    ]
+    sequences = ["".join(tokenizer.convert_ids_to_tokens(seq)) for seq in random_tokens]
+    esm_input = tokenizer.batch_encode_plus(
+        sequences,
+        add_special_tokens=True,
+        padding=True,
+        truncation=True,
+        return_tensors="pt",
+    )
+    esm_input = {k: v.to(device) for k, v in esm_input.items()}
+    return esm_input
 
 def benchmark_torch_memory(f, *args, **kwargs):
     torch.cuda.reset_peak_memory_stats()
@@ -51,7 +72,7 @@ def benchmark_inference_time(f, *args, **kwargs):
             "facebook/esm2_t33_650M_UR50D",
             8,
             torch.float16,
-            [100, 200, 300, 400, 500, 600, 700, 800, 1000],
+            [100, 200, 300, 400, 500, 600, 700],
             10,
         )
     ],

tests/benchmark_faesmfold_vs_esmfold.py

@@ -70,17 +70,15 @@ def get_faesmfold(device):
     [
         (
             torch.float16,
-            [100,200,300,400,500],
-            8,
+            [100,300,400],
+            10,
         )
     ],
 )
 def test_esmfold_vs_faesmfold_benchmark(dtype, max_seq_lengths, repeats):
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    esmfold = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1").to(device).eval()
-    esmfold.esm = esmfold.esm.to(dtype)
-    fa_esmfold = get_faesmfold(device)
+
     esm_memory_usage, fa_esm_memory_usage = [], []
     esm_inference_times, fa_esm_inference_times = [], []
 
@@ -91,6 +89,9 @@ def test_esmfold_vs_faesmfold_benchmark(dtype, max_seq_lengths, repeats):
         esm_time_fold, fa_esm_time_fold = [], []
 
         for _ in tqdm(range(repeats)):
+            esmfold = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1").to(device).eval()
+            esmfold.esm = esmfold.esm.to(dtype)
+            fa_esmfold = get_faesmfold(device)
 
             def esm_forward():
 
@@ -104,6 +105,7 @@ def esm_forward():
             def fa_esm_forward():
                 fa_esmfold.esm.half()
                 fa_esmfold.infer_pdb(inputs)
+
             fa_esmfold.to(device)
             fa_esm_memory_fold.append(benchmark_torch_memory(fa_esm_forward))
             fa_esm_time_fold.append(benchmark_inference_time(fa_esm_forward))