Update app.py

app.py

@@ -150,55 +150,18 @@ def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset)
         "weight_decay": 0.2,
         # Add other hyperparameters as needed
     }
-    # The base model you will train a LoRA on top of
-    #base_model_path = "facebook/esm2_t12_35M_UR50D"  
-    
-    # Define labels and model
-    #id2label = {0: "No binding site", 1: "Binding site"}
-    #label2id = {v: k for k, v in id2label.items()}
 
-    
     base_model = AutoModelForTokenClassification.from_pretrained(base_model_path, num_labels=len(id2label), id2label=id2label, label2id=label2id)
-
-    '''
-    # Load the data from pickle files (replace with your local paths)
-    with open("./datasets/train_sequences_chunked_by_family.pkl", "rb") as f:
-        train_sequences = pickle.load(f)
-    
-    with open("./datasets/test_sequences_chunked_by_family.pkl", "rb") as f:
-        test_sequences = pickle.load(f)
-    
-    with open("./datasets/train_labels_chunked_by_family.pkl", "rb") as f:
-        train_labels = pickle.load(f)
-    
-    with open("./datasets/test_labels_chunked_by_family.pkl", "rb") as f:
-        test_labels = pickle.load(f)
-    '''
     
     # Tokenization
     tokenizer = AutoTokenizer.from_pretrained(base_model_path) #("facebook/esm2_t12_35M_UR50D")
-    #max_sequence_length = 1000
     
     train_tokenized = tokenizer(train_sequences, padding=True, truncation=True, max_length=max_sequence_length, return_tensors="pt", is_split_into_words=False)
     test_tokenized = tokenizer(test_sequences, padding=True, truncation=True, max_length=max_sequence_length, return_tensors="pt", is_split_into_words=False)
     
-    # Directly truncate the entire list of labels
-    #train_labels = truncate_labels(train_labels, max_sequence_length)
-    #test_labels = truncate_labels(test_labels, max_sequence_length)
-    
     train_dataset = Dataset.from_dict({k: v for k, v in train_tokenized.items()}).add_column("labels", train_labels)
     test_dataset = Dataset.from_dict({k: v for k, v in test_tokenized.items()}).add_column("labels", test_labels)
     
-    '''
-    # Compute Class Weights
-    classes = [0, 1]  
-    flat_train_labels = [label for sublist in train_labels for label in sublist]
-    class_weights = compute_class_weight(class_weight='balanced', classes=classes, y=flat_train_labels)
-    accelerator = Accelerator()
-    class_weights = torch.tensor(class_weights, dtype=torch.float32).to(accelerator.device)  
-    print(" class_weights:", class_weights)
-    '''
-    
     # Convert the model into a PeftModel
     peft_config = LoraConfig(
         task_type=TaskType.TOKEN_CLS, 
@@ -217,7 +180,7 @@ def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset)
     test_dataset = accelerator.prepare(test_dataset)
 
     model_name_base = base_model_path.split("/")[1]
-    timestamp = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
+    timestamp = datetime.now().strftime('%Y-%m-%d_%H')
 
     # Training setup
     training_args = TrainingArguments(
@@ -262,9 +225,6 @@ def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset)
 
     # Train and Save Model
     trainer.train()
-    #save_path = os.path.join("lora_binding_sites", f"best_model_esm2_t12_35M_lora_{timestamp}")
-    #trainer.save_model(save_path)
-    #tokenizer.save_pretrained(save_path)
 
     return save_path
 
@@ -279,8 +239,8 @@ MODEL_OPTIONS = [
 ]  # models users can choose from
 
 PEFT_MODEL_OPTIONS = [
-    "AmelieSchreiber/esm2_t12_35M_lora_binding_sites_v2_cp3",
     "wangjin2000/esm2_t6_8M-lora-binding-sites_2024-07-02_09-26-54",
+    "AmelieSchreiber/esm2_t12_35M_lora_binding_sites_v2_cp3",
 ]  # finetuned models 
 
 
