Update app.py

app.py

@@ -122,8 +122,6 @@ def finetune(base_model_path, peptide_length):   #, train_dataset, test_dataset)
 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)
-    print("125: original_input:", original_input)
-    print("126: original_input.size:", original_input.shape)
     length_of_binder = len(binder_seq)
     
     # Prepare a batch with each row having one masked token from the binder sequence
@@ -134,14 +132,11 @@ def compute_pseudo_perplexity(model, tokenizer, protein_seq, binder_seq):
     # 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("135: masked_inputs:",masked_inputs)
-    print("136: masked_inputs.shape:",masked_inputs.shape)
+
     # Get model predictions and calculate loss
     with torch.no_grad():
         outputs = model(masked_inputs, labels=labels)
         loss = outputs.loss
-        print("140: logits:", outputs.logits)
-        print("141: logits.size:", outputs.logits.shape)
         
     # Loss is already averaged by the model
     avg_loss = loss.item()
@@ -163,23 +158,17 @@ 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("164: inputs:",inputs)
-        #print("165: inputs.shape:",inputs.size)
+
         with torch.no_grad():
             logits = model(**inputs).logits
-            print("166: logits:", logits)
-            print("167: logits.size:", logits.shape)
+
         mask_token_indices = (inputs["input_ids"] == tokenizer.mask_token_id).nonzero(as_tuple=True)[1]
         logits_at_masks = logits[0, mask_token_indices]
         
         # Apply top-k sampling
         top_k_logits, top_k_indices = logits_at_masks.topk(top_k, dim=-1)
-        print("171:top_k_logits, top_k_indices:", top_k_logits, top_k_indices)
-        print("174: top_k_logits.shape:", top_k_logits.shape)
         probabilities = torch.nn.functional.softmax(top_k_logits, dim=-1)                                
         predicted_indices = Categorical(probabilities).sample()
-        print("174:predicted_indices:", predicted_indices)
-        print("178: predicted_indices.shape:", predicted_indices.shape)
         predicted_token_ids = top_k_indices.gather(-1, predicted_indices.unsqueeze(-1)).squeeze(-1)
         generated_binder = tokenizer.decode(predicted_token_ids, skip_special_tokens=True).replace(' ', '')
         # Compute PPL for the generated binder
@@ -293,7 +282,7 @@ with demo:
                     interactive = True,
                 )
                 with gr.Row():
-                    peptide_length=gr.Slider(minimum=10, maximum=100, step=1, label="Peptide Maximum Length", value=50)
+                    peptide_length=gr.Slider(minimum=10, maximum=100, step=1, label="Peptide Maximum Length", value=15)
                     num_pred_peptides=gr.Slider(minimum=1, maximum=10, step=1, label="Number of Predicted Peptides", value=4)
             with gr.Column(scale=5, variant="compact"): 
                     name = gr.Dropdown(