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
+from rdkit.Chem import Draw, Descriptors
 from transformers import BartForConditionalGeneration, BartTokenizer
 from admet_ai import ADMETModel
 import safe
@@ -59,51 +59,24 @@ creativity = st.sidebar.slider(
 num_molecules = st.sidebar.number_input(
     'Number of Molecules to Generate:',
     min_value=1,
-    max_value=5,
-    value=5,
-    help="Select the number of molecules you want to generate."
-)
-
-# **String Format Option (SMILES or SAFE)**
-string_format = st.sidebar.radio(
-    'String Format:',
-    ('SMILES', 'SAFE'),
-    help="Choose the format for displaying molecule strings."
+    max_value=3,  # Reduced from 5 to 3
+    value=3,
+    help="Select the number of molecules you want to generate (up to 3)."
 )
 
 # **Function to Generate Molecule Images**
-def generate_molecule_image(input_string, use_safe_visualization=True):
+def generate_molecule_image(smiles):
     """
-    Generates an image of the molecule from the input string.
-    Supports SAFE visualization if enabled.
+    Generates an image of the molecule from the SMILES string.
     """
     try:
-        if use_safe_visualization:
-            try:
-                # **Attempt to decode as SAFE string**
-                smiles = safe.decode(input_string)
-                # **Encode back to SAFE string**
-                safe_string = safe.encode(smiles)
-            except Exception as e:
-                # **Handle decoding errors**
-                st.error(f"Error decoding SAFE string: {e}")
-                return None
-            # **Generate SVG image with fragment highlights**
-            svg_str = safe.to_image(safe_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))
+        mol = Chem.MolFromSmiles(smiles)
+        if mol:
+            img = Draw.MolToImage(mol, size=(200, 200))
         else:
-            # **Generate standard molecule image**
-            mol = Chem.MolFromSmiles(input_string)
-            if mol:
-                img = Draw.MolToImage(mol, size=(200, 200))  # Adjusted size
-            else:
-                img = None
+            img = None
         return img
     except Exception as e:
-        # **Collect exceptions for later reporting**
         st.error(f"Error generating molecule image: {e}")
         return None
 
@@ -111,10 +84,15 @@ def generate_molecule_image(input_string, use_safe_visualization=True):
 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
+    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>
     """
-    components.html(f"""
-        <button onclick="navigator.clipboard.writeText('{text}')" style="padding:5px;">Copy</button>
-    """, height=45)
+    components.html(button_html, height=45)
 
 # **Generate Molecules Button**
 if st.button('Generate Molecules'):
@@ -188,40 +166,72 @@ if st.button('Generate Molecules'):
         # **Set 'Molecule Name' as index**
         df_results.set_index('Molecule Name', inplace=True)
         
-        # **Check if df_results is empty after merging**
-        if df_results.empty:
-            st.error("No valid molecules were generated after predictions. Please try adjusting the generation parameters.")
+        # **Select only desired ADMET properties**
+        admet_properties = [
+            'molecular weight', 'logP', 'hydrogen_bond_acceptors',
+            'hydrogen_bond_donors', 'QED', 'ClinTox', 'hERG', 'BBB_Martins'
+        ]
+        df_results_filtered = df_results[[
+            'SMILES', 'Valid'] + admet_properties]
+        
+        # **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**
             st.subheader('Generated Molecules')
             
             # **Determine number of columns per row**
-            cols_per_row = min(5, len(df_results))
+            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.iterrows()):
+            for idx, (mol_name, row) in enumerate(df_results_filtered.iterrows()):
                 smiles = row['SMILES']
-                img = generate_molecule_image(
-                    smiles,
-                    use_safe_visualization=(string_format == 'SAFE')
-                )
+                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
+                
                 with cols[idx % cols_per_row]:
                     if img is not None and isinstance(img, Image.Image):
                         st.image(img, caption=mol_name)
                     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 molecule string in chosen format**
-                    string_to_display = safe.encode(smiles) if string_format == 'SAFE' else smiles
-                    st.code(string_to_display)
+                    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)
+                    
                     # **Copy-to-clipboard functionality**
-                    st_copy_button(string_to_display, key=f'copy_{mol_name}')
-                    # **Display ADMET properties**
+                    st_copy_button(
+                        safe_string if st.session_state[toggle_key] else smiles,
+                        key=f'copy_{mol_name}'
+                    )
+                    
+                    # **Display selected ADMET properties**
                     st.write("**ADMET Properties:**")
-                    st.write(row.drop(['SMILES', 'Valid']))
+                    # 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.")
-
-