Update prediction.py

prediction.py

@@ -8,10 +8,9 @@ from rdkit import Chem
 from rdkit.Chem import Descriptors
 from rdkit.Chem import AllChem
 from datetime import datetime
-from db import get_database  # Ensure this module is available
+from db import get_database
 import random
-
- 
+import os  # <-- Added for debugging file paths
 
 # ------------------------ Ensuring Deterministic Behavior ------------------------
 random.seed(42)
@@ -20,7 +19,6 @@ torch.manual_seed(42)
 torch.backends.cudnn.deterministic = True
 torch.backends.cudnn.benchmark = False
 
-# Check if CUDA is available for GPU acceleration
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # ------------------------ Load ChemBERTa Model + Tokenizer ------------------------
@@ -29,7 +27,7 @@ def load_chemberta():
     tokenizer = AutoTokenizer.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
     model = AutoModel.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
     model.eval()
-    model.to(device)  # Send model to GPU if available
+    model.to(device)
     return tokenizer, model
 
 # ------------------------ Load Scalers ------------------------
@@ -51,7 +49,7 @@ class TransformerRegressor(nn.Module):
             d_model=embedding_dim,
             nhead=8,
             dim_feedforward=ff_dim,
-            dropout=0.0,  # No dropout for consistency
+            dropout=0.0,
             batch_first=True
         )
         self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
@@ -72,18 +70,15 @@ class TransformerRegressor(nn.Module):
 # ------------------------ Load Model ------------------------
 @st.cache_resource
 def load_model():
-    # Initialize the model architecture first
     model = TransformerRegressor()
-
-    # Load the state_dict (weights) from the saved model file
     try:
-        state_dict = torch.load("transformer_model(1).bin",map_location=device) # Ensure loading on the correct device
+        print("Files in working directory:", os.listdir())  # <-- Debug print
+        state_dict = torch.load("transformer_model.bin", map_location=device)
         model.load_state_dict(state_dict)
         model.eval()
-        model.to(device)  # Send model to GPU if available
+        model.to(device)
     except Exception as e:
         raise ValueError(f"Failed to load model: {e}")
-
     return model
 
 # ------------------------ Descriptors ------------------------
@@ -91,7 +86,6 @@ def compute_descriptors(smiles: str):
     mol = Chem.MolFromSmiles(smiles)
     if mol is None:
         raise ValueError("Invalid SMILES string.")
-    
     descriptors = [
         Descriptors.MolWt(mol),
         Descriptors.MolLogP(mol),
@@ -119,7 +113,7 @@ def get_chemberta_embedding(smiles: str, tokenizer, chemberta):
     inputs = tokenizer(smiles, return_tensors="pt")
     with torch.no_grad():
         outputs = chemberta(**inputs)
-    return outputs.last_hidden_state.mean(dim=1)  # Use average instead of CLS token
+    return outputs.last_hidden_state.mean(dim=1)
 
 # ------------------------ Save to DB ------------------------
 def save_to_db(smiles, predictions):
@@ -141,7 +135,6 @@ def show():
 
     if st.button("Predict"):
         try:
-            # Load the model
             model = load_model()
 
             mol = Chem.MolFromSmiles(smiles_input)
@@ -149,10 +142,8 @@ def show():
                 st.error("Invalid SMILES string.")
                 return
 
-            # Load the ChemBERTa tokenizer and model
             tokenizer, chemberta = load_chemberta()
 
-            # Compute Descriptors, Fingerprints, and Embedding
             descriptors = compute_descriptors(smiles_input)
             descriptors_tensor = torch.tensor(descriptors, dtype=torch.float32).unsqueeze(0)
 
@@ -161,7 +152,6 @@ def show():
 
             embedding = get_chemberta_embedding(smiles_input, tokenizer, chemberta)
 
-            # Combine Inputs and Make Prediction
             combined_input = torch.cat([embedding, descriptors_tensor, fingerprint_tensor], dim=1)
             combined = combined_input.unsqueeze(1)
 
@@ -171,7 +161,6 @@ def show():
             preds_np = preds.numpy()
             keys = list(scalers.keys())
 
-            # Rescale Predictions
             preds_rescaled = np.concatenate([
                 scalers[keys[i]].inverse_transform(preds_np[:, [i]])
                 for i in range(6)
@@ -183,7 +172,6 @@ def show():
             for key, val in results.items():
                 st.markdown(f"**{key}**: {val}")
 
-            # Save the results to the database
             save_to_db(smiles_input, results)
 
         except Exception as e: