Update app.py

app.py

@@ -4,9 +4,8 @@ import base64
 from io import BytesIO
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 from rdkit import Chem
-from rdkit.Chem import AllChem
+from rdkit.Chem import Draw, AllChem
 import torch
-import py3Dmol
 
 # Load models
 bio_gpt = pipeline("text-generation", model="microsoft/BioGPT-Large")
@@ -14,26 +13,28 @@ chemberta_tokenizer = AutoTokenizer.from_pretrained("seyonec/ChemBERTa-zinc-base
 chemberta_model = AutoModelForCausalLM.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
 compliance_qa = pipeline("question-answering", model="nlpaueb/legal-bert-base-uncased")
 
-# Functions
-def extract_insights(prompt):
-    try:
-        result = bio_gpt(prompt, max_length=200, do_sample=True)
-        return result[0]['generated_text']
-    except Exception as e:
-        return f"Error: {str(e)}"
+# --- Functions ---
+def extract_insights(disease, symptoms):
+    prompt = f"Recent treatments for {disease} with symptoms: {symptoms}"
+    result = bio_gpt(prompt, max_length=200, do_sample=True)
+    return result[0]['generated_text']
 
 def generate_molecule():
     sample_smiles = ["CCO", "CCN", "C1=CC=CC=C1", "C(C(=O)O)N", "CC(C)CC"]
     return random.choice(sample_smiles)
 
 def predict_properties(smiles):
-    try:
-        inputs = chemberta_tokenizer(smiles, return_tensors="pt")
-        with torch.no_grad():
-            outputs = chemberta_model(**inputs)
-        return round(outputs.logits.mean().item(), 3)
-    except Exception as e:
-        return f"Error: {str(e)}"
+    inputs = chemberta_tokenizer(smiles, return_tensors="pt")
+    with torch.no_grad():
+        outputs = chemberta_model(**inputs)
+    return round(outputs.logits.mean().item(), 3)
+
+def visualize_2d(smiles):
+    mol = Chem.MolFromSmiles(smiles)
+    img = Draw.MolToImage(mol, size=(300, 300))
+    buf = BytesIO()
+    img.save(buf, format="PNG")
+    return base64.b64encode(buf.getvalue()).decode()
 
 def mol_to_3d_html(smiles):
     try:
@@ -41,53 +42,66 @@ def mol_to_3d_html(smiles):
         mol = Chem.AddHs(mol)
         AllChem.EmbedMolecule(mol, AllChem.ETKDG())
         AllChem.UFFOptimizeMolecule(mol)
-        block = Chem.MolToMolBlock(mol)
+        mol_block = Chem.MolToMolBlock(mol)
 
-        view = py3Dmol.view(width=400, height=400)
-        view.addModel(block, "mol")
-        view.setStyle({"stick": {}})
-        view.setBackgroundColor("white")
-        view.zoomTo()
-
-        return view._make_html()
+        encoded_block = mol_block.replace("\n", "\\n")
+        viewer_div = f"""
+        <div id="molviewer" style="width: 400px; height: 400px;"></div>
+        <script src="https://3Dmol.org/build/3Dmol-min.js"></script>
+        <script>
+        let element = document.getElementById("molviewer");
+        let config = {{ backgroundColor: "white" }};
+        let viewer = $3Dmol.createViewer(element, config);
+        let molData = `{encoded_block}`;
+        viewer.addModel(molData, "mol");
+        viewer.setStyle({{}}, {{stick:{{}}}});
+        viewer.zoomTo();
+        viewer.render();
+        </script>
+        """
+        return viewer_div
     except Exception as e:
-        return f"<p>Error generating 3D structure: {str(e)}</p>"
+        return f"<p>Error generating 3D molecule: {str(e)}</p>"
 
-def check_compliance(question, context):
-    try:
-        return compliance_qa(question=question, context=context)['answer']
-    except Exception as e:
-        return f"Error: {str(e)}"
+def check_compliance():
+    context = "FDA requires extensive testing for new drug candidates including Phase I, II, and III clinical trials."
+    question = "What does FDA require for drug testing?"
+    return compliance_qa(question=question, context=context)['answer']
 
-# Main function
-def run_pipeline(disease, symptoms):
-    insights = extract_insights(f"Recent treatments for {disease} with symptoms: {symptoms}")
+# --- Gradio UI ---
+def run_discovery(disease, symptoms):
+    insights = extract_insights(disease, symptoms)
     smiles = generate_molecule()
-    score = predict_properties(smiles)
-    mol3d_html = mol_to_3d_html(smiles)
-    compliance = check_compliance(
-        "What does FDA require for drug testing?",
-        "FDA requires extensive testing for new drug candidates including Phase I, II, and III clinical trials."
+    prop_score = predict_properties(smiles)
+    img_base64 = visualize_2d(smiles)
+    compliance = check_compliance()
+    html_3d = mol_to_3d_html(smiles)
+
+    return (
+        insights,
+        f"SMILES: {smiles}",
+        f"ChemBERTa Property Score: {prop_score}",
+        f"<img src='data:image/png;base64,{img_base64}' width='300'/>",
+        html_3d,
+        compliance
     )
-    return insights, smiles, mol3d_html, score, compliance
 
-# Gradio Interface
 demo = gr.Interface(
-    fn=run_pipeline,
+    fn=run_discovery,
     inputs=[
         gr.Textbox(label="🦠 Disease", value="lung cancer"),
         gr.Textbox(label="🩺 Symptoms", value="shortness of breath, weight loss")
     ],
     outputs=[
         gr.Textbox(label="📜 Literature Insights"),
-        gr.Textbox(label="🧪 Generated SMILES"),
-        gr.HTML(label="🧬 3D Molecule Viewer"),
-        gr.Textbox(label="📊 Molecular Property Score (ChemBERTa)"),
-        gr.Textbox(label="⚖️ Legal Compliance (FDA)")
+        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")
     ],
     title="🧬 AI-Driven Drug Discovery System",
-    description="Input a disease and its symptoms to discover potential drug candidates using LLMs and molecule modeling."
+    description="Enter disease and symptoms to generate drug candidates using BioGPT, ChemBERTa, and LegalBERT. View 2D and animated 3D molecules!"
 )
 
 demo.launch()
-