Update prediction.py

prediction.py

@@ -20,15 +20,15 @@ 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()
+# # ------------------------ 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()
 
 
 # ------------------------ Load Scalers ------------------------
@@ -67,6 +67,13 @@ 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():
@@ -80,6 +87,7 @@ def load_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)