Update prediction.py

prediction.py

@@ -6,26 +6,24 @@ from rdkit.Chem import Descriptors
 from transformers import AutoTokenizer, AutoModel
 import numpy as np
 
-# Load tokenizer and embedding model
+# Load tokenizer and model for embeddings
 tokenizer = AutoTokenizer.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
 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]
+# Load saved scalers (for inverse_transform)
+scaler_tensile_strength = joblib.load("scaler_Tensile_strength_Mpa_.joblib")  # Scaler for Tensile Strength
+scaler_ionization_energy = joblib.load("scaler_lonization_Energy_eV_.joblib")  # Scaler for Ionization Energy
+scaler_electron_affinity = joblib.load("scaler_Electron_Affinity_eV_.joblib")  # Scaler for Electron Affinity
+scaler_logp = joblib.load("scaler_LogP.joblib")  # Scaler for LogP
+scaler_refractive_index = joblib.load("scaler_Refractive_Index.joblib")  # Scaler for Refractive Index
+scaler_molecular_weight = joblib.load("scaler_Molecular_Weight_g_mol_.joblib")  # Scaler for Molecular Weight
 
-# Descriptor function (must match training order)
+# Descriptor function with exact order from training
 def compute_descriptors(smiles):
     mol = Chem.MolFromSmiles(smiles)
     if mol is None:
-        raise ValueError("Invalid SMILES string.")
+        raise ValueError("Invalid SMILES")
+
     return np.array([
         Descriptors.MolWt(mol),
         Descriptors.MolLogP(mol),
@@ -39,7 +37,7 @@ def compute_descriptors(smiles):
         Descriptors.MolMR(mol)
     ], dtype=np.float32)
 
-# Model class must match training
+# Define your model class exactly like in 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__()
@@ -57,9 +55,10 @@ class TransformerRegressor(nn.Module):
         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)
+        output = self.regressor(encoded)
+        return output
 
-# Load trained model
+# Load model
 model = TransformerRegressor()
 model.load_state_dict(torch.load("transformer_model.pt", map_location=torch.device("cpu")))
 model.eval()
@@ -67,29 +66,33 @@ 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
+        # Get embedding
         inputs = tokenizer(smiles, return_tensors="pt")
         with torch.no_grad():
             outputs = embedding_model(**inputs)
-        embedding = outputs.last_hidden_state[:, 0, :]  # (1, 768)
+        emb = outputs.last_hidden_state[:, 0, :]  # [CLS] token, shape (1, 768)
 
-        # Predict
+        # Forward pass
         with torch.no_grad():
-            preds = model(embedding, descriptors_tensor)
+            preds = model(emb, 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))
-        ]
+        # Inverse transform predictions using respective scalers
+        preds_np = preds.numpy()
+        preds_rescaled = np.concatenate([
+            scaler_tensile_strength.inverse_transform(preds_np[:, [0]]),
+            scaler_ionization_energy.inverse_transform(preds_np[:, [1]]),
+            scaler_electron_affinity.inverse_transform(preds_np[:, [2]]),
+            scaler_logp.inverse_transform(preds_np[:, [3]]),
+            scaler_refractive_index.inverse_transform(preds_np[:, [4]]),
+            scaler_molecular_weight.inverse_transform(preds_np[:, [5]])
+        ], axis=1)
 
-        # 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)))
+        # Round and format
+        keys = ["Tensile Strength", "Ionization Energy", "Electron Affinity", "logP", "Refractive Index", "Molecular Weight"]
+        results = dict(zip(keys, preds_rescaled.flatten().round(4)))
 
         return results