Upload 2 files

app.py

@@ -0,0 +1,122 @@
+import streamlit as st
+import torch
+import os
+from rdkit import Chem
+from rdkit.Chem import Draw
+from transformers import BartForConditionalGeneration, BartTokenizer
+from admet_ai import ADMETModel
+import safe
+import io
+from PIL import Image
+import cairosvg
+import pandas as pd
+
+# Page Configuration
+st.set_page_config(page_title='Beta-Lactam Molecule Generator', layout='wide')
+
+# Load Models
+@st.cache_resource(show_spinner="Loading Models...", ttl=600)
+def load_models():
+    # Load your molecule generation model
+    model_name = "your-new-beta-lactam-model-path"  # Replace with your actual model path
+    access_token = os.getenv("HUGGING_FACE_TOKEN")
+    model = BartForConditionalGeneration.from_pretrained(model_name, use_auth_token=access_token)
+    tokenizer = BartTokenizer.from_pretrained(model_name, use_auth_token=access_token)
+    # Load ADMET-AI model
+    admet_model = ADMETModel()
+    return model, tokenizer, admet_model
+
+model, tokenizer, admet_model = load_models()
+
+# Set Generation Parameters
+st.sidebar.header('Generation Parameters')
+creativity = st.sidebar.slider('Creativity (Temperature):', 0.0, 2.0, 1.0, step=0.1)
+num_molecules = st.sidebar.number_input('Number of Molecules to Generate:', min_value=1, max_value=5, value=5)
+
+# String Format Option
+string_format = st.sidebar.radio('String Format:', ('SMILES', 'SAFE'))
+
+# Generate Molecules Button
+if st.button('Generate Molecules'):
+    st.info("Generating molecules... Please wait.")
+    # Generate molecules
+    core_smiles = "C1C(=O)N(C)C(=O)C1"  # Beta-lactam core structure
+    output_ids = model.generate(
+        tokenizer(core_smiles, return_tensors='pt').input_ids,
+        max_length=128,
+        temperature=creativity,
+        do_sample=True,
+        top_k=50,
+        num_return_sequences=num_molecules
+    )
+    generated_smiles = [tokenizer.decode(ids, skip_special_tokens=True) for ids in output_ids]
+    molecule_names = [f"Mol{str(i).zfill(2)}" for i in range(1, num_molecules + 1)]
+    generated_molecules = dict(zip(molecule_names, generated_smiles))
+    
+    # ADMET Predictions
+    preds = admet_model.predict(smiles=list(generated_molecules.values()))
+    preds['Molecule Name'] = molecule_names
+    preds.set_index('Molecule Name', inplace=True)
+    
+    # Display Molecules
+    st.subheader('Generated Molecules')
+    cols_per_row = min(5, num_molecules)
+    cols = st.columns(cols_per_row)
+    for idx, mol_name in enumerate(molecule_names):
+        smiles = generated_molecules[mol_name]
+        img = generate_molecule_image(smiles, use_safe_visualization=(string_format == 'SAFE'))
+        with cols[idx % cols_per_row]:
+            if isinstance(img, Image.Image):
+                st.image(img, caption=mol_name)
+            else:
+                st.error(f"Could not generate image for {mol_name}")
+            # Display molecule string
+            string_to_display = safe.encode(smiles) if string_format == 'SAFE' else smiles
+            st.code(string_to_display)
+            # Copy-to-clipboard functionality
+            st_copy_button(string_to_display, key=f'copy_{mol_name}')
+            # Display ADMET properties
+            st.write("**ADMET Properties:**")
+            st.write(preds.loc[mol_name])
+else:
+    st.write("Click the 'Generate Molecules' button to generate beta-lactam molecules.")
+
+# Function Definitions
+def generate_molecule_image(input_string, use_safe_visualization=True):
+    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:
+                # If decoding fails, assume input is SMILES and encode to SAFE
+                safe_string = safe.encode(input_string)
+            # 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))
+        else:
+            # Generate standard molecule image
+            mol = Chem.MolFromSmiles(input_string)
+            if mol:
+                img = Draw.MolToImage(mol, size=(200, 200))  # Adjusted size
+            else:
+                img = None
+        return img
+    except Exception as e:
+        # Collect exceptions for later reporting
+        return e
+
+import streamlit.components.v1 as components
+
+def st_copy_button(text, key):
+    """Creates a copy-to-clipboard button."""
+    components.html(f"""
+        <button onclick="navigator.clipboard.writeText('{text}')" style="padding:5px;">Copy</button>
+    """, height=45)
+
+

requirements.txt

@@ -0,0 +1,10 @@
+streamlit
+rdkit-pypi
+numpy
+torch
+transformers
+admet-ai
+safe-encoding
+pandas
+cairosvg
+Pillow