ProSST PDB quantizer fix: load cpu weights if gpu not available

Author: niklases

pypef/llm/inference.py

@@ -6,6 +6,7 @@
 
 import numpy as np
 
+from pypef.utils.helpers import get_device
 from pypef.llm.utils import get_batches
 from pypef.llm.esm_lora_tune import esm_setup, esm_tokenize_sequences
 from pypef.llm.prosst_lora_tune import prosst_setup, prosst_tokenize_sequences
@@ -43,6 +44,8 @@ def inference(
     """
     Inference of base models.
     """
+    if device is None:
+        device = get_device()
     if llm == 'esm':
         logger.info("Zero-shot LLM inference on test set using ESM1v...")
         llm_dict = esm_setup(sequences)

pypef/llm/prosst_structure/quantizer.py

@@ -508,8 +508,15 @@ def __init__(
         self.subgraph_depth = subgraph_depth
         self.subgraph_interval = subgraph_interval
         self.anchor_nodes = anchor_nodes
+        if device is None:
+            self.device = get_device()
+        else:
+            self.device = device
         if model_path is None:
-            self.model_path = str(Path(__file__).parent / "static" / "AE.pt")
+            if self.device is 'cpu':
+                self.model_path = str(Path(__file__).parent / "static" / "AE_CPU.pt")
+            else:
+                self.model_path = str(Path(__file__).parent / "static" / "AE.pt")
         else:
             self.model_path = model_path
 
@@ -522,11 +529,6 @@ def __init__(
             self.cluster_dir = cluster_dir
             self.cluster_model = cluster_model
 
-        if device is None:
-            self.device = get_device()
-        else:
-            self.device = device
-
         # Load model
         node_dim = (256, 32)
         edge_dim = (64, 2)