Update app.py

app.py

@@ -6,131 +6,96 @@ from rdkit.Chem.Draw import rdMolDraw2D
 import base64
 import re
 import py3Dmol
+import time
 
-# Drug discovery function
+# Load model once
+bio_gpt = pipeline("text-generation", model="microsoft/BioGPT-Large")
 
 def drug_discovery(disease, symptoms):
-    bio_gpt = pipeline("text-generation", model="microsoft/BioGPT-Large")
-
-    # Detailed medical prompt
+    # Simplified and efficient medical prompt
     prompt = (
-        f"Act as a biomedical researcher. For the disease '{disease}' with symptoms '{symptoms}', provide a detailed summary of:\n"
-        "- Causes\n- Diagnosis methods\n- Treatment options\n- Common medications (include drug names)\n"
-        "- Any FDA-approved therapies or hospital protocols\n\nKeep it concise but detailed."
+        f"You're a biomedical AI. A new disease shows symptoms: '{symptoms}'. "
+        f"Suggest 5 generic drug names and 5 SMILES strings that could help treat this. "
+        f"List drug names first, then SMILES strings in separate lines like:\n"
+        f"Drugs: Aspirin, Ibuprofen, Paracetamol, ...\n"
+        f"SMILES: C1=CC=CC=C1 C(C(=O)O)N ..."
     )
 
     try:
-        result = bio_gpt(prompt, max_length=512, do_sample=True, temperature=0.7)[0]['generated_text']
-    except Exception as e:
-        result = f"Could not generate literature due to an error: {e}"
-
-    result = re.sub(r"<\s*/?\s*(TITLE|FREETEXT)\s*>", "", result)
-    result = re.sub(r"^.*?(?=Causes|Diagnosis|Treatment|Common medications)", "", result, flags=re.IGNORECASE)
-
-    # Generate SMILES
-    molecule_prompt = f"Give 5 different valid drug-like SMILES strings that can treat {disease} with symptoms: {symptoms}. Only list SMILES separated by spaces."
-    try:
-        smiles_result = bio_gpt(molecule_prompt, max_length=100)[0]['generated_text']
+        start = time.time()
+        result = bio_gpt(prompt, max_length=150, do_sample=True, temperature=0.6)[0]['generated_text']
     except Exception as e:
