Update models/peptide_classifiers.py

models/peptide_classifiers.py

@@ -509,7 +509,7 @@ class AffinityModel(nn.Module):
 
 class HemolysisModel:
     def __init__(self, device):
-        self.predictor = xgb.Booster(model_file='/scratch/pranamlab/tong/collection/classifiers/ckpt/wt_hemolysis.json')
+        self.predictor = xgb.Booster(model_file='../classifier_ckpt/wt_hemolysis.json')
         
         self.model = EsmModel.from_pretrained("facebook/esm2_t33_650M_UR50D").to(device)
         self.model.eval()
@@ -569,7 +569,7 @@ class MLPClassifier(nn.Module):
 
 class NonfoulingModel:
     def __init__(self, device):
-        ckpt = torch.load('/scratch/pranamlab/tong/collection/classifiers/ckpt/wt_nonfouling.pt', weights_only=False, map_location=device)
+        ckpt = torch.load('../classifier_ckpt/wt_nonfouling.pt', weights_only=False, map_location=device)
         best_params = ckpt["best_params"]
         self.predictor = MLPClassifier(in_dim=1280, hidden=int(best_params["hidden"]), dropout=float(best_params.get("dropout", 0.1)))
         self.predictor.load_state_dict(ckpt["state_dict"])
@@ -595,7 +595,7 @@ class NonfoulingModel:
 class SolubilityModel:
     def __init__(self, device):
         # change model path
-        self.predictor = xgb.Booster(model_file='/scratch/pranamlab/tong/checkpoints/MOG-DFM/classifier_ckpt/best_model_solubility.json')
+        self.predictor = xgb.Booster(model_file='../classifier_ckpt/best_model_solubility.json')
         
         self.model = EsmModel.from_pretrained("facebook/esm2_t33_650M_UR50D").to(device)
         self.model.eval()
@@ -675,21 +675,6 @@ class PeptideCNN(nn.Module):
             return features
         return self.predictor(features)  # Output shape: (B, 1)
 
-# class HalfLifeModel:
-#     def __init__(self, device):
-#         input_dim = 1280
-#         hidden_dims = [input_dim // 2, input_dim // 4]
-#         output_dim = input_dim // 8
-#         dropout_rate = 0.3
-#         self.model = PeptideCNN(input_dim, hidden_dims, output_dim, dropout_rate).to(device)
-#         self.model.load_state_dict(torch.load('/scratch/pranamlab/tong/checkpoints/MOG-DFM/classifier_ckpt/best_model_half_life.pth', map_location=device, weights_only=False))
-#         self.model.eval()
-
-#     def __call__(self, x):
-#         prediction = self.model(x, return_features=False)
-#         halflife = torch.clamp(prediction.squeeze(-1), max=2.0, min=0.0)
-#         return halflife / 2
-
 
 # -----------------------------
 # Model definition (must match training)
@@ -758,7 +743,7 @@ class HalfLifeModel:
     def __init__(
         self,
         device,
-        ckpt_path = "/scratch/pranamlab/tong/PeptiVerse/src/halflife/finetune_stability_transformer_log/final_model.pt",
+        ckpt_path = "../classifier_ckpt/wt_halflife.pt",
     ):
         self.device = device