Update esm_scripts/extract.py

esm_scripts/extract.py

@@ -132,38 +132,31 @@ def run(args):
 
 
 def run_demo(protein_name, protein_seq, model, alphabet, include, 
-             repr_layers=-1, truncation_seq_length=1022, toks_per_batch=4096):
-
+             repr_layers=[-1], truncation_seq_length=1022, toks_per_batch=4096):
     dataset = FastaBatchedDataset([protein_name], [protein_seq])
     batches = dataset.get_batch_indices(toks_per_batch, extra_toks_per_seq=1)
     data_loader = torch.utils.data.DataLoader(
         dataset, collate_fn=alphabet.get_batch_converter(truncation_seq_length), batch_sampler=batches
     )
     print(f"Read sequences")
-
     # output_dir.mkdir(parents=True, exist_ok=True)
     return_contacts = "contacts" in include
-
     assert all(-(model.num_layers + 1) <= i <= model.num_layers for i in repr_layers)
     repr_layers = [(i + model.num_layers + 1) % (model.num_layers + 1) for i in repr_layers]
-
     with torch.no_grad():
         for batch_idx, (labels, strs, toks) in enumerate(data_loader):
             print(
                 f"Processing {batch_idx + 1} of {len(batches)} batches ({toks.size(0)} sequences)"
             )
-            if torch.cuda.is_available() and not nogpu:
+            if torch.cuda.is_available():
                 toks = toks.to(device="cuda", non_blocking=True)
-
             out = model(toks, repr_layers=repr_layers, return_contacts=return_contacts)
-
             logits = out["logits"].to(device="cpu")
             representations = {
                 layer: t.to(device="cpu") for layer, t in out["representations"].items()
             }
             if return_contacts:
                 contacts = out["contacts"].to(device="cpu")
-
             for i, label in enumerate(labels):
                 result = {"label": label}
                 truncate_len = min(truncation_seq_length, len(strs[i]))
@@ -185,7 +178,6 @@ def run_demo(protein_name, protein_seq, model, alphabet, include,
                     }
                 if return_contacts:
                     result["contacts"] = contacts[i, : truncate_len, : truncate_len].clone()
-
                 return result['representations'][36]