Update prediction.py

prediction.py

@@ -1,30 +1,32 @@
 import torch
 import torch.nn as nn
 import joblib
+import numpy as np
 from rdkit import Chem
 from rdkit.Chem import Descriptors
 from transformers import AutoTokenizer, AutoModel
-import numpy as np
 
-# Load tokenizer and model for embeddings
+# Load ChemBERTa tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
 embedding_model = AutoModel.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
 
-# Load saved scalers (for inverse_transform)
-scaler_tensile_strength = joblib.load("scaler_Tensile_strength_Mpa_.joblib")
-scaler_ionization_energy = joblib.load("scaler_Ionization_Energy_eV_.joblib")
-scaler_electron_affinity = joblib.load("scaler_Electron_Affinity_eV_.joblib")
-scaler_logp = joblib.load("scaler_LogP.joblib")
-scaler_refractive_index = joblib.load("scaler_Refractive_Index.joblib")
-scaler_molecular_weight = joblib.load("scaler_Molecular_Weight_g_mol_.joblib")
+# Load saved scalers for inverse transformations
+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")
+}
 
-# Descriptor function with exact order from training
-def compute_descriptors(smiles):
+# Descriptor calculation
+def compute_descriptors(smiles: str):
     mol = Chem.MolFromSmiles(smiles)
     if mol is None:
-        raise ValueError("Invalid SMILES")
-
-    return np.array([
+        raise ValueError("Invalid SMILES string.")
+    
+    descriptors = [
         Descriptors.MolWt(mol),
         Descriptors.MolLogP(mol),
         Descriptors.TPSA(mol),
@@ -35,9 +37,10 @@ def compute_descriptors(smiles):
         Descriptors.HeavyAtomCount(mol),
         Descriptors.RingCount(mol),
         Descriptors.MolMR(mol)
-    ], dtype=np.float32)
+    ]
+    return np.array(descriptors, dtype=np.float32)
 
-# Define your model class exactly like in training
+# Transformer regression model definition
 class TransformerRegressor(nn.Module):
     def __init__(self, input_dim, hidden_dim, num_layers, output_dim):
         super().__init__()
@@ -55,49 +58,48 @@ class TransformerRegressor(nn.Module):
     def forward(self, x):
         x = self.feat_proj(x)
         x = self.transformer_encoder(x)
-        x = x.mean(dim=1)
+        x = x.mean(dim=1)  # Global average pooling
         return self.regression_head(x)
 
-# Set model hyperparameters (must match training config)
-input_dim = 768  # ChemBERTa embedding size
+# Model hyperparameters (must match training)
+input_dim = 768
 hidden_dim = 256
 num_layers = 2
-output_dim = 6   # Number of properties predicted
+output_dim = 6
 
-# Load model
+# Load trained model
+device = torch.device("cpu")
 model = TransformerRegressor(input_dim, hidden_dim, num_layers, output_dim)
-model.load_state_dict(torch.load("transformer_model.pt", map_location=torch.device("cpu")))
+model.load_state_dict(torch.load("transformer_model.pt", map_location=device))
 model.eval()
 
-# Main prediction function
-def predict_properties(smiles):
+# Prediction function
+def predict_properties(smiles: str):
     try:
-        descriptors = compute_descriptors(smiles)
-        descriptors_tensor = torch.tensor(descriptors).unsqueeze(0)
+        # Validate SMILES and compute descriptors
+        _ = compute_descriptors(smiles)
 
-        # Get embedding
+        # ChemBERTa embedding (CLS token)
         inputs = tokenizer(smiles, return_tensors="pt")
         with torch.no_grad():
             outputs = embedding_model(**inputs)
-        emb = outputs.last_hidden_state[:, 0, :]  # CLS token output (1, 768)
+        embedding = outputs.last_hidden_state[:, 0, :]  # Shape: (1, 768)
 
-        # Forward pass
+        # Forward pass through model
         with torch.no_grad():
-            preds = model(emb)
+            preds = model(embedding)
 
         preds_np = preds.numpy()
+
+        # Inverse transform each property
+        keys = list(scalers.keys())
         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]])
+            scalers[keys[i]].inverse_transform(preds_np[:, [i]])
+            for i in range(output_dim)
         ], axis=1)
 
-        keys = ["Tensile Strength", "Ionization Energy", "Electron Affinity", "logP", "Refractive Index", "Molecular Weight"]
-        results = dict(zip(keys, preds_rescaled.flatten().round(4)))
-
+        # Create dictionary of results
+        results = {key: round(val, 4) for key, val in zip(keys, preds_rescaled.flatten())}
         return results
 
     except Exception as e: