Update app.py

app.py

@@ -1,4 +1,4 @@
-import streamlit as st
+import gradio as gr
 from transformers import pipeline
 from rdkit import Chem
 from rdkit.Chem import AllChem
@@ -8,59 +8,60 @@ import base64
 from io import BytesIO
 import py3Dmol
 
-st.set_page_config(page_title="Drug Discovery using LLMs", layout="wide")
+# Function to generate literature and 3D molecule view
+def drug_discovery(disease, symptoms):
+    # 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']
 
-st.title("🧬 Drug Discovery using LLMs")
+    # Generate SMILES (placeholder)
+    smiles = "CC(C)CC"  # Example: isopentane
 
-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")
+    # Generate RDKit molecule
+    mol = Chem.MolFromSmiles(smiles)
+    AllChem.Compute2DCoords(mol)
 
-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']
+    # Draw 2D image
+    drawer = rdMolDraw2D.MolDraw2DCairo(300, 300)
+    drawer.DrawMolecule(mol)
+    drawer.FinishDrawing()
+    img_data = drawer.GetDrawingText()
 
-        st.subheader("📚 Literature Insights")
-        st.write(literature)
+    # Convert binary to base64
+    img_base64 = base64.b64encode(img_data).decode("utf-8")
+    img_html = f'<img src="data:image/png;base64,{img_base64}" alt="2D Molecule">'
 
-        # 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}")
+    # 3D molecule
+    mol3d = Chem.AddHs(mol)
+    AllChem.EmbedMolecule(mol3d)
+    AllChem.UFFOptimizeMolecule(mol3d)
+    mb = Chem.MolToMolBlock(mol3d)
 
-        # 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)
+    viewer = py3Dmol.view(width=400, height=400)
+    viewer.addModel(mb, "mol")
+    viewer.setStyle({"stick": {}})
+    viewer.setBackgroundColor("white")
+    viewer.zoomTo()
+    viewer_html = viewer._make_html()
 
-        # Generate 3D coordinates
-        mol3d = Chem.AddHs(mol)
-        AllChem.EmbedMolecule(mol3d)
-        AllChem.UFFOptimizeMolecule(mol3d)
+    return literature, smiles, img_html, viewer_html
 
-        # Extract 3D coordinates for py3Dmol
-        mb = Chem.MolToMolBlock(mol3d)
+# Gradio UI
+disease_input = gr.Textbox(label="Enter Disease", value="lung cancer")
+symptom_input = gr.Textbox(label="Enter Symptoms", value="shortness of breath, weight loss")
+lit_output = gr.Textbox(label="Literature Insights")
+smiles_output = gr.Textbox(label="SMILES String")
+img_output = gr.HTML(label="2D Molecule")
+viewer_output = gr.HTML(label="3D Molecule Viewer")
 
-        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)
+iface = gr.Interface(
+    fn=drug_discovery,
+    inputs=[disease_input, symptom_input],
+    outputs=[lit_output, smiles_output, img_output, viewer_output],
+    title="🧬 Drug Discovery using LLMs",
+    description="Enter a disease and symptoms to generate literature insights and view the molecule in 2D and interactive 3D."
+)
 
-        st.success("✅ Done")
+iface.launch()
 
-else:
-    st.info("👈 Enter disease and symptoms, then hit Submit to discover a molecule.")