Update app.py

app.py

@@ -2,7 +2,7 @@ import streamlit as st
 import torch
 import os
 from rdkit import Chem
-from rdkit.Chem import Draw, Descriptors
+from rdkit.Chem import Draw
 from transformers import BartForConditionalGeneration, BartTokenizer
 from admet_ai import ADMETModel
 import safe
@@ -11,6 +11,7 @@ from PIL import Image
 import cairosvg
 import pandas as pd
 import streamlit.components.v1 as components
+import json  # For safely encoding text in JavaScript
 
 # **Page Configuration**
 st.set_page_config(
@@ -65,16 +66,26 @@ num_molecules = st.sidebar.number_input(
 )
 
 # **Function to Generate Molecule Images**
-def generate_molecule_image(smiles):
+def generate_molecule_image(input_string, use_safe=False):
     """
-    Generates an image of the molecule from the SMILES string.
+    Generates an image of the molecule from the input string.
+    If use_safe is True, input_string is treated as a SAFE string.
     """
     try:
-        mol = Chem.MolFromSmiles(smiles)
-        if mol:
-            img = Draw.MolToImage(mol, size=(200, 200))
+        if use_safe and input_string is not None:
+            # Generate image from SAFE encoding
+            svg_str = safe.to_image(input_string)
+            # Convert SVG to PNG bytes
+            png_bytes = cairosvg.svg2png(bytestring=svg_str.encode('utf-8'))
+            # Create an image object
+            img = Image.open(io.BytesIO(png_bytes))
         else:
-            img = None
+            # Generate standard molecule image
+            mol = Chem.MolFromSmiles(input_string)
+            if mol:
+                img = Draw.MolToImage(mol, size=(200, 200))
+            else:
+                img = None
         return img
     except Exception as e:
         st.error(f"Error generating molecule image: {e}")
@@ -82,17 +93,17 @@ def generate_molecule_image(smiles):
 
 # **Function to Create Copy-to-Clipboard Button**
 def st_copy_button(text, key):
-    """
-    Creates a copy-to-clipboard button for the given text.
-    Adjusted to position the button without overlapping the text.
-    """
-    # Adjusted styling to position the button
+    """Creates a copy-to-clipboard button placed appropriately."""
+    # Safely encode the text for JavaScript
+    escaped_text = json.dumps(text)
     button_html = f"""
-        <div style="display: flex; justify-content: flex-end;">
-            <button onclick="navigator.clipboard.writeText('{text}')" style="padding:5px; margin-top: -40px; position: relative; z-index: 1;">Copy</button>
-        </div>
+    <div style="text-align: right; margin-top: -10px; margin-bottom: 10px;">
+        <button onclick="navigator.clipboard.writeText({escaped_text})" style="
+            padding:5px;
+        ">Copy</button>
+    </div>
     """
-    components.html(button_html, height=45)
+    components.html(button_html, height=35)
 
 # **Generate Molecules Button**
 if st.button('Generate Molecules'):
@@ -134,68 +145,64 @@ if st.button('Generate Molecules'):
     })
     
     # **Invalid SMILES Check**
-    from rdkit import Chem
-    
-    # **Function to validate SMILES**
+    # Function to validate SMILES
     def is_valid_smile(smile):
         return Chem.MolFromSmiles(smile) is not None
     
-    # **Apply validation function**
+    # Apply validation function
     df_molecules['Valid'] = df_molecules['SMILES'].apply(is_valid_smile)
     df_valid = df_molecules[df_molecules['Valid']].copy()
     
-    # **Inform user if any molecules were invalid**
+    # Inform user if any molecules were invalid
     invalid_molecules = df_molecules[~df_molecules['Valid']]
     if not invalid_molecules.empty:
         st.warning(f"{len(invalid_molecules)} generated molecules were invalid and excluded from predictions.")
     
-    # **Check if there are valid molecules to proceed**
+    # Check if there are valid molecules to proceed
     if df_valid.empty:
         st.error("No valid molecules were generated. Please try adjusting the generation parameters.")
     else:
-        # **ADMET Predictions**
+        # ADMET Predictions
         preds = admet_model.predict(smiles=df_valid['SMILES'].tolist())
         
-        # **Ensure 'SMILES' is a column in preds**
+        # Ensure 'SMILES' is a column in preds
         if 'SMILES' not in preds.columns:
             preds['SMILES'] = df_valid['SMILES'].values
         