@@ -297,21 +257,12 @@ with open("./datasets/train_labels_chunked_by_family.pkl", "rb") as f:
 with open("./datasets/test_labels_chunked_by_family.pkl", "rb") as f:
     test_labels = pickle.load(f)
 
-## Tokenization
-#tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t12_35M_UR50D")
 max_sequence_length = 1000
 
-#train_tokenized = tokenizer(train_sequences, padding=True, truncation=True, max_length=max_sequence_length, return_tensors="pt", is_split_into_words=False)
-#test_tokenized = tokenizer(test_sequences, padding=True, truncation=True, max_length=max_sequence_length, return_tensors="pt", is_split_into_words=False)
-
 # Directly truncate the entire list of labels
 train_labels = truncate_labels(train_labels, max_sequence_length)
 test_labels = truncate_labels(test_labels, max_sequence_length)
 
-#train_dataset = Dataset.from_dict({k: v for k, v in train_tokenized.items()}).add_column("labels", train_labels)
-#test_dataset = Dataset.from_dict({k: v for k, v in test_tokenized.items()}).add_column("labels", test_labels)
-
-
 # Compute Class Weights
 classes = [0, 1]  
 flat_train_labels = [label for sublist in train_labels for label in sublist]
@@ -324,48 +275,6 @@ id2label = {0: "No binding site", 1: "Binding site"}
 label2id = {v: k for k, v in id2label.items()}
 
 '''
-# inference
-# Path to the saved LoRA model
-model_path = "AmelieSchreiber/esm2_t12_35M_lora_binding_sites_v2_cp3"
-# ESM2 base model
-base_model_path = "facebook/esm2_t12_35M_UR50D"
-
-# Load the model
-base_model = AutoModelForTokenClassification.from_pretrained(base_model_path)
-loaded_model = PeftModel.from_pretrained(base_model, model_path)
-
-# Ensure the model is in evaluation mode
-loaded_model.eval()
-
-# Protein sequence for inference
-protein_sequence = "MAVPETRPNHTIYINNLNEKIKKDELKKSLHAIFSRFGQILDILVSRSLKMRGQAFVIFKEVSSATNALRSMQGFPFYDKPMRIQYAKTDSDIIAKMKGT"  # Replace with your actual sequence
-
-# Tokenize the sequence
-inputs = tokenizer(protein_sequence, return_tensors="pt", truncation=True, max_length=1024, padding='max_length')
-
-# Run the model
-with torch.no_grad():
-    logits = loaded_model(**inputs).logits
-
-# Get predictions
-tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])  # Convert input ids back to tokens
-predictions = torch.argmax(logits, dim=2)
-
-
-# Define labels
-id2label = {
-    0: "No binding site",
-    1: "Binding site"
-}
-
-# Print the predicted labels for each token
-for token, prediction in zip(tokens, predictions[0].numpy()):
-    if token not in ['<pad>', '<cls>', '<eos>']:
-        print((token, id2label[prediction]))
-
-# train
-saved_path = train_function_no_sweeps(base_model_path,train_dataset, test_dataset)
-
 # debug result
 dubug_result = saved_path  #predictions  #class_weights
 '''
@@ -376,12 +285,9 @@ with demo:
     gr.Markdown("# DEMO FOR ESM2Bind")
     #gr.Textbox(dubug_result)
     
-    #gr.Markdown("## Finetune Pre-trained Model")
     with gr.Column():
         gr.Markdown("## Select a base model and a corresponding PEFT finetune model")
-        #gr.Markdown(
-        #    """ Pick a base model and press **Finetune Pre-trained Model!"""
-        #)
+
         with gr.Row():   
             with gr.Column(scale=5, variant="compact"):
                 base_model_name = gr.Dropdown(
@@ -462,6 +368,7 @@ with demo:
         inputs=[base_model_name,PEFT_model_name,input_seq], 
         outputs = [output_text],
     )
+    
     # "Finetune Pre-trained Model" actions
     finetune_button.click(
         fn = train_function_no_sweeps,