Update app.py

app.py

@@ -1,40 +1,25 @@
 import gradio as gr
-import random
+from transformers import pipeline
+from rdkit import Chem
+from rdkit.Chem import AllChem, Draw
+from rdkit.Chem import Descriptors
 import base64
 from io import BytesIO
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
-from rdkit import Chem
-from rdkit.Chem import Draw, AllChem
-import torch
 
 # Load models
 bio_gpt = pipeline("text-generation", model="microsoft/BioGPT-Large")
-chemberta_tokenizer = AutoTokenizer.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
-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(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)
+chemberta = pipeline("text-classification", model="seyonec/ChemBERTa-zinc-base-chemprop")
+fda_gpt = pipeline("text-generation", model="EleutherAI/gpt-neo-1.3B")
 
-def predict_properties(smiles):
-    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):
+def mol_to_svg(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()
+    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()
 
 def mol_to_3d_html(smiles):
     try:
@@ -43,65 +28,59 @@ def mol_to_3d_html(smiles):
         AllChem.EmbedMolecule(mol, AllChem.ETKDG())
         AllChem.UFFOptimizeMolecule(mol)
         mol_block = Chem.MolToMolBlock(mol)
+        mol_block = mol_block.replace("\n", "\\n")
 
-        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>
+        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 element = document.getElementById("molviewer");
-        let config = {{ backgroundColor: "white" }};
-        let viewer = $3Dmol.createViewer(element, config);
-        let molData = `{encoded_block}`;
-        viewer.addModel(molData, "mol");
+        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 viewer_div
+        return html
     except Exception as e:
-        return f"<p>Error generating 3D molecule: {str(e)}</p>"
+        return f"<p>Error rendering 3D molecule: {str(e)}</p>"
 
-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']
+def drug_discovery(disease, symptoms):
+    prompt = f"Disease: {disease}\nSymptoms: {symptoms}\nLiterature Review:"
+    literature = bio_gpt(prompt, max_length=250)[0]['generated_text']
 
-# --- Gradio UI ---
-def run_discovery(disease, symptoms):
-    insights = extract_insights(disease, symptoms)
-    smiles = generate_molecule()
-    prop_score = predict_properties(smiles)
-    img_base64 = visualize_2d(smiles)
-    compliance = check_compliance()
+    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)
 
-    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
-    )
+    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']
+
+    return literature, smiles, f"ChemBERTa Property Score: {property_score:.3f}", svg, html_3d, fda_output
+
+inputs = [
+    gr.Textbox(label="🦠 Disease", placeholder="e.g. lung cancer"),
+    gr.Textbox(label="🧾 Symptoms", placeholder="e.g. shortness of breath, weight loss")
+]
+
+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")
+]
 
 demo = gr.Interface(
-    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="🧪 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="Enter disease and symptoms to generate drug candidates using BioGPT, ChemBERTa, and LegalBERT. View 2D and animated 3D molecules!"
+    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."
 )
 
 demo.launch()