Update app.py

app.py

@@ -106,7 +106,7 @@ def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset)
     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)
@@ -119,22 +119,23 @@ def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset)
     
     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
+    #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_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]
@@ -142,6 +143,7 @@ def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset)
     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(
@@ -188,7 +190,7 @@ def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset)
         fp16=True,
         #report_to='wandb'
         report_to=None,
-        hub_token = access_token, #jw 20240701
+        hub_token = HF_TOKEN, #jw 20240701
     )
 
     # Initialize Trainer
@@ -211,7 +213,7 @@ def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset)
     return save_path
 
 # Constants & Globals
-access_token = os.getenv("hf_access_tokens")
+HF_TOKEN = os.environ.get("HF_TOKEN")
 
 MODEL_OPTIONS = [
     "facebook/esm2_t6_8M_UR50D",
@@ -233,19 +235,19 @@ 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")
+## 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)
+#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)
+#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