Update app.py

app.py

@@ -1,448 +1,253 @@
-#!/usr/bin/env python3
-"""
-ChemBERTa SMILES Utilities Dashboard
-A Streamlit application for molecular prediction and visualization
-"""
+# app.py
+# To run this app, save the code as app.py and run:
+# streamlit run app.py
+#
+# You also need to install the following libraries:
+# pip install streamlit torch transformers bitsandbytes rdkit-pypi py3Dmol pandas
 
 import streamlit as st
+import streamlit.components.v1 as components
 import torch
 from transformers import AutoModelForMaskedLM, AutoTokenizer, pipeline, BitsAndBytesConfig
 from rdkit import Chem
-from rdkit.Chem import Draw, rdFMCS, AllChem
+from rdkit.Chem import Draw, AllChem
 from rdkit.Chem.Draw import MolToImage
 import pandas as pd
-import io
-import base64
 import logging
-import streamlit.components.v1 as components
-import sys
-import os
-
-# Check if running in Streamlit context
-def is_streamlit_context():
-    """Check if we're running in a Streamlit context"""
-    try:
-        import streamlit.runtime.scriptrunner as sr
-        return sr.get_script_run_ctx() is not None
-    except:
-        return False
-
-# Only proceed if we're in a Streamlit context or being run by streamlit
-if not is_streamlit_context() and __name__ == "__main__":
-    print("This app must be run with: streamlit run app.py")
-    print("Please use the command: streamlit run app.py --server.port=7860 --server.address=0.0.0.0")
-    sys.exit(1)
 
 # Set up logging to monitor quantization effects
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Page configuration - only if in streamlit context
-if is_streamlit_context() or 'streamlit' in sys.modules:
-    st.set_page_config(
-        page_title="ChemBERTa SMILES Utilities Dashboard",
-        page_icon="🧪",
-        layout="wide"
-    )
-
-# --- Quantization Configuration ---
-@st.cache_resource
-def get_quantization_config():
+# --- Page Configuration ---
+st.set_page_config(
+    page_title="ChemBERTa SMILES Utilities",
+    page_icon="🔬",
+    layout="wide",
+)
+
+# --- Model Loading (Cached for Performance) ---
+
+@st.cache_resource(show_spinner="Loading ChemBERTa model...")
+def load_models():
     """
-    Configure 8-bit quantization for model optimization.
-    Falls back gracefully if bitsandbytes is not available.
+    Load the tokenizer and model, wrapped in a Streamlit cache resource decorator
+    to ensure it only runs once per session.
     """
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+    quantization_config = None
+
     try:
-        # Only use quantization on CUDA
-        if not torch.cuda.is_available():
-            logger.info("CUDA not available, skipping quantization")
-            return None
-            
-        # 8-bit quantization configuration - good balance of speed and quality
         quantization_config = BitsAndBytesConfig(
             load_in_8bit=True,
             bnb_8bit_compute_dtype=torch.float16,
-            bnb_8bit_use_double_quant=True,  # Nested quantization for better compression
+            bnb_8bit_use_double_quant=True,
         )
-        logger.info("8-bit quantization configuration loaded successfully")
-        return quantization_config
+        logger.info("8-bit quantization configuration created.")
     except ImportError:
-        logger.warning("bitsandbytes not available, falling back to standard loading")
-        return None
+        logger.warning("bitsandbytes not available, falling back to standard loading.")
     except Exception as e:
-        logger.warning(f"Quantization setup failed: {e}, using standard loading")
-        return None
-
-def get_torch_dtype():
-    """Get appropriate torch dtype based on available hardware."""
-    if torch.cuda.is_available():
-        return torch.float16  # Use half precision on GPU
-    else:
-        return torch.float32  # Keep full precision on CPU
-
-# --- Optimized Model Loading ---
-@st.cache_resource
-def load_optimized_models():
-    """Load models with quantization and other optimizations."""
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    torch_dtype = get_torch_dtype()
-    quantization_config = get_quantization_config()
-
-    logger.info(f"Loading models on device: {device} with dtype: {torch_dtype}")
+        logger.warning(f"Quantization setup failed: {e}, using standard loading.")
 
