Update prediction.py

prediction.py

@@ -5,15 +5,18 @@ from transformers import AutoTokenizer, AutoModel
 from rdkit import Chem
 from rdkit.Chem import AllChem, Descriptors
 from torch import nn
-import pandas as pd
-import requests
-import datetime
+from datetime import datetime
 from db import get_database  # Assuming you have a file db.py with get_database function to connect to MongoDB
 
-# Model Setup
-tokenizer = AutoTokenizer.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
-chemberta = AutoModel.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
-chemberta.eval()
+# Load tokenizer and 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()
+    return tokenizer, model
+
+tokenizer, chemberta = load_chemberta()
 
 # Define your model architecture
 class TransformerRegressor(nn.Module):
@@ -35,10 +38,16 @@ class TransformerRegressor(nn.Module):
         aggregated = encoded.mean(dim=1)
         return self.regression_head(aggregated)
 
-# Load model
-model = TransformerRegressor()
-model.load_state_dict(torch.load("transformer_model.pt", map_location=torch.device('cpu')))
-model.eval()
+# Load your saved model
+@st.cache_resource
+def load_regression_model():
+    model = TransformerRegressor()
+    state_dict = torch.load("transformer_model.pt", map_location=torch.device("cpu"))
+    model.load_state_dict(state_dict)
+    model.eval()
+    return model
+
+model = load_regression_model()
 
 # Feature Functions
 descriptor_fns = [Descriptors.MolWt, Descriptors.MolLogP, Descriptors.TPSA,
@@ -69,15 +78,7 @@ def embed_smiles(smiles_list):
     outputs = chemberta(**inputs)
     return outputs.last_hidden_state[:, 0, :]
 
-# Function to validate SMILES string
-def is_valid_smiles(smiles):
-    """ Validate if the input is a valid SMILES string using RDKit """
-    mol = Chem.MolFromSmiles(smiles)
-    return mol is not None
-
 # Function to save prediction to MongoDB
-from datetime import datetime
-
 def save_to_db(smiles, predictions):
     # Convert all prediction values to native Python float
     predictions_clean = {k: float(v) for k, v in predictions.items()}
@@ -88,11 +89,10 @@ def save_to_db(smiles, predictions):
         "timestamp": datetime.now()
     }
 
-    db = get_database()
+    db = get_database()  # Connect to MongoDB
     collection = db["polymer_predictions"]
     collection.insert_one(doc)
 
-
 # Prediction Page UI
 def show():
     st.markdown("<h1 style='text-align: center; color: #4CAF50;'>🔬 Polymer Property Prediction</h1>", unsafe_allow_html=True)
@@ -130,4 +130,4 @@ def show():
 
             # Save the prediction to MongoDB
             save_to_db(smiles_input, predictions)
-            st.success("Prediction saved successfully!")
+            st.success("Prediction saved successfully!")