Update predict.py

DeepRank · Sep 20, 2023 · 5345b09 · 5345b09
1 parent 54397be
commit 5345b09
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/deeprank_gnn/predict.py b/src/deeprank_gnn/predict.py
@@ -39,10 +39,13 @@
 # Constants
 # TODO: Make these configurable
 ESM_MODEL = "esm2_t33_650M_UR50D"
-GNN_ESM_MODEL = "paper_pretrained_models/scoring_of_docking_models/gnn_esm/treg_yfnat_b64_e20_lr0.001_foldall_esm.pth.tar"
+GNN_ESM_MODEL = (
+    Path(__file__).parent
+    / "paper_pretrained_models/scoring_of_docking_models/gnn_esm/treg_yfnat_b64_e20_lr0.001_foldall_esm.pth.tar"
+)
 TOKS_PER_BATCH = 4096
 REPR_LAYERS = [0, 32, 33]
-TRUNCATION_SEQ_LENGTH = 2000
+TRUNCATION_SEQ_LENGTH = 2500
 INCLUDE = ["mean", "per_tok"]
 NPROC = mp.cpu_count() - 1 if mp.cpu_count() > 1 else 1
 BATCH_SIZE = 64