-    # Model names
     model_name = "seyonec/PubChem10M_SMILES_BPE_450k"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-    try:
-        # Load tokenizer (doesn't need quantization)
-        fill_mask_tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-        # Load model with quantization if available
-        model_kwargs = {
-            "torch_dtype": torch_dtype,
-        }
-
-        if quantization_config is not None and torch.cuda.is_available():
-            model_kwargs["quantization_config"] = quantization_config
-            model_kwargs["device_map"] = "auto"
-        else:
-            # For CPU or non-quantized loading
-            model_kwargs["device_map"] = None
-
-        # Masked LM Model
-        fill_mask_model = AutoModelForMaskedLM.from_pretrained(
-            model_name,
-            **model_kwargs
-        )
-
-        # Move to device if not using device_map
-        if model_kwargs["device_map"] is None and torch.cuda.is_available():
-            fill_mask_model.to(device)
-
-        # Set model to evaluation mode for inference
-        fill_mask_model.eval()
-
-        # Create pipeline with proper device handling
-        pipeline_device = 0 if torch.cuda.is_available() else -1
-
-        fill_mask_pipeline = pipeline(
-            'fill-mask',
-            model=fill_mask_model,
-            tokenizer=fill_mask_tokenizer,
-            device=pipeline_device,
-        )
+    model_kwargs = {"torch_dtype": torch_dtype}
+    if quantization_config and torch.cuda.is_available():
+        model_kwargs["quantization_config"] = quantization_config
+        model_kwargs["device_map"] = "auto"
+    elif torch.cuda.is_available():
+        model_kwargs["device_map"] = "auto"
 
-        logger.info("Models loaded successfully with optimizations")
-        return fill_mask_tokenizer, fill_mask_model, fill_mask_pipeline
-
-    except Exception as e:
-        logger.error(f"Error loading optimized models: {e}")
-        # Fallback to standard loading
-        logger.info("Falling back to standard model loading...")
-        return load_standard_models(model_name)
-
-def load_standard_models(model_name):
-    """Fallback standard model loading without quantization."""
     try:
-        fill_mask_tokenizer = AutoTokenizer.from_pretrained(model_name)
-        fill_mask_model = AutoModelForMaskedLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.float32
-        )
-        
-        # Determine device for standard loading
+        model = AutoModelForMaskedLM.from_pretrained(model_name, **model_kwargs)
+        model.eval()
+        pipeline_device = model.device.index if hasattr(model.device, 'type') and model.device.type == "cuda" else -1
+        fill_mask_pipeline = pipeline('fill-mask', model=model, tokenizer=tokenizer, device=pipeline_device)
+        logger.info("Models loaded successfully with optimizations.")
+        return tokenizer, fill_mask_pipeline
+    except Exception as e:
+        logger.error(f"Error loading optimized models: {e}. Retrying with standard loading.")
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForMaskedLM.from_pretrained(model_name)
         device_idx = 0 if torch.cuda.is_available() else -1
-        
         if torch.cuda.is_available():
-            fill_mask_model.to("cuda")
-            
-        fill_mask_pipeline = pipeline(
-            'fill-mask', 
-            model=fill_mask_model, 
-            tokenizer=fill_mask_tokenizer, 
-            device=device_idx
-        )
+            model.to("cuda")
+        fill_mask_pipeline = pipeline('fill-mask', model=model, tokenizer=tokenizer, device=device_idx)
+        return tokenizer, fill_mask_pipeline
+
+# Load the models once
+fill_mask_tokenizer, fill_mask_pipeline = load_models()
+
+
+# --- Molecule & Visualization Helpers ---
 
-        return fill_mask_tokenizer, fill_mask_model, fill_mask_pipeline
-    except Exception as e:
-        logger.error(f"Failed to load models: {e}")
-        if is_streamlit_context():
-            st.error(f"Failed to load models: {e}")
-        return None, None, None
-
-# --- Memory Management Utilities ---
-def clear_gpu_cache():
-    """Clear CUDA cache to free up memory."""
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
-# --- Helper Functions ---
 def get_mol(smiles):
     """Converts SMILES to RDKit Mol object and Kekulizes it."""
     mol = Chem.MolFromSmiles(smiles)
