Update prediction.py

prediction.py

@@ -9,8 +9,16 @@ from rdkit.Chem import Descriptors
 from rdkit.Chem import AllChem
 from datetime import datetime
 from db import get_database  # This must be available in your repo
+import random
 
-# Load ChemBERTa model + tokenizer
+# ------------------------ Ensuring Deterministic Behavior ------------------------
+random.seed(42)
+np.random.seed(42)
+torch.manual_seed(42)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+
+# ------------------------ Load ChemBERTa Model + Tokenizer ------------------------
 @st.cache_resource
 def load_chemberta():
     tokenizer = AutoTokenizer.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
@@ -20,7 +28,7 @@ def load_chemberta():
 
 tokenizer, chemberta = load_chemberta()
 
-# Load scalers
+# ------------------------ Load Scalers ------------------------
 scalers = {
     "Tensile Strength": joblib.load("scaler_Tensile_strength_Mpa_.joblib"),
     "Ionization Energy": joblib.load("scaler_Ionization_Energy_eV_.joblib"),
@@ -30,7 +38,7 @@ scalers = {
     "Molecular Weight": joblib.load("scaler_Molecular_Weight_g_mol_.joblib")
 }
 
-# Model Definition
+# ------------------------ 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__()
@@ -39,7 +47,7 @@ class TransformerRegressor(nn.Module):
             d_model=embedding_dim,
             nhead=8,
             dim_feedforward=ff_dim,
-            dropout=0.1,
+            dropout=0.0,  # No dropout for consistency
             batch_first=True
         )
         self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
@@ -57,7 +65,7 @@ class TransformerRegressor(nn.Module):
         x = x.mean(dim=1)
         return self.regression_head(x)
 
-# Load model
+# ------------------------ Load Model ------------------------
 @st.cache_resource
 def load_model():
     model = TransformerRegressor()
@@ -67,7 +75,7 @@ def load_model():
 
 model = load_model()
 
-# Descriptor computation
+# ------------------------ Descriptors ------------------------
 def compute_descriptors(smiles: str):
     mol = Chem.MolFromSmiles(smiles)
     if mol is None:
@@ -87,22 +95,22 @@ def compute_descriptors(smiles: str):
     ]
     return np.array(descriptors, dtype=np.float32)
 
-# Fingerprint computation
+# ------------------------ Fingerprints ------------------------
 def get_morgan_fingerprint(smiles, radius=2, n_bits=1280):
     mol = Chem.MolFromSmiles(smiles)
     if mol is None:
         raise ValueError("Invalid SMILES string.")
     fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius, nBits=n_bits)
-    return np.array(fp, dtype=np.float32).reshape(1, -1)  # (1, 1280)
+    return np.array(fp, dtype=np.float32).reshape(1, -1)
 
-# Embedding function
+# ------------------------ Embedding ------------------------
 def get_chemberta_embedding(smiles: str):
     inputs = tokenizer(smiles, return_tensors="pt")
     with torch.no_grad():
         outputs = chemberta(**inputs)
-    return outputs.last_hidden_state[:, 0, :]  # CLS token (1, 768)
+    return outputs.last_hidden_state.mean(dim=1)  # Use average instead of CLS token
 
-# Save prediction to MongoDB
+# ------------------------ Save to DB ------------------------
 def save_to_db(smiles, predictions):
     predictions_clean = {k: float(v) for k, v in predictions.items()}
     doc = {
@@ -113,7 +121,7 @@ def save_to_db(smiles, predictions):
     db = get_database()
     db["polymer_predictions"].insert_one(doc)
 
-# Main Streamlit UI + prediction
+# ------------------------ Streamlit App ------------------------
 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)
@@ -128,22 +136,22 @@ def show():
                 return
 
             descriptors = compute_descriptors(smiles_input)
-            descriptors_tensor = torch.tensor(descriptors, dtype=torch.float32).unsqueeze(0)  # (1, 10)
+            descriptors_tensor = torch.tensor(descriptors, dtype=torch.float32).unsqueeze(0)
 
-            fingerprint = get_morgan_fingerprint(smiles_input)  # (1, 1280)
-            fingerprint_tensor = torch.tensor(fingerprint, dtype=torch.float32)  # (1, 1280)
+            fingerprint = get_morgan_fingerprint(smiles_input)
+            fingerprint_tensor = torch.tensor(fingerprint, dtype=torch.float32)
 
-            embedding = get_chemberta_embedding(smiles_input)  # (1, 768)
+            embedding = get_chemberta_embedding(smiles_input)
 
-            combined_input = torch.cat([embedding, descriptors_tensor, fingerprint_tensor], dim=1)  # (1, 2058)
-            combined = combined_input.unsqueeze(1)  # (1, 1, 2058)
+            combined_input = torch.cat([embedding, descriptors_tensor, fingerprint_tensor], dim=1)
+            combined = combined_input.unsqueeze(1)
 
             with torch.no_grad():
                 preds = model(combined)
 
             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)
@@ -151,12 +159,10 @@ def show():
 
             results = {key: round(val, 4) for key, val in zip(keys, preds_rescaled.flatten())}
 
-            # Display results
             st.success("Predicted Properties:")
             for key, val in results.items():
                 st.markdown(f"**{key}**: {val}")
 
-            # Save to MongoDB
             save_to_db(smiles_input, results)
 
         except Exception as e: