Update prediction.py

prediction.py

@@ -1,88 +1,97 @@
-import streamlit as st
 import torch
+import torch.nn as nn
 import joblib
-import pandas as pd
-import numpy as np
 from rdkit import Chem
-from rdkit.Chem import AllChem
+from rdkit.Chem import Descriptors
 from transformers import AutoTokenizer, AutoModel
-import os
+import numpy as np
 
-# Load ChemBERTa model and tokenizer
+# Load tokenizer and embedding model
 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
-
-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):
+embedding_model = AutoModel.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
+
+# Load individual scalers
+target_keys = [
+    "Tensile_strength(Mpa)", 
+    "Ionization_Energy(eV)", 
+    "Electron_Affinity(eV)", 
+    "LogP", 
+    "Refractive_Index", 
+    "Molecular_Weight(g/mol)"
+]
+scalers = [joblib.load(f"scaler_{key.replace('/', '_').replace(' ', '_').replace('(', '').replace(')', '').replace('[', '').replace(']', '').replace('__', '_')}.joblib") for key in target_keys]
+
+# Descriptor function (must match training order)
+def compute_descriptors(smiles):
     mol = Chem.MolFromSmiles(smiles)
     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:
-            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)
+        raise ValueError("Invalid SMILES string.")
+    return np.array([
+        Descriptors.MolWt(mol),
+        Descriptors.MolLogP(mol),
+        Descriptors.TPSA(mol),
+        Descriptors.NumRotatableBonds(mol),
+        Descriptors.NumHDonors(mol),
+        Descriptors.NumHAcceptors(mol),
+        Descriptors.FractionCSP3(mol),
+        Descriptors.HeavyAtomCount(mol),
+        Descriptors.RingCount(mol),
+        Descriptors.MolMR(mol)
+    ], dtype=np.float32)
+
+# Model class must match training
+class TransformerRegressor(nn.Module):
+    def __init__(self, input_dim=768, descriptor_dim=10, d_model=768, nhead=4, num_layers=2, num_targets=6):
+        super().__init__()
+        self.descriptor_proj = nn.Linear(descriptor_dim, d_model)
+        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, batch_first=True)
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        self.regressor = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(2 * d_model, 256),
+            nn.ReLU(),
+            nn.Linear(256, num_targets)
+        )
+
+    def forward(self, embedding, descriptors):
+        desc_proj = self.descriptor_proj(descriptors).unsqueeze(1)  # (B, 1, d_model)
+        stacked = torch.cat([embedding.unsqueeze(1), desc_proj], dim=1)  # (B, 2, d_model)
+        encoded = self.transformer(stacked)  # (B, 2, d_model)
+        return self.regressor(encoded)
+
+# Load trained model
+model = TransformerRegressor()
+model.load_state_dict(torch.load("transformer_model.pt", map_location=torch.device("cpu")))
+model.eval()
+
+# Main prediction function
+def predict_properties(smiles):
+    try:
+        # Compute descriptors
+        descriptors = compute_descriptors(smiles)
+        descriptors_tensor = torch.tensor(descriptors).unsqueeze(0)  # (1, 10)
+
+        # Get embedding from ChemBERTa
+        inputs = tokenizer(smiles, return_tensors="pt")
+        with torch.no_grad():
+            outputs = embedding_model(**inputs)
+        embedding = outputs.last_hidden_state[:, 0, :]  # (1, 768)
+
+        # Predict
+        with torch.no_grad():
+            preds = model(embedding, descriptors_tensor)
+
+        # Inverse transform each prediction
+        preds_np = preds.numpy().flatten()
+        preds_rescaled = [
+            scalers[i].inverse_transform([[preds_np[i]]])[0][0] for i in range(len(scalers))
+        ]
+
+        # Prepare results
+        readable_keys = ["Tensile Strength", "Ionization Energy", "Electron Affinity", "logP", "Refractive Index", "Molecular Weight"]
+        results = dict(zip(readable_keys, np.round(preds_rescaled, 4)))
+
+        return results
+
+    except Exception as e:
+        return {"error": str(e)}