-        # **Merge predictions with valid molecules**
+        # Merge predictions with valid molecules
         df_results = pd.merge(df_valid, preds, on='SMILES', how='inner')
         
-        # **Set 'Molecule Name' as index**
+        # Set 'Molecule Name' as index
         df_results.set_index('Molecule Name', inplace=True)
         
-        # **Select only desired ADMET properties**
+        # Select only desired ADMET properties
         admet_properties = [
-            'molecular_weight', 'logP', 'hydrogen_bond_acceptors',
+            'molecular weight', 'logP', 'hydrogen_bond_acceptors',
             'hydrogen_bond_donors', 'QED', 'ClinTox', 'hERG', 'BBB_Martins'
         ]
-        df_results_filtered = df_results[[
-            'SMILES', 'Valid'] + admet_properties]
+        df_results_filtered = df_results[['SMILES', 'Valid'] + admet_properties]
         
-        # **Check if df_results_filtered is empty after filtering**
+        # Check if df_results_filtered is empty after filtering
         if df_results_filtered.empty:
             st.error("No valid ADMET predictions were obtained. Please try adjusting the generation parameters.")
         else:
-            # **Display Molecules**
+            # Display Molecules
             st.subheader('Generated Molecules')
-            
-            # **Determine number of columns per row**
             cols_per_row = min(3, len(df_results_filtered))  # Max 3 columns
-            
-            # **Create columns in Streamlit**
             cols = st.columns(cols_per_row)
             
-            # **Iterate over each molecule to display**
             for idx, (mol_name, row) in enumerate(df_results_filtered.iterrows()):
                 smiles = row['SMILES']
-                img = generate_molecule_image(smiles)
                 
-                # **Initialize session state for toggle buttons**
-                toggle_key = f'toggle_{mol_name}'
-                if toggle_key not in st.session_state:
-                    st.session_state[toggle_key] = False  # False means SMILES is displayed
+                # Attempt to encode to SAFE
+                try:
+                    safe_string = safe.encode(smiles)
+                except Exception as e:
+                    safe_string = None
+                    st.error(f"Could not convert to SAFE encoding for {mol_name}: {e}")
+                
+                # Generate molecule image (SMILES or SAFE)
+                img = generate_molecule_image(smiles)
                 
                 with cols[idx % cols_per_row]:
                     if img is not None and isinstance(img, Image.Image):
@@ -203,35 +210,25 @@ if st.button('Generate Molecules'):
                     else:
                         st.error(f"Could not generate image for {mol_name}")
                     
-                    # **Toggle Button to Switch Between SMILES and SAFE**
-                    if st.button(
-                        'Toggle to SAFE' if not st.session_state[toggle_key] else 'Toggle to SMILES',
-                        key=toggle_key + '_button'
-                    ):
-                        st.session_state[toggle_key] = not st.session_state[toggle_key]
+                    # Display SMILES string
+                    st.write("**SMILES:**")
+                    st.text(smiles)
+                    st_copy_button(smiles, key=f'copy_smiles_{mol_name}')
                     
-                    # **Display molecule string in chosen format**
-                    if st.session_state[toggle_key]:
-                        # **Attempt to encode to SAFE**
-                        try:
-                            safe_string = safe.encode(smiles)
-                            st.code(safe_string)
-                        except Exception as e:
-                            st.error(f"Could not convert to SAFE encoding: {e}")
-                            st.code(smiles)
-                    else:
-                        st.code(smiles)
+                    # Display SAFE encoding if available
+                    if safe_string:
+                        st.write("**SAFE Encoding:**")
+                        st.text(safe_string)
+                        st_copy_button(safe_string, key=f'copy_safe_{mol_name}')
+                        # Optionally display SAFE visualization
+                        safe_img = generate_molecule_image(safe_string, use_safe=True)
+                        if safe_img is not None:
+                            st.image(safe_img, caption=f"{mol_name} (SAFE Visualization)")
                     
-                    # **Copy-to-clipboard functionality**
-                    st_copy_button(
-                        safe_string if st.session_state[toggle_key] else smiles,
-                        key=f'copy_{mol_name}'
-                    )
-                    
-                    # **Display selected ADMET properties**
+                    # Display selected ADMET properties
                     st.write("**ADMET Properties:**")
-                    # Drop 'SMILES' and 'Valid' columns for display
                     admet_data = row.drop(['SMILES', 'Valid'])
                     st.write(admet_data)
 else:
     st.write("Click the 'Generate Molecules' button to generate beta-lactam molecules.")
+