LLM / app.py
anaghanagesh's picture
Update app.py
82f018e verified
import gradio as gr
import base64
import html
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem.Draw import rdMolDraw2D
from agents import run_pipeline
# =========================================
# 2D + 3D MOLECULE GENERATION
# =========================================
def generate_molecule_visuals(smiles):
try:
# ==========================
# CREATE MOLECULE
# ==========================
mol = Chem.MolFromSmiles(smiles)
if mol is None:
return "", ""
# ==========================
# 2D MOLECULE
# ==========================
AllChem.Compute2DCoords(mol)
drawer = rdMolDraw2D.MolDraw2DCairo(
400,
400
)
drawer.DrawMolecule(mol)
drawer.FinishDrawing()
img_data = drawer.GetDrawingText()
img_base64 = base64.b64encode(
img_data
).decode("utf-8")
molecule_2d = f"""
<div style='text-align:center;'>
<img
src="data:image/png;base64,{img_base64}"
style="
width:320px;
border-radius:12px;
box-shadow:0 0 15px rgba(0,255,255,0.5);
"
/>
<h3 style='color:white;'>
🔬 2D Molecule Structure
</h3>
</div>
"""
# ==========================
# 3D MOLECULE
# ==========================
mol3d = Chem.AddHs(mol)
AllChem.EmbedMolecule(
mol3d
)
AllChem.UFFOptimizeMolecule(
mol3d
)
molblock = Chem.MolToMolBlock(
mol3d
)
escaped_molblock = html.escape(
molblock
)
viewer_html = f"""
<html>
<head>
<script src="https://3Dmol.org/build/3Dmol-min.js"></script>
<style>
body {{
margin:0;
overflow:hidden;
background:black;
}}
#viewer {{
width:100vw;
height:100vh;
}}
</style>
</head>
<body>
<div id="viewer"></div>
<script>
let viewer = $3Dmol.createViewer(
document.getElementById("viewer"),
{{
backgroundColor: "black"
}}
);
let moldata = `{escaped_molblock}`;
viewer.addModel(
moldata,
"mol"
);
viewer.setStyle(
{{}},
{{
stick: {{
colorscheme:
"cyanCarbon"
}},
sphere: {{
scale:0.3
}}
}}
);
viewer.zoomTo();
viewer.render();
viewer.spin(true);
</script>
</body>
</html>
"""
molecule_3d = f"""
<div style='text-align:center;'>
<iframe
srcdoc='{viewer_html}'
width='520'
height='520'
style='
border:none;
border-radius:12px;
overflow:hidden;
'>
</iframe>
<h3 style='color:white;'>
🧬 Animated 3D Molecule
</h3>
</div>
"""
return molecule_2d, molecule_3d
except Exception as e:
print("MOLECULE ERROR:", e)
return "", ""
# =========================================
# MAIN ANALYSIS FUNCTION
# =========================================
def analyze(symptoms, followup, audio):
try:
combined_input = symptoms
# FOLLOW-UP INPUT
if followup and followup.strip():
combined_input += (
" " + followup
)
# RUN PIPELINE
result = run_pipeline(
combined_input,
audio
)
response = result["response"]
pubmed = "\n".join(
result["pubmed"]
)
molecule_info = ""
molecule_2d = ""
molecule_3d = ""
# ==========================
# MOLECULE GENERATION
# ==========================
if result["molecules"]:
molecule = result[
"molecules"
][0]
molecule_info = f"""
Drug:
{molecule['name']}
SMILES:
{molecule['smiles']}
Molecular Weight:
{molecule['molecular_weight']}
"""
molecule_2d, molecule_3d = (
generate_molecule_visuals(
molecule["smiles"]
)
)
return (
response,
pubmed,
molecule_info,
molecule_2d,
molecule_3d
)
except Exception as e:
print("APP ERROR:", e)
return (
str(e),
"",
"",
"",
""
)
# =========================================
# CUSTOM CSS
# =========================================
custom_css = """
body {
background:
linear-gradient(
to right,
#141e30,
#243b55
);
color:white;
}
textarea {
font-size:16px !important;
}
"""
# =========================================
# UI
# =========================================
with gr.Blocks(
theme=gr.themes.Soft(),
css=custom_css
) as demo:
gr.Markdown(
"""
# 🧬 AI Biomedical Assistant
### Retrieval-Augmented Biomedical Intelligence System
"""
)
symptoms_input = gr.Textbox(
label="🩺 Symptoms",
lines=4,
placeholder="""
Examples:
fever, nausea, headache,
stomach pain, chest pain,
skin irritation...
"""
)
followup_input = gr.Textbox(
label="💬 Follow-Up Information",
lines=2
)
audio_input = gr.Audio(
sources=["microphone"],
type="filepath",
label="🎙️ Record Symptoms"
)
analyze_btn = gr.Button(
"🧠 Analyze"
)
diagnosis_output = gr.Markdown(
label="🧠 Biomedical Report"
)
pubmed_output = gr.Textbox(
label="📄 PubMed IDs",
lines=5
)
molecule_output = gr.Textbox(
label="🧪 Drug Information",
lines=8
)
molecule_2d_output = gr.HTML(
label="🔬 2D Molecule"
)
molecule_3d_output = gr.HTML(
label="🧬 Animated 3D Molecule"
)
analyze_btn.click(
fn=analyze,
inputs=[
symptoms_input,
followup_input,
audio_input
],
outputs=[
diagnosis_output,
pubmed_output,
molecule_output,
molecule_2d_output,
molecule_3d_output
]
)
demo.queue()
demo.launch()