-    if mol is None:
-        return None
-    try:
-        Chem.Kekulize(mol)
-    except: # Kekulization can fail for some structures
-        pass
+    if mol:
+        try:
+            Chem.Kekulize(mol)
+        except Exception:
+            pass
     return mol
 
 def find_matches_one(mol, submol_smarts):
-    """Finds all matching atoms for a SMARTS pattern in a molecule."""
-    if not mol or not submol_smarts:
-        return []
+    """Finds all matching atoms for a SMARTS pattern."""
+    if not mol or not submol_smarts: return []
     submol = Chem.MolFromSmarts(submol_smarts)
-    if not submol:
-        return []
-    matches = mol.GetSubstructMatches(submol)
-    return matches
+    return mol.GetSubstructMatches(submol) if submol else []
 
 def get_image_with_highlight(mol, atomset=None, size=(300, 300)):
-    """Draws molecule with optional atom highlighting."""
-    if mol is None:
-        return None
-    highlight_color = (0, 1, 0, 0.5) # Green with some transparency
-
-    # Ensure atomset contains integers if not None or empty
-    valid_atomset = []
-    if atomset:
-        try:
-            valid_atomset = [int(a) for a in atomset]
-        except (ValueError, TypeError):
-            logger.warning(f"Invalid atom in atomset: {atomset}. Proceeding without highlighting problematic atoms.")
-            valid_atomset = [int(a) for a in atomset if str(a).isdigit()] # Filter out non-integers
-
-    img = MolToImage(mol, size=size, fitImage=True,
-                     highlightAtoms=valid_atomset if valid_atomset else [],
-                     highlightAtomColors={i: highlight_color for i in valid_atomset} if valid_atomset else {})
-    return img
-
-def generate_3d_structure(mol):
-    """Generate 3D coordinates for a molecule."""
-    if mol is None:
-        return None
-    
-    # Create a copy to avoid modifying the original
-    mol_3d = Chem.Mol(mol)
-    
-    # Add hydrogens
-    mol_3d = Chem.AddHs(mol_3d)
-    
-    # Generate 3D coordinates
+    """Draws a 2D molecule image with optional atom highlighting."""
+    if mol is None: return None
+    valid_atomset = [int(a) for a in atomset if str(a).isdigit()] if atomset else []
+    return MolToImage(mol, size=size, fitImage=True,
+                      highlightAtoms=valid_atomset,
+                      highlightAtomColors={i: (0, 1, 0, 0.5) for i in valid_atomset})
+
+def generate_3d_view_html(smiles):
+    """Generates an interactive 3D molecule view using py3Dmol."""
+    if not smiles: return None
+    mol = get_mol(smiles)
+    if not mol: return "<p>Invalid SMILES for 3D view.</p>"
     try:
