Update prediction.py

prediction.py

@@ -1,133 +1,88 @@
 import streamlit as st
 import torch
+import joblib
+import pandas as pd
 import numpy as np
-from transformers import AutoTokenizer, AutoModel
 from rdkit import Chem
-from rdkit.Chem import AllChem, Descriptors
-from torch import nn
-from datetime import datetime
-from db import get_database  # Assuming you have a file db.py with get_database function to connect to MongoDB
-
-# Load tokenizer and ChemBERTa model
-@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):
-    def __init__(self, emb_dim=768, feat_dim=2058, output_dim=6, nhead=8, num_layers=2):
-        super().__init__()
-        self.feat_proj = nn.Linear(feat_dim, emb_dim)
-        encoder_layer = nn.TransformerEncoderLayer(d_model=768, nhead=8, dim_feedforward=1024, dropout=0.1, batch_first=True)
-        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
-        self.regression_head = nn.Sequential(
-            nn.Linear(emb_dim, 256), nn.ReLU(),
-            nn.Linear(256, 128), nn.ReLU(),
-            nn.Linear(128, output_dim)
-        )
-
-    def forward(self, x, feat):
-        feat_emb = self.feat_proj(feat)
-        stacked = torch.stack([x, feat_emb], dim=1)
-        encoded = self.transformer_encoder(stacked)
-        aggregated = encoded.mean(dim=1)
-        return self.regression_head(aggregated)
-
-# 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,
-                  Descriptors.NumRotatableBonds, Descriptors.NumHAcceptors,
-                  Descriptors.NumHDonors, Descriptors.RingCount,
-                  Descriptors.FractionCSP3, Descriptors.HeavyAtomCount,
-                  Descriptors.NHOHCount]
-
-def fix_smiles(s):
-    try:
-        mol = Chem.MolFromSmiles(s.strip())
-        if mol:
-            return Chem.MolToSmiles(mol)
-    except:
-        return None
-    return None
+from rdkit.Chem import AllChem
+from transformers import AutoTokenizer, AutoModel
+import os
+
+# Load ChemBERTa model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
+chemberta_model = AutoModel.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+chemberta_model.to(device)
+chemberta_model.eval()
+
+# Load models
+model_dir = "saved_model"
+model_paths = [os.path.join(model_dir, f) for f in os.listdir(model_dir) if f.endswith(".pkl") and "scaler" not in f]
+models = [joblib.load(p) for p in model_paths]
+
+# Load input and target scalers
+input_scaler_path = os.path.join(model_dir, "scaler.pkl")
+input_scaler = joblib.load(input_scaler_path) if os.path.exists(input_scaler_path) else None
 
-def compute_features(smiles):
+target_scaler_path = os.path.join(model_dir, "target_scaler.pkl")
+target_scaler = joblib.load(target_scaler_path) if os.path.exists(target_scaler_path) else None
+
+# Properties
+PROPERTY_NAMES = ["Tensile Strength", "Ionization Energy", "Electron Affinity", "logP", "Refractive Index", "Molecular Weight"]
+
+def smiles_to_fingerprint(smiles, radius=2, nBits=2048):
     mol = Chem.MolFromSmiles(smiles)
-    if not mol:
-        return [0]*10 + [0]*2048
-    desc = [fn(mol) for fn in descriptor_fns]
-    fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=2048)
-    return desc + list(fp)
-
-def embed_smiles(smiles_list):
-    inputs = tokenizer(smiles_list, return_tensors="pt", padding=True, truncation=True, max_length=128)
-    outputs = chemberta(**inputs)
-    return outputs.last_hidden_state[:, 0, :]
-
-# Function to save prediction to MongoDB
-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()}
-    
-    doc = {
-        "smiles": smiles,
-        "predictions": predictions_clean,
-        "timestamp": datetime.now()
-    }
-
-    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)
-    st.markdown("<hr style='border: 1px solid #ccc;'>", unsafe_allow_html=True)
-
-    smiles_input = st.text_input("Enter SMILES Representation of Polymer")
-
-    if st.button("Predict"):
-        fixed = fix_smiles(smiles_input)
-        if not fixed:
-            st.error("Invalid SMILES string.")
+    if mol is None:
+        return None
+    return np.array(AllChem.GetMorganFingerprintAsBitVect(mol, radius, nBits))
+
+def smiles_to_chemberta_embedding(smiles):
+    inputs = tokenizer(smiles, return_tensors="pt", padding=True, truncation=True).to(device)
+    with torch.no_grad():
+        outputs = chemberta_model(**inputs)
+        return outputs.last_hidden_state[:, 0, :].squeeze().cpu().numpy()
+
+def create_features(smiles):
+    fp = smiles_to_fingerprint(smiles)
+    if fp is None:
+        return None
+    emb = smiles_to_chemberta_embedding(smiles)
+    return np.concatenate([fp, emb])
+
+# Streamlit UI
+st.title("TransPolymer Property Predictor")
+user_input = st.text_input("Enter SMILES:")
+
+if st.button("Predict"):
+    if not user_input.strip():
+        st.error("Please enter a valid SMILES string.")
+    else:
+        features = create_features(user_input)
+        if features is None:
+            st.error("Invalid SMILES format.")
         else:
-            features = compute_features(fixed)
-            features_tensor = torch.tensor(features, dtype=torch.float32).unsqueeze(0)
-            embedding = embed_smiles([fixed])
-
-            with torch.no_grad():
-                pred = model(embedding, features_tensor)
-                result = pred.numpy().flatten()
-
-            properties = [
-                "Tensile Strength",
-                "Ionization Energy",
-                "Electron Affinity",
-                "logP",
-                "Refractive Index",
-                "Molecular Weight"
-            ]
-
-            predictions = {}
-            st.success("Predicted Polymer Properties:")
-            for prop, val in zip(properties, result):
-                st.write(f"**{prop}**: {val:.4f}")
-                predictions[prop] = val
-
-            # Save the prediction to MongoDB
-            save_to_db(smiles_input, predictions)
-            st.success("Prediction saved successfully!")
+            if input_scaler:
+                features = input_scaler.transform([features])
+            else:
+                features = [features]
+
+            raw_preds = np.mean([model.predict(features) for model in models], axis=0).flatten()
+
+            if target_scaler:
+                predictions = target_scaler.inverse_transform([raw_preds])[0]
+            else:
+                predictions = raw_preds
+
+            result_df = pd.DataFrame([predictions], columns=PROPERTY_NAMES)
+            result_df.insert(0, "SMILES", user_input)
+
+            st.success("Predicted Properties:")
+            st.dataframe(result_df.style.format(precision=4))
+
+            # Optional: save to CSV
+            history_path = "prediction_history.csv"
+            if os.path.exists(history_path):
+                existing = pd.read_csv(history_path)
+                result_df = pd.concat([existing, result_df], ignore_index=True)
+            result_df.to_csv(history_path, index=False)