Update app.py

app.py

@@ -123,19 +123,15 @@ def compute_pseudo_perplexity(model, tokenizer, protein_seq, binder_seq):
     sequence = protein_seq + binder_seq 
     original_input = tokenizer.encode(sequence, return_tensors='pt').to(model.device)
     length_of_binder = len(binder_seq)
-    print("line 126, protein_seq",len(protein_seq))
-    print("line 127, length_of_binder",length_of_binder)
     
     # Prepare a batch with each row having one masked token from the binder sequence
     masked_inputs = original_input.repeat(length_of_binder, 1)
     positions_to_mask = torch.arange(-length_of_binder - 1, -1, device=model.device)
     masked_inputs[torch.arange(length_of_binder), positions_to_mask] = tokenizer.mask_token_id
-    print("line 131 : masked_inputs:", masked_inputs.shape)
     
     # Prepare labels for the masked tokens
     labels = torch.full_like(masked_inputs, -100)
     labels[torch.arange(length_of_binder), positions_to_mask] = original_input[0, positions_to_mask]
-    print("line 136 : labels:", labels.shape)
     
     # Get model predictions and calculate loss
     with torch.no_grad():
@@ -170,16 +166,10 @@ def generate_peptide_for_single_sequence(model, tokenizer, protein_seq, peptide_
         
         # Apply top-k sampling
         top_k_logits, top_k_indices = logits_at_masks.topk(top_k, dim=-1)
-        print("line 169 : top_k_logits:", top_k_logits)
-        print("line 170 : top_k_indices:", top_k_indices)
-        probabilities = torch.nn.functional.softmax(top_k_logits, dim=-1)
-        print("line 172 : probabilities:", probabilities)                                     
-        predicted_indices = Categorical(probabilities).sample()
-        print("line 174 : predicted_indices:", predicted_indices)           
+        probabilities = torch.nn.functional.softmax(top_k_logits, dim=-1)                                
+        predicted_indices = Categorical(probabilities).sample()       
         predicted_token_ids = top_k_indices.gather(-1, predicted_indices.unsqueeze(-1)).squeeze(-1)
-        print("line 176 : predicted_token_ids:", predicted_token_ids)
         generated_binder = tokenizer.decode(predicted_token_ids, skip_special_tokens=True).replace(' ', '')
-        print("line 178 : generated_binder:", generated_binder)
         # Compute PPL for the generated binder
         ppl_value = compute_pseudo_perplexity(model, tokenizer, protein_seq, generated_binder)
         
@@ -212,7 +202,7 @@ def predict_peptide(base_model_path, finetuned_model_path, input_seqs, peptide_l
         results_df = pd.DataFrame(results, columns=['Input Sequence', 'Binder', 'Pseudo Perplexity'])
     print(results_df)
 
-    #combine target protein and peptide with 20 G amino acids.
+    #combine target protein and predicted peptide with 20 G amino acids.
     separator = 'G' * 20
     peptide_lp = results_df['Binder'][results_df['Pseudo Perplexity'].idxmin()] #Choosing the one with the lowest perplexity
     print("lowest perplesity:", peptide_lp)