Update app.py

app.py

@@ -140,7 +140,8 @@ def train_function_no_sweeps(base_model_path, train_dataset, test_dataset):
         no_cuda=False,
         seed=8893,
         fp16=True,
-        report_to='wandb'
+        #report_to='wandb'
+        report_to=None
     )
 
     # Initialize Trainer
@@ -160,6 +161,8 @@ def train_function_no_sweeps(base_model_path, train_dataset, test_dataset):
     trainer.save_model(save_path)
     tokenizer.save_pretrained(save_path)
 
+    return save_path
+    
 # 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)
@@ -217,6 +220,9 @@ inputs = tokenizer(protein_sequence, return_tensors="pt", truncation=True, max_l
 with torch.no_grad():
     logits = loaded_model(**inputs).logits
 
+# train
+saved_path = train_function_no_sweeps(base_model_path,train_dataset, test_dataset)
+
 # Get predictions
 tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])  # Convert input ids back to tokens
 predictions = torch.argmax(logits, dim=2)
@@ -233,7 +239,7 @@ for token, prediction in zip(tokens, predictions[0].numpy()):
         print((token, id2label[prediction]))
         
 # debug result
-dubug_result = predictions  #class_weights
+dubug_result = saved_path  #predictions  #class_weights
 
 demo = gr.Blocks(title="DEMO FOR ESM2Bind")