Update app.py

app.py

@@ -125,6 +125,8 @@ def finetune(base_model_path, peptide_length):   #, train_dataset, test_dataset)
 def compute_pseudo_perplexity(model, tokenizer, protein_seq, binder_seq):
     start = time.time()
     sequence = protein_seq + binder_seq 
+    
+    print("129:model.device in compute_pseudo_perplexity",model.device)
     original_input = tokenizer.encode(sequence, return_tensors='pt').to(model.device)
     length_of_binder = len(binder_seq)
     
@@ -155,8 +157,14 @@ def compute_pseudo_perplexity(model, tokenizer, protein_seq, binder_seq):
 def compute_plddt_iptm(protein_seq, binder_seq):
     start = time.time()
     # always the ESMFold model
-    model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1")
+    #model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1")
     tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1")
+
+    # based on https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/protein_folding.ipynb 
+    model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=True)
+    model = model.cuda()
+    model.esm = model.esm.half()     # Uncomment to switch the stem to float16
+    torch.backends.cuda.matmul.allow_tf32 = True
     
     sequence = protein_seq + binder_seq
     inputs = tokenizer(sequence, return_tensors='pt', add_special_tokens=False)
@@ -195,7 +203,8 @@ def generate_peptide_for_single_sequence(model, tokenizer, protein_seq, peptide_
         masked_peptide = '<mask>' * peptide_length
         input_sequence = protein_seq + masked_peptide
         inputs = tokenizer(input_sequence, return_tensors="pt").to(model.device)
-
+        print("198:model.device in generate_:",model.device)
+        
         with torch.no_grad():
             logits = model(**inputs).logits