Update esm_scripts/extract.py

esm_scripts/extract.py

@@ -131,7 +131,7 @@ def run(args):
                 )
 
 
-def run_demo(model_location, fasta_file, output_dir, include, nogpu, 
+def run_demo(protein_name, protein_seq, model_location, include, nogpu, 
              repr_layers=-1, truncation_seq_length=1022, toks_per_batch=4096):
     model, alphabet = pretrained.load_model_and_alphabet(model_location)
     model.eval()
@@ -143,14 +143,14 @@ def run_demo(model_location, fasta_file, output_dir, include, nogpu,
         model = model.cuda()
         print("Transferred model to GPU")
 
-    dataset = FastaBatchedDataset.from_file(fasta_file)
+    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 {fasta_file} with {len(dataset)} sequences")
 
-    output_dir.mkdir(parents=True, exist_ok=True)
+    # 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)
@@ -194,6 +194,8 @@ def run_demo(model_location, fasta_file, output_dir, include, nogpu,
                     }
                 if return_contacts:
                     result["contacts"] = contacts[i, : truncate_len, : truncate_len].clone()
+
+                return result['representations'][36]
                 
 
 def main():