-        AllChem.EmbedMolecule(mol_3d, randomSeed=42)
-        AllChem.UFFOptimizeMolecule(mol_3d)
-        return mol_3d
-    except:
-        # If 3D generation fails, return None
-        return None
-
-def mol_to_3d_html(mol):
-    """Convert molecule to 3D HTML representation using py3Dmol."""
-    if mol is None:
-        return None
-    
-    # Generate SDF string
-    sdf = Chem.MolToMolBlock(mol)
-    
-    # Create 3D viewer HTML
-    html_template = """
-    <div id="3dmolviewer_{id}" style="height: 400px; width: 100%; position: relative;" class="viewer_3Dmoljs"></div>
-    <script src="https://cdnjs.cloudflare.com/ajax/libs/3Dmol/2.0.4/3Dmol-min.js"></script>
-    <script>
-    let viewer_{id} = $3Dmol.createViewer(document.getElementById('3dmolviewer_{id}'), {{
-        defaultcolors: $3Dmol.rasmolElementColors
-    }});
-    viewer_{id}.addModel(`{sdf}`, 'sdf');
-    viewer_{id}.setStyle({{}}, {{stick: {{}}}});
-    viewer_{id}.zoomTo();
-    viewer_{id}.render();
-    </script>
-    """
-    
-    import random
-    viewer_id = random.randint(1000, 9999)
-    
-    html_content = html_template.format(id=viewer_id, sdf=sdf.replace('`', '\\`'))
-    
-    return html_content
+        mol_3d = Chem.AddHs(mol)
+        AllChem.EmbedMolecule(mol_3d, randomSeed=42, useRandomCoords=True)
+        AllChem.MMFFOptimizeMolecule(mol_3d)
+        sdf_data = Chem.MolToMolBlock(mol_3d)
+
+        viewer = py3Dmol.view(width=350, height=350)
+        viewer.setBackgroundColor('#FFFFFF')
+        viewer.addModel(sdf_data, "sdf")
+        viewer.setStyle({'stick': {}, 'sphere': {'scale': 0.25}})
+        viewer.zoomTo()
+        return viewer._make_html()
+    except Exception as e:
+        logger.error(f"Failed to generate 3D view for {smiles}: {e}")
+        return f"<p>Error generating 3D view: {e}</p>"
 
-# --- Streamlit Interface Functions ---
+# --- Core Application Logic ---
 
-def predict_and_visualize_masked_smiles(smiles_mask, substructure_smarts_highlight="CC=CC"):
+def run_masked_smiles_prediction(smiles_mask, substructure_smarts_highlight):
     """
-    Predicts masked tokens in a SMILES string, shows scores, and visualizes molecules.
+    Handles the logic for the masked SMILES prediction tab.
     """
-    # Load models when needed
-    try:
-        models = load_optimized_models()
-        if models[0] is None:  # Check if loading failed
-            st.error("Failed to load models. Please check the logs.")
-            return
-        fill_mask_tokenizer, fill_mask_model, fill_mask_pipeline = models
-    except Exception as e:
-        st.error(f"Error loading models: {str(e)}")
-        return
-    
     if fill_mask_tokenizer.mask_token not in smiles_mask:
-        st.error("Error: Input SMILES must contain a mask token (e.g., <mask>).")
+        st.error(f"Error: Input SMILES must contain a mask token (e.g., {fill_mask_tokenizer.mask_token}).")
         return
 
-    try:
-        # Use torch.no_grad() for inference to save memory
-        with torch.no_grad():
-            predictions = fill_mask_pipeline(smiles_mask, top_k=10)
-    except Exception as e:
-        clear_gpu_cache()
-        st.error(f"Error during prediction: {str(e)}")
-        return
+    with st.spinner("Predicting completions..."):
+        try:
+            with torch.no_grad():
+                predictions = fill_mask_pipeline(smiles_mask, top_k=10)
+        except Exception as e:
+            st.error(f"An error occurred during prediction: {e}")
+            if torch.cuda.is_available(): torch.cuda.empty_cache()
+            return
 
-    results_data = []
-    valid_predictions = []
-    valid_predictions_count = 0
+        results = []
+        for pred in predictions:
+            if len(results) >= 5: break
+            predicted_smiles = pred['sequence']
+            mol = get_mol(predicted_smiles)
+            if mol:
+                atom_matches = find_matches_one(mol, substructure_smarts_highlight)
+                results.append({
+                    "smiles": predicted_smiles,
+                    "score": f"{pred['score']:.4f}",
+                    "image_2d": get_image_with_highlight(mol, atomset=atom_matches[0] if atom_matches else []),
+                    "html_3d": generate_3d_view_html(predicted_smiles)
+                })
+        
+        if torch.cuda.is_available(): torch.cuda.empty_cache()
+        st.session_state.prediction_results = results
 
-    for pred in predictions:
-        if valid_predictions_count >= 5:
-            break
 
-        predicted_smiles = pred['sequence']
-        score = pred['score']
+# --- Streamlit UI Definition ---
 
