Update esm_scripts/extract.py

esm_scripts/extract.py

@@ -131,17 +131,8 @@ def run(args):
                 )
 
 
-def run_demo(protein_name, protein_seq, model_location, include, nogpu=False, 
+def run_demo(protein_name, protein_seq, model, alphabet, include, 
              repr_layers=-1, truncation_seq_length=1022, toks_per_batch=4096):
-    model, alphabet = pretrained.load_model_and_alphabet(model_location)
-    model.eval()
-    if isinstance(model, MSATransformer):
-        raise ValueError(
-            "This script currently does not handle models with MSA input (MSA Transformer)."
-        )
-    if torch.cuda.is_available() and not nogpu:
-        model = model.cuda()
-        print("Transferred model to GPU")
 
     dataset = FastaBatchedDataset([protein_name], [protein_seq])
     batches = dataset.get_batch_indices(toks_per_batch, extra_toks_per_seq=1)