Update app.py

app.py

@@ -1,86 +1,66 @@
-import gradio as gr
+import streamlit as st
 from transformers import pipeline
 from rdkit import Chem
-from rdkit.Chem import AllChem, Draw
-from rdkit.Chem import Descriptors
+from rdkit.Chem import AllChem
+from rdkit.Chem.Draw import rdMolDraw2D
+from rdkit.Chem import rdDepictor
 import base64
 from io import BytesIO
+import py3Dmol
 
-# Load models
-bio_gpt = pipeline("text-generation", model="microsoft/BioGPT-Large")
-chemberta = pipeline("text-classification", model="seyonec/ChemBERTa-zinc-base-chemprop")
-fda_gpt = pipeline("text-generation", model="EleutherAI/gpt-neo-1.3B")
+st.set_page_config(page_title="Drug Discovery using LLMs", layout="wide")
 
-def mol_to_svg(smiles):
-    mol = Chem.MolFromSmiles(smiles)
-    mol = Chem.AddHs(mol)
-    AllChem.EmbedMolecule(mol, AllChem.ETKDG())
-    AllChem.UFFOptimizeMolecule(mol)
-    drawer = Draw.MolDraw2DSVG(300, 300)
-    drawer.DrawMolecule(mol)
-    drawer.FinishDrawing()
-    return drawer.GetDrawingText()
+st.title("🧬 Drug Discovery using LLMs")
 
-def mol_to_3d_html(smiles):
-    try:
-        mol = Chem.MolFromSmiles(smiles)
-        mol = Chem.AddHs(mol)
-        AllChem.EmbedMolecule(mol, AllChem.ETKDG())
-        AllChem.UFFOptimizeMolecule(mol)
-        mol_block = Chem.MolToMolBlock(mol)
-        mol_block = mol_block.replace("\n", "\\n")
+with st.sidebar:
+    st.header("🦠 Disease")
+    disease = st.text_input("Enter disease", "lung cancer")
+    st.header("🧾 Symptoms")
+    symptoms = st.text_input("Enter symptoms", "shortness of breath, weight loss")
 
-        html = f"""
-        <div id=\"viewer\" style=\"width: 400px; height: 400px; position: relative;\"></div>
-        <script src=\"https://3Dmol.org/build/3Dmol-min.js\"></script>
-        <script>
-        let viewer = $3Dmol.createViewer("viewer", {{ backgroundColor: "white" }});
-        let molBlock = `{mol_block}`;
-        viewer.addModel(molBlock, "mol");
-        viewer.setStyle({{}}, {{stick:{{}}}});
-        viewer.zoomTo();
-        viewer.animate({{loop: "backAndForth"}});
-        viewer.render();
-        </script>
-        """
-        return html
-    except Exception as e:
-        return f"<p>Error rendering 3D molecule: {str(e)}</p>"
+if st.button("Submit"):
+    with st.spinner("🔎 Discovering insights and potential drugs..."):
+        # BioGPT pipeline
+        bio_gpt = pipeline("text-generation", model="microsoft/BioGPT-Large")
+        prompt = f"Recent treatments for {disease} with symptoms: {symptoms}."
+        literature = bio_gpt(prompt, max_length=200)[0]['generated_text']
 
-def drug_discovery(disease, symptoms):
-    prompt = f"Disease: {disease}\nSymptoms: {symptoms}\nLiterature Review:"
-    literature = bio_gpt(prompt, max_length=250)[0]['generated_text']
+        st.subheader("📚 Literature Insights")
+        st.write(literature)
 
-    smiles = "CC(C)CC"  # Just an example SMILES string, normally generated by a molecule generator
-    property_score = chemberta(smiles)[0]['score']
-    svg = mol_to_svg(smiles)
-    html_3d = mol_to_3d_html(smiles)
+        # Generate a dummy SMILES for demo (could replace with MoLeR/Diffusion)
+        st.subheader("🧪 SMILES String")
+        smiles = "CC(C)CC"  # Isopentane (dummy drug structure)
+        st.write(f"SMILES: {smiles}")
 
-    fda_prompt = f"A drug with SMILES {smiles} is under review. List FDA regulatory issues."
-    fda_output = fda_gpt(fda_prompt, max_length=100)[0]['generated_text']
+        # Draw 2D molecule (optional, not shown if 3D works)
+        mol = Chem.MolFromSmiles(smiles)
+        AllChem.Compute2DCoords(mol)
+        rdDepictor.Compute2DCoords(mol)
+        drawer = rdMolDraw2D.MolDraw2DCairo(300, 300)
+        drawer.DrawMolecule(mol)
+        drawer.FinishDrawing()
+        img_data = drawer.GetDrawingText()
+        st.image(img_data, caption="2D Structure", use_column_width=False)
 
-    return literature, smiles, f"ChemBERTa Property Score: {property_score:.3f}", svg, html_3d, fda_output
+        # Generate 3D coordinates
+        mol3d = Chem.AddHs(mol)
+        AllChem.EmbedMolecule(mol3d)
+        AllChem.UFFOptimizeMolecule(mol3d)
 
-inputs = [
-    gr.Textbox(label="🦠 Disease", placeholder="e.g. lung cancer"),
-    gr.Textbox(label="🧾 Symptoms", placeholder="e.g. shortness of breath, weight loss")
-]
+        # Extract 3D coordinates for py3Dmol
+        mb = Chem.MolToMolBlock(mol3d)
 
-outputs = [
-    gr.Textbox(label="📜 Literature Insights"),
-    gr.Textbox(label="🧪 SMILES String"),
-    gr.Textbox(label="🧬 Property Score"),
-    gr.HTML(label="🧫 2D Molecule Structure"),
-    gr.HTML(label="🔬 3D Molecule Viewer"),
-    gr.Textbox(label="⚖️ FDA Compliance Summary")
-]
+        st.subheader("🔬 3D Structure Viewer")
+        xyzview = py3Dmol.view(width=400, height=400)
+        xyzview.addModel(mb, "mol")
+        xyzview.setStyle({"stick": {}})
+        xyzview.setBackgroundColor("white")
+        xyzview.zoomTo()
+        xyzview.show()
+        st.components.v1.html(xyzview._make_html(), height=400)
 
-demo = gr.Interface(
-    fn=drug_discovery,
-    inputs=inputs,
-    outputs=outputs,
-    title="💊 AI-Driven Drug Discovery using LLMs",
-    description="Enter a disease and its symptoms. This app generates literature insights, a possible molecule in SMILES format, scores its drug-likeness, and shows 2D & 3D views."
-)
+        st.success("✅ Done")
 
-demo.launch()
+else:
+    st.info("👈 Enter disease and symptoms, then hit Submit to discover a molecule.")