-        mol = get_mol(predicted_smiles)
-        if mol:
-            results_data.append({"Predicted SMILES": predicted_smiles, "Score": f"{score:.4f}"})
-            valid_predictions.append((mol, predicted_smiles, score))
-            valid_predictions_count += 1
+st.title("🔬 ChemBERTa SMILES Utilities Dashboard (2D & 3D)")
+st.markdown("A tool to predict masked tokens in SMILES strings and visualize molecules, powered by ChemBERTa and Streamlit.")
 
-    if valid_predictions_count == 0:
-        st.warning("No valid molecules found for top predictions.")
-        return
+tab1, tab2 = st.tabs(["Masked SMILES Prediction", "Molecule Viewer (2D & 3D)"])
 
-    # Display results table
-    df_results = pd.DataFrame(results_data)
-    st.subheader("Top Predictions & Scores")
-    st.dataframe(df_results, use_container_width=True)
+# --- Tab 1: Masked SMILES Prediction ---
+with tab1:
+    st.header("Predict and Visualize Masked SMILES")
+    st.markdown("Enter a SMILES string with a `<mask>` token to predict possible completions.")
 
-    # Display molecule visualizations
-    st.subheader("Predicted Molecule Visualizations")
-    
-    for i, (mol, smiles, score) in enumerate(valid_predictions):
-        st.write(f"**Prediction {i+1}:** {smiles} (Score: {score:.4f})")
-        
-        col1, col2 = st.columns(2)
-        
-        with col1:
-            st.write("**2D Structure:**")
-            atom_matches_indices = []
-            if substructure_smarts_highlight:
-                matches = find_matches_one(mol, substructure_smarts_highlight)
-                if matches:
-                    atom_matches_indices = list(matches[0])
-            
-            img_2d = get_image_with_highlight(mol, atomset=atom_matches_indices)
-            if img_2d:
-                st.image(img_2d, use_column_width=True)
-        
-        with col2:
-            st.write("**3D Structure:**")
-            mol_3d = generate_3d_structure(mol)
-            if mol_3d:
-                html_3d = mol_to_3d_html(mol_3d)
-                if html_3d:
-                    components.html(html_3d, height=450)
-            else:
-                st.write("3D structure generation failed for this molecule.")
-        
-        st.divider()
-
-    # Clear cache after inference
-    clear_gpu_cache()
-    st.success("Prediction successful!")
-
-def display_molecule_image(smiles_string):
-    """
-    Displays both 2D and 3D images of a molecule from its SMILES string.
-    """
-    if not smiles_string:
-        st.error("Please enter a SMILES string.")
-        return
-    
-    mol = get_mol(smiles_string)
-    if mol is None:
-        st.error("Invalid SMILES string.")
-        return
-    
-    st.success("Molecule displayed successfully!")
-    
-    col1, col2 = st.columns(2)
-    
-    with col1:
-        st.subheader("2D Structure")
-        img_2d = MolToImage(mol, size=(400, 400), fitImage=True)
-        st.image(img_2d, use_column_width=True)
-    
-    with col2:
-        st.subheader("3D Structure")
-        mol_3d = generate_3d_structure(mol)
-        if mol_3d:
-            html_3d = mol_to_3d_html(mol_3d)
-            if html_3d:
-                components.html(html_3d, height=450)
-        else:
-            st.write("3D structure generation failed for this molecule.")
-
-# --- Main Streamlit App ---
-def main():
-    # Only run if in Streamlit context
-    if not is_streamlit_context():
-        return
-        
-    # Initialize session state
-    if 'initialized' not in st.session_state:
-        st.session_state.initialized = True
-    
-    st.title("🧪 ChemBERTa SMILES Utilities Dashboard")
-    
-    # Sidebar for navigation
-    st.sidebar.title("Navigation")
-    tab_selection = st.sidebar.selectbox(
-        "Choose a tool:",
-        ["Masked SMILES Prediction", "Molecule Viewer"]
-    )
-    
-    if tab_selection == "Masked SMILES Prediction":
-        st.header("Masked SMILES Prediction")
-        st.markdown("Enter a SMILES string with a `<mask>` token (e.g., `C1=CC=CC<mask>C1`) to predict possible completions.")
-        
+    with st.form(key="prediction_form"):
         col1, col2 = st.columns(2)
         with col1:
             smiles_input_masked = st.text_input(
-                "SMILES String with Mask", 
-                value="C1=CC=CC<mask>C1"
+                "SMILES String with Mask",
+                value="C1=CC=CC<mask>C1",
+                help=f"The mask token is `{fill_mask_tokenizer.mask_token}`"
             )
         with col2:
             substructure_input = st.text_input(
-                "Substructure to Highlight (SMARTS)", 
-                value="C=C"
+                "Substructure to Highlight (SMARTS)",
+                value="C=C",
+                help="Enter a SMARTS pattern to highlight in the 2D images."
             )
         
