Update prediction.py

prediction.py

@@ -24,9 +24,10 @@ def load_chemberta():
     tokenizer = AutoTokenizer.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
     model = AutoModel.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
     model.eval()
+    model.to(device)  # Send model to GPU if available
     return tokenizer, model
-
-tokenizer, chemberta = load_chemberta()
+    # Check if CUDA is available for GPU acceleration
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # ------------------------ Load Scalers ------------------------
 scalers = {
@@ -65,7 +66,6 @@ class TransformerRegressor(nn.Module):
         x = x.mean(dim=1)
         return self.regression_head(x)
 
-# ------------------------ Load Model ------------------------
 # ------------------------ Load Model ------------------------
 @st.cache_resource
 def load_model():
@@ -73,18 +73,15 @@ def load_model():
     model = TransformerRegressor()
 
     # Load the state_dict (weights) from the saved model file
-    state_dict = torch.load("transformer_model.bin", map_location=torch.device("cpu"))
+    state_dict = torch.load("transformer_model.bin", map_location=device)  # Ensure loading on the correct device
 
     # Load the state_dict into the model
     model.load_state_dict(state_dict)
 
     # Set the model to evaluation mode
     model.eval()
+    model.to(device)  # Send model to GPU if available
     return model
-
-# Load the model
-model = load_model()
-
 # ------------------------ Descriptors ------------------------
 def compute_descriptors(smiles: str):
     mol = Chem.MolFromSmiles(smiles)