Update app.py

app.py

@@ -1,101 +1,126 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForMaskedLM, AutoModelForCausalLM
 from rdkit import Chem
 from rdkit.Chem import AllChem
 from rdkit.Chem.Draw import rdMolDraw2D
 import base64
 import re
 import py3Dmol
-import time
+import random
 
-# Load model once
-bio_gpt = pipeline("text-generation", model="microsoft/BioGPT-Large")
+# Load multiple models
+bio_gpt = pipeline("text-generation", model="microsoft/BioGPT-Large", device=0)
+pubmed_bert = pipeline("fill-mask", model="microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext")
 
+# Helper function to clean and parse generated SMILES
+def extract_valid_smiles(generated_text, count=3):
+    smiles_matches = re.findall(r"(?<![A-Za-z0-9])[A-Za-z0-9@+\-\[\]\(\)=#$]{5,}(?![A-Za-z0-9])", generated_text)
+    valid_smiles = []
+    for match in smiles_matches:
+        mol = Chem.MolFromSmiles(match)
+        if mol:
+            valid_smiles.append(match)
+        if len(valid_smiles) >= count:
+            break
+    return valid_smiles
+
+# Drug discovery function
 def drug_discovery(disease, symptoms):
-    # Simplified and efficient medical prompt
+    # Prompt to BioGPT for biomedical insights
     prompt = (
-        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 ..."
+        f"Imagine a novel disease '{disease}' with symptoms: {symptoms}.\n"
+        f"Write a short research summary covering:\n"
+        f"- Hypothetical causes\n- Suggested diagnostic methods\n- Possible treatments\n"
+        f"- Names of potential experimental drugs\n"
     )
 
     try:
-        start = time.time()
-        result = bio_gpt(prompt, max_length=150, do_sample=True, temperature=0.6)[0]['generated_text']
-    except Exception as e:
-        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)
+        literature = bio_gpt(prompt, max_length=512, temperature=0.7)[0]["generated_text"]
+    except:
+        literature = "⚠️ Error: Could not retrieve literature using BioGPT."
 
-    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"
+    # Prompt for SMILES + Drug Names using BioGPT
+    molecule_prompt = (
+        f"Give 3 unique experimental drug-like SMILES strings with hypothetical drug names for treating '{disease}' "
+        f"with symptoms: {symptoms}. Format: <SMILES> - <DrugName>"
+    )
 
-    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]
+    try:
+        smiles_response = bio_gpt(molecule_prompt, max_length=100)[0]["generated_text"]
+    except:
+        smiles_response = "C1=CC=CC=C1 - DemoDrug"
 
-    if not smiles_list:
-        smiles_list = ["C1=CC=CC=C1"]
+    entries = re.findall(r"([A-Za-z0-9@+\-\[\]\(\)=#$]{5,})\s*-\s*(\w+)", smiles_response)
+    results = []
 
-    img_html, viewer_htmls = "", ""
-    for smiles in smiles_list:
+    for smiles, name in entries[:3]:  # Limit to 3 molecules
         mol = Chem.MolFromSmiles(smiles)
+        if not mol:
+            continue
+
+        # 2D drawing
         AllChem.Compute2DCoords(mol)
-        drawer = rdMolDraw2D.MolDraw2DCairo(250, 250)
+        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="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>'''
+        img_html = f'''<img src="data:image/png;base64,{img_base64}" width="200" style="margin:10px; border:1px solid #ccc; border-radius:12px;">'''
 
-        # 3D View
+        # 3D molecule view
         mol3d = Chem.AddHs(mol)
-        AllChem.EmbedMolecule(mol3d, randomSeed=42)
+        AllChem.EmbedMolecule(mol3d)
         AllChem.UFFOptimizeMolecule(mol3d)
         mb = Chem.MolToMolBlock(mol3d)
-        viewer = py3Dmol.view(width=240, height=240)
+
+        viewer = py3Dmol.view(width=300, height=300)
         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>'''
+        viewer_html = f'''<iframe srcdoc="{viewer_html_raw.replace('"', '&quot;')}" 
+                          width="320" height="320" frameborder="0" 
+                          style="border-radius: 12px; box-shadow: 0 6px 20px rgba(0,255,255,0.35);"></iframe>'''
+
+        results.append((f"{name}", smiles, img_html, viewer_html))
 
-    duration = round(time.time() - start, 2)
-    literature_summary = f"📋 Drug candidates (auto-generated in {duration}s):\n{drug_names}"
+    # If no valid molecule found
+    if not results:
+        results.append(("DemoDrug", "C1=CC=CC=C1", "", ""))
 
-    return literature_summary, ", ".join(smiles_list), img_html, viewer_htmls
+    # Combine outputs
+    combined_molecules = ""
+    combined_imgs = ""
+    combined_3d = ""
+    for drug_name, smiles, img, viewer in results:
+        combined_molecules += f"💊 {drug_name}: {smiles}\n"
+        combined_imgs += img
+        combined_3d += viewer
 
-# Gradio UI setup
+    return literature.strip(), combined_molecules.strip(), combined_imgs, combined_3d
 
+# Gradio UI
 iface = gr.Interface(
     fn=drug_discovery,
     inputs=[
-        gr.Textbox(label="🧬 Enter Unknown Disease or Name", value="X-disease"),
-        gr.Textbox(label="📝 Symptoms", value="fever, joint pain")
+        gr.Textbox(label="🦠 Enter Novel Disease Name", value="Neurospike Fever"),
+        gr.Textbox(label="🩺 Enter Symptoms", value="fever, neural tremors, fatigue"),
     ],
     outputs=[
-        gr.Textbox(label="🔖 AI Literature Summary"),
-        gr.Textbox(label="🧪 SMILES List"),
+        gr.Textbox(label="📖 Biomedical Insights"),
+        gr.Textbox(label="🧪 Molecule Names + SMILES"),
         gr.HTML(label="🖼️ 2D Molecules"),
-        gr.HTML(label="🔬 3D 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.",
+    title="🧬 Experimental Drug Discovery for Unknown Diseases",
+    description="Enter an unknown disease and symptoms. The AI (BioGPT + PubMedBERT) will generate biomedical research, name possible experimental drugs, and visualize 2D + 3D molecular structures.",
     css="""
-        body { background-color: #111; color: #eee; }
-        .gradio-container { animation: fadeIn 1.5s ease-in-out; }
-    """
+    body { background-color: #0d1117; color: #ffffff; font-family: Segoe UI, sans-serif; }
+    .gradio-container { padding: 20px; }
+    iframe, img { margin: 8px; }
+    """,
+    allow_flagging="never"
 )
 
 iface.launch(share=True)