Update prediction.py

prediction.py

@@ -5,10 +5,9 @@ import numpy as np
 import joblib
 from transformers import AutoTokenizer, AutoModel
 from rdkit import Chem
-from rdkit.Chem import Descriptors
-from rdkit.Chem import AllChem
+from rdkit.Chem import Descriptors, AllChem
 from datetime import datetime
-from db import get_database  # This must be available in your repo
+from db import get_database
 import random
 
 # ------------------------ Ensuring Deterministic Behavior ------------------------
@@ -17,38 +16,25 @@ np.random.seed(42)
 torch.manual_seed(42)
 torch.backends.cudnn.deterministic = True
 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 ------------------------
 @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
+    model.to(device)
     return tokenizer, model
 
-
-# ------------------------ Load Scalers ------------------------
-scalers = {
-    "Tensile Strength": joblib.load("scaler_Tensile_strength_Mpa_.joblib"),
-    "Ionization Energy": joblib.load("scaler_Ionization_Energy_eV_.joblib"),
-    "Electron Affinity": joblib.load("scaler_Electron_Affinity_eV_.joblib"),
-    "logP": joblib.load("scaler_LogP.joblib"),
-    "Refractive Index": joblib.load("scaler_Refractive_Index.joblib"),
-    "Molecular Weight": joblib.load("scaler_Molecular_Weight_g_mol_.joblib")
-}
-
-# ------------------------ Transformer Model ------------------------
+# ------------------------ Load Transformer Model ------------------------
 class TransformerRegressor(nn.Module):
     def __init__(self, input_dim=2058, embedding_dim=768, ff_dim=1024, num_layers=2, output_dim=6):
         super().__init__()
         self.feat_proj = nn.Linear(input_dim, embedding_dim)
         encoder_layer = nn.TransformerEncoderLayer(
-            d_model=embedding_dim,
-            nhead=8,
-            dim_feedforward=ff_dim,
-            dropout=0.0,  # No dropout for consistency
+            d_model=embedding_dim, nhead=8,
+            dim_feedforward=ff_dim, dropout=0.0,
             batch_first=True
         )
         self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
@@ -68,26 +54,33 @@ class TransformerRegressor(nn.Module):
 
 @st.cache_resource
 def load_model():
-    # Initialize the model architecture first
     model = TransformerRegressor()
-
-    # 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
+    state_dict = torch.load("transformer_model(1).bin", map_location=device)
     model.load_state_dict(state_dict)
-
-    # Set the model to evaluation mode
     model.eval()
-    model.to(device)  # Send model to GPU if available
+    model.to(device)
     return model
+
+# ✅ Load tokenizer/model globally
+tokenizer, chemberta = load_chemberta()
+model = load_model()
+
+# ------------------------ Load Scalers ------------------------
+scalers = {
+    "Tensile Strength": joblib.load("scaler_Tensile_strength_Mpa_.joblib"),
+    "Ionization Energy": joblib.load("scaler_Ionization_Energy_eV_.joblib"),
+    "Electron Affinity": joblib.load("scaler_Electron_Affinity_eV_.joblib"),
+    "logP": joblib.load("scaler_LogP.joblib"),
+    "Refractive Index": joblib.load("scaler_Refractive_Index.joblib"),
+    "Molecular Weight": joblib.load("scaler_Molecular_Weight_g_mol_.joblib")
+}
+
 # ------------------------ Descriptors ------------------------
 def compute_descriptors(smiles: str):
     mol = Chem.MolFromSmiles(smiles)
     if mol is None:
         raise ValueError("Invalid SMILES string.")
-    
-    descriptors = [
+    return np.array([
         Descriptors.MolWt(mol),
         Descriptors.MolLogP(mol),
         Descriptors.TPSA(mol),
@@ -98,8 +91,7 @@ def compute_descriptors(smiles: str):
         Descriptors.HeavyAtomCount(mol),
         Descriptors.RingCount(mol),
         Descriptors.MolMR(mol)
-    ]
-    return np.array(descriptors, dtype=np.float32)
+    ], dtype=np.float32)
 
 # ------------------------ Fingerprints ------------------------
 def get_morgan_fingerprint(smiles, radius=2, n_bits=1280):
@@ -111,23 +103,22 @@ def get_morgan_fingerprint(smiles, radius=2, n_bits=1280):
 
 # ------------------------ Embedding ------------------------
 def get_chemberta_embedding(smiles: str):
-    inputs = tokenizer(smiles, return_tensors="pt")
+    inputs = tokenizer(smiles, return_tensors="pt").to(device)
     with torch.no_grad():
         outputs = chemberta(**inputs)
-    return outputs.last_hidden_state.mean(dim=1)  # Use average instead of CLS token
+    return outputs.last_hidden_state.mean(dim=1).cpu()
 
 # ------------------------ Save to DB ------------------------
 def save_to_db(smiles, predictions):
-    predictions_clean = {k: float(v) for k, v in predictions.items()}
     doc = {
         "smiles": smiles,
-        "predictions": predictions_clean,
+        "predictions": {k: float(v) for k, v in predictions.items()},
         "timestamp": datetime.now()
     }
     db = get_database()
     db["polymer_predictions"].insert_one(doc)
 
-# ------------------------ Streamlit App ------------------------
+# ------------------------ Streamlit UI ------------------------
 def show():
     st.markdown("<h1 style='text-align: center; color: #4CAF50;'>🔬 Polymer Property Prediction</h1>", unsafe_allow_html=True)
     st.markdown("<hr style='border: 1px solid #ccc;'>", unsafe_allow_html=True)
@@ -150,17 +141,14 @@ def show():
             embedding = get_chemberta_embedding(smiles_input)
 
             combined_input = torch.cat([embedding, descriptors_tensor, fingerprint_tensor], dim=1)
-            combined = combined_input.unsqueeze(1)
+            combined = combined_input.unsqueeze(1).to(device)
 
             with torch.no_grad():
-                preds = model(combined)
+                preds = model(combined).cpu().numpy()
 
-            preds_np = preds.numpy()
             keys = list(scalers.keys())
-
             preds_rescaled = np.concatenate([
-                scalers[keys[i]].inverse_transform(preds_np[:, [i]])
-                for i in range(6)
+                scalers[key].inverse_transform(preds[:, [i]]) for i, key in enumerate(keys)
             ], axis=1)
 
             results = {key: round(val, 4) for key, val in zip(keys, preds_rescaled.flatten())}
@@ -172,4 +160,6 @@ def show():
             save_to_db(smiles_input, results)
 
         except Exception as e:
-            st.error(f"Prediction failed: {e}")
+            st.error(f"Prediction failed: {e}")
+
+