-        if st.button("Predict and Visualize", type="primary"):
-            with st.spinner("Predicting masked SMILES..."):
-                predict_and_visualize_masked_smiles(smiles_input_masked, substructure_input)
-    
-    elif tab_selection == "Molecule Viewer":
-        st.header("Molecule Viewer")
-        st.markdown("Enter a SMILES string to display its 2D and 3D structure.")
-        
-        smiles_input_viewer = st.text_input(
-            "SMILES String", 
-            value="C1=CC=CC=C1"
-        )
+        predict_button = st.form_submit_button("Predict and Visualize", use_container_width=True)
+
+    if predict_button:
+        run_masked_smiles_prediction(smiles_input_masked, substructure_input)
+
+    if 'prediction_results' in st.session_state and st.session_state.prediction_results:
+        results = st.session_state.prediction_results
+        st.subheader("Top 5 Valid Predictions")
+
+        # Display results in a table
+        df_data = [{"Predicted SMILES": r["smiles"], "Score": r["score"]} for r in results]
+        st.dataframe(pd.DataFrame(df_data), use_container_width=True)
+
+        st.markdown("---")
         
-        if st.button("View Molecule", type="primary"):
-            with st.spinner("Generating molecule structures..."):
-                display_molecule_image(smiles_input_viewer)
+        # Display molecule visualizations
+        for i, res in enumerate(results):
+            st.markdown(f"**Prediction {i+1}:** `{res['smiles']}` (Score: {res['score']})")
+            col1, col2 = st.columns(2)
+            with col1:
+                st.subheader("2D Structure")
+                if res["image_2d"]:
+                    st.image(res["image_2d"], use_column_width=True)
+                else:
+                    st.warning("Could not generate 2D image.")
+            with col2:
+                st.subheader("3D Interactive Structure")
+                if res["html_3d"]:
+                    components.html(res["html_3d"], height=370)
+                else:
+                    st.warning("Could not generate 3D view.")
+            st.markdown("---")
+
+# --- Tab 2: Molecule Viewer ---
+with tab2:
+    st.header("Visualize a Molecule from SMILES")
+    st.markdown("Enter a single SMILES string to display its 2D and 3D structures side-by-side.")
+
+    with st.form(key="viewer_form"):
+        smiles_input_viewer = st.text_input("SMILES String", value="CC(=O)Oc1ccccc1C(=O)O") # Aspirin
+        view_button = st.form_submit_button("View Molecule", use_container_width=True)
+
+    if view_button and smiles_input_viewer:
+        with st.spinner("Generating views..."):
+            mol = get_mol(smiles_input_viewer)
+            if not mol:
+                st.error("Invalid SMILES string provided.")
+            else:
+                st.subheader(f"Visualizations for: `{smiles_input_viewer}`")
+                col1, col2 = st.columns(2)
+                with col1:
+                    st.subheader("2D Structure")
+                    img_2d = MolToImage(mol, size=(450, 450), fitImage=True)
+                    st.image(img_2d, use_column_width=True)
+                with col2:
+                    st.subheader("3D Interactive Structure")
+                    html_3d = generate_3d_view_html(smiles_input_viewer)
+                    components.html(html_3d, height=470)
 
-if __name__ == "__main__":
-    main()