-        smiles_result = "C1=CC=CC=C1"
-
-    smiles_matches = re.findall(r"(?<![A-Za-z0-9])[A-Za-z0-9@+\-\[\]\(\)=#$]{5,}(?![A-Za-z0-9])", smiles_result)
-    smiles = None
-    for match in smiles_matches:
-        mol_test = Chem.MolFromSmiles(match)
-        if mol_test:
-            smiles = match
-            break
-    if not smiles:
-        smiles = "C1=CC=CC=C1"
-
-    mol = Chem.MolFromSmiles(smiles)
-    if not mol:
-        return "Invalid SMILES generated", smiles, "", ""
-
-    AllChem.Compute2DCoords(mol)
-    drawer = rdMolDraw2D.MolDraw2DCairo(300, 300)
-    drawer.DrawMolecule(mol)
-    drawer.FinishDrawing()
-    img_data = drawer.GetDrawingText()
-    img_base64 = base64.b64encode(img_data).decode("utf-8")
-    img_html = f'''<div style="text-align:center; margin-top: 10px; animation: fadeIn 2s ease-in-out;">
-        <img src="data:image/png;base64,{img_base64}" alt="2D Molecule" 
-             style="border-radius: 16px; box-shadow: 0 6px 20px rgba(0,255,255,0.3); border: 1px solid #444;">
-        <div style='font-family: Arial, sans-serif; color: #eeeeee; margin-top: 8px; animation: slideUp 1.5s ease-in-out;'>💊 Visualized Drug Molecule (2D)</div>
-    </div>'''
-
-    mol3d = Chem.AddHs(mol)
-    AllChem.EmbedMolecule(mol3d)
-    AllChem.UFFOptimizeMolecule(mol3d)
-    mb = Chem.MolToMolBlock(mol3d)
-
-    viewer = py3Dmol.view(width=420, height=420)
-    viewer.addModel(mb, "mol")
-    viewer.setStyle({"stick": {"colorscheme": "cyanCarbon"}})
-    viewer.setBackgroundColor("black")
-    viewer.zoomTo()
-    viewer.spin(True)
-    viewer_html_raw = viewer._make_html()
-
-    viewer_html = f'''
-    <div style="text-align:center; margin-top: 20px; animation: zoomIn 2s ease-in-out;">
-        <iframe srcdoc="{viewer_html_raw.replace('"', '&quot;')}" 
-                width="440" height="440" frameborder="0" 
-                style="border-radius: 16px; box-shadow: 0 8px 30px rgba(0,255,255,0.35);"></iframe>
-        <div style='font-family: Arial, sans-serif; color: #eeeeee; margin-top: 8px; animation: slideUp 1.5s ease-in-out;'>🧬 Animated 3D Molecule (Stick View)</div>
-    </div>'''
-
-    return result.strip(), smiles, img_html, viewer_html
-
-# Gradio UI
-
-disease_input = gr.Textbox(label="🏥 Enter Disease (e.g., lung cancer)", value="lung cancer")
-symptom_input = gr.Textbox(label="📈 Enter Symptoms (e.g., cough, weight loss)", value="shortness of breath, weight loss")
-lit_output = gr.Textbox(label="🔖 Literature Insights from BioGPT")
-smiles_output = gr.Textbox(label="🧪 SMILES Representation")
-img_output = gr.HTML(label="🖼️ Molecule 2D Visualization")
-viewer_output = gr.HTML(label="🔬 3D Drug Molecule Animation")
-
-custom_css = """
-@keyframes fadeIn {
-    from {opacity: 0;}
-    to {opacity: 1;}
-}
-
-@keyframes slideUp {
-    from {transform: translateY(40px); opacity: 0;}
-    to {transform: translateY(0); opacity: 1;}
-}
-
-@keyframes zoomIn {
-    from {transform: scale(0.5); opacity: 0;}
-    to {transform: scale(1); opacity: 1;}
-}
-
-body {
-    background: linear-gradient(to right, #0f0f0f, #1a1a1a, #000000);
-    color: #eeeeee;
-    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
-}
-
-.gradio-container {
-    animation: fadeIn 1.5s ease-in-out;
-}
-
-.gradio-container .block-label {
-    color: #ffffff;
-}
-"""
+        return f"BioGPT error: {e}", "", "", ""
+
+    # Extract drug names and SMILES
+    drugs_match = re.search(r"Drugs:\s*(.+)", result)
+    smiles_match = re.search(r"SMILES:\s*(.+)", result)
+
+    drug_names = drugs_match.group(1).strip() if drugs_match else "Unknown"
+    raw_smiles = smiles_match.group(1).strip() if smiles_match else "C1=CC=CC=C1"
+
+    smiles_list = re.findall(r"(?<![A-Za-z0-9])[A-Za-z0-9@+\-\[\]\(\)=#$]{5,}(?![A-Za-z0-9])", raw_smiles)
+    smiles_list = list({sm for sm in smiles_list if Chem.MolFromSmiles(sm)})[:3]
+
+    if not smiles_list:
+        smiles_list = ["C1=CC=CC=C1"]
+
+    img_html, viewer_htmls = "", ""
+    for smiles in smiles_list:
+        mol = Chem.MolFromSmiles(smiles)
+        AllChem.Compute2DCoords(mol)
+        drawer = rdMolDraw2D.MolDraw2DCairo(250, 250)
+        drawer.DrawMolecule(mol)
+        drawer.FinishDrawing()
+        img_data = drawer.GetDrawingText()
+        img_base64 = base64.b64encode(img_data).decode("utf-8")
+        img_html += f'''
+        <div style="display:inline-block; margin:10px;">
+            <img src="data:image/png;base64,{img_base64}" width="120" height="120">
+            <p style="color:white; font-size:12px;">{smiles}</p>
+        </div>'''
+
+        # 3D View
+        mol3d = Chem.AddHs(mol)
+        AllChem.EmbedMolecule(mol3d, randomSeed=42)
+        AllChem.UFFOptimizeMolecule(mol3d)
+        mb = Chem.MolToMolBlock(mol3d)
+        viewer = py3Dmol.view(width=240, height=240)
+        viewer.addModel(mb, "mol")
+        viewer.setStyle({"stick": {"colorscheme": "cyanCarbon"}})
+        viewer.setBackgroundColor("black")
+        viewer.zoomTo()
+        viewer_html_raw = viewer._make_html()
+        viewer_htmls += f'''
+        <div style="display:inline-block; margin:10px;">
+            <iframe srcdoc="{viewer_html_raw.replace('"', '&quot;')}" width="240" height="240" frameborder="0"></iframe>
+        </div>'''
+
+    duration = round(time.time() - start, 2)
+    literature_summary = f"📋 Drug candidates (auto-generated in {duration}s):\n{drug_names}"
+
+    return literature_summary, ", ".join(smiles_list), img_html, viewer_htmls
+
+# Gradio UI setup
 
 iface = gr.Interface(
     fn=drug_discovery,
-    inputs=[disease_input, symptom_input],
-    outputs=[lit_output, smiles_output, img_output, viewer_output],
-    title="🏥 AI-Powered Drug Discovery for Hospitals",
-    description="This hospital-themed platform takes a disease and symptoms as input, retrieves biomedical insights using BioGPT, and visualizes potential drug molecules in 2D and animated 3D. Ideal for clinical research and pharma innovation.",
-    theme="default",
-    css=custom_css
+    inputs=[
+        gr.Textbox(label="🧬 Enter Unknown Disease or Name", value="X-disease"),
+        gr.Textbox(label="📝 Symptoms", value="fever, joint pain")
+    ],
+    outputs=[
+        gr.Textbox(label="🔖 AI Literature Summary"),
+        gr.Textbox(label="🧪 SMILES List"),
+        gr.HTML(label="🖼️ 2D Molecules"),
+        gr.HTML(label="🔬 3D Molecules")
+    ],
+    title="🧪 Drug Discovery for Unknown Diseases",
+    description="BioGPT + RDKit-powered system to suggest potential drug molecules for unknown or rare diseases.",
+    css="""
+        body { background-color: #111; color: #eee; }
+        .gradio-container { animation: fadeIn 1.5s ease-in-out; }
+    """
 )
 
-iface.launch(share=True)
+iface.launch(share=True)