Update prediction.py

prediction.py

@@ -20,15 +20,13 @@ torch.backends.cudnn.benchmark = False
 # Check if CUDA is available for GPU acceleration
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # ------------------------ Load ChemBERTa Model + Tokenizer ------------------------
-# # ------------------------ Load ChemBERTa Model + Tokenizer ------------------------
-# @st.cache_resource
-# 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)
-#     return tokenizer, model
-# tokenizer, chemberta = load_chemberta()
+@st.cache_resource
+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
 
 
 # ------------------------ Load Scalers ------------------------
@@ -67,27 +65,22 @@ class TransformerRegressor(nn.Module):
         x = self.transformer_encoder(x)
         x = x.mean(dim=1)
         return self.regression_head(x)
-@st.cache_resource
-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)
-    return tokenizer, model
 
 @st.cache_resource
 def load_model():
+    # Initialize the model architecture first
     model = TransformerRegressor()
-    state_dict = torch.load("transformer_model(1).bin", map_location=device)
+
+    # Load the state_dict (weights) from the saved model file
+    state_dict = torch.load("transformer_model(1).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)
+    model.to(device)  # Send model to GPU if available
     return model
-
-# Call them to load the actual models
-tokenizer, chemberta = load_chemberta()
-model = load_model()
-
 # ------------------------ Descriptors ------------------------
 def compute_descriptors(smiles: str):
     mol = Chem.MolFromSmiles(smiles)