Add Gradio UI with single and batch prediction

app.py

@@ -5,23 +5,33 @@ from typing import List, Dict, Any
 ROOT = Path(__file__).parent
 sys.path.insert(0, str(ROOT))
 
+import gradio as gr
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel, Field
 import numpy as np
 
 from src import EnhancedFeatureExtractor, Tox21Ensemble
 
-app = FastAPI(
-    title="Rasayan Tox21 Classifier",
-    description="Self-Normalizing Neural Network ensemble for Tox21 toxicity prediction",
-    version="1.0.0"
-)
-
 TASKS = [
     "NR-AR", "NR-AR-LBD", "NR-AhR", "NR-Aromatase", "NR-ER", "NR-ER-LBD",
     "NR-PPAR-gamma", "SR-ARE", "SR-ATAD5", "SR-HSE", "SR-MMP", "SR-p53"
 ]
 
+TASK_DESCRIPTIONS = {
+    "NR-AR": "Androgen Receptor",
+    "NR-AR-LBD": "Androgen Receptor LBD",
+    "NR-AhR": "Aryl Hydrocarbon Receptor",
+    "NR-Aromatase": "Aromatase (CYP19A1)",
+    "NR-ER": "Estrogen Receptor",
+    "NR-ER-LBD": "Estrogen Receptor LBD",
+    "NR-PPAR-gamma": "PPARγ",
+    "SR-ARE": "Antioxidant Response",
+    "SR-ATAD5": "DNA Damage (ATAD5)",
+    "SR-HSE": "Heat Shock Response",
+    "SR-MMP": "Mitochondrial Toxicity",
+    "SR-p53": "Genotoxicity (p53)"
+}
+
 FEATURE_KEYS = [
     "ecfps", "maccs", "rdkit_descrs", "tox", "rdkit_filters",
     "similarity", "max_similarity", "db_similarity"
@@ -38,6 +48,239 @@ ensemble = Tox21Ensemble(ROOT / "checkpoints" / "ensemble.pt")
 print("Model loaded successfully!")
 
 
+def predict_toxicity(smiles_input: str) -> tuple:
+    if not smiles_input.strip():
+        return None, "Please enter at least one SMILES"
+
+    lines = [s.strip() for s in smiles_input.strip().split('\n') if s.strip()]
+
+    if len(lines) > 100:
+        return None, "Maximum 100 molecules per request"
+
+    try:
+        features_dict, valid = extractor.extract_features(lines)
+
+        features = np.concatenate(
+            [features_dict[k] for k in FEATURE_KEYS if k in features_dict],
+            axis=1
+        )
+        features = np.nan_to_num(features, nan=0.0, posinf=0.0, neginf=0.0)
+
+        probs = ensemble.predict(features)
+
+        results = []
+        for i, smi in enumerate(lines):
+            if valid[i]:
+                row = {"SMILES": smi[:50] + "..." if len(smi) > 50 else smi}
+                for j, task in enumerate(TASKS):
+                    score = float(probs[i, j])
+                    row[task] = f"{score:.1%}"
+                results.append(row)
+            else:
+                row = {"SMILES": smi[:50] + "..." if len(smi) > 50 else smi}
+                for task in TASKS:
+                    row[task] = "Invalid"
+                results.append(row)
+
+        import pandas as pd
+        df = pd.DataFrame(results)
+
+        return df, f"Processed {len(lines)} molecule(s)"
+
+    except Exception as e:
+        return None, f"Error: {str(e)}"
+
+
+def predict_single(smiles: str) -> str:
+    if not smiles.strip():
+        return "Enter a SMILES string"
+
+    try:
+        features_dict, valid = extractor.extract_features([smiles])
+
+        if not valid[0]:
+            return "Invalid SMILES structure"
+
+        features = np.concatenate(
+            [features_dict[k] for k in FEATURE_KEYS if k in features_dict],
+            axis=1
+        )
+        features = np.nan_to_num(features, nan=0.0, posinf=0.0, neginf=0.0)
+
+        probs = ensemble.predict(features)
+
+        lines = []
+        lines.append("═" * 45)
+        lines.append("  TOXICITY PREDICTION RESULTS")
+        lines.append("═" * 45)
+
+        sorted_results = sorted(
+            [(task, float(probs[0, j])) for j, task in enumerate(TASKS)],
+            key=lambda x: -x[1]
+        )
+
+        for task, score in sorted_results:
+            desc = TASK_DESCRIPTIONS[task]
+            bar_len = int(score * 20)
+            bar = "█" * bar_len + "░" * (20 - bar_len)
+
+            if score >= 0.7:
+                risk = "HIGH"
+            elif score >= 0.4:
+                risk = "MED "
+            elif score >= 0.2:
+                risk = "LOW "
+            else:
+                risk = "MIN "
+
+            lines.append(f"{task:15} {bar} {score:5.1%} [{risk}]")
+            lines.append(f"  └─ {desc}")
+
+        lines.append("═" * 45)
+
+        return "\n".join(lines)
+
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+
+EXAMPLES = [
+    ["CCO"],
+    ["CC(=O)Nc1ccc(O)cc1"],
+    ["c1ccc2c(c1)cc3ccc4cccc5ccc2c3c45"],
+    ["CC12CCC3C(C1CCC2O)CCC4=CC(=O)CCC34C"],
+    ["CC12CCC3c4ccc(O)cc4CCC3C1CCC2O"],
+]
+
+with gr.Blocks(
+    title="Rasayan Tox21 Classifier",
+    theme=gr.themes.Soft()
+) as demo:
+    gr.Markdown("""
+# ☠️ Rasayan Tox21 Classifier
+
+Predict molecular toxicity across **12 Tox21 endpoints** using a Self-Normalizing Neural Network ensemble.
+
+| Model | Features | Training |
+|-------|----------|----------|
+| 10-fold SNN Ensemble | 11,377 molecular descriptors | 40-fold CV, AUC: 0.882 |
+    """)
+
+    with gr.Tabs():
+        with gr.TabItem("Single Molecule"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    single_input = gr.Textbox(
+                        label="SMILES",
+                        placeholder="Enter SMILES (e.g., CCO for ethanol)",
+                        lines=1
+                    )
+                    single_btn = gr.Button("Predict", variant="primary")
+                    gr.Examples(
+                        examples=EXAMPLES,
+                        inputs=single_input,
+                        label="Example Molecules"
+                    )
+
+                with gr.Column(scale=2):
+                    single_output = gr.Textbox(
+                        label="Toxicity Profile",
+                        lines=30,
+                        show_copy_button=True
+                    )
+
+            single_btn.click(
+                fn=predict_single,
+                inputs=single_input,
+                outputs=single_output
+            )
+
+        with gr.TabItem("Batch Processing"):
+            gr.Markdown("Enter multiple SMILES (one per line, max 100)")
+
+            batch_input = gr.Textbox(
+                label="SMILES List",
+                placeholder="CCO\nCC(=O)Nc1ccc(O)cc1\nc1ccccc1",
+                lines=5
+            )
+            batch_btn = gr.Button("Process Batch", variant="primary")
+            batch_status = gr.Textbox(label="Status", lines=1)
+            batch_output = gr.Dataframe(
+                label="Results",
+                wrap=True
+            )
+
+            batch_btn.click(
+                fn=predict_toxicity,
+                inputs=batch_input,
+                outputs=[batch_output, batch_status]
+            )
+
+        with gr.TabItem("About"):
+            gr.Markdown("""
+## Model Architecture
+
+**Self-Normalizing Neural Networks (SNNs)** with SELU activation and AlphaDropout.
+
+| Component | Details |
+|-----------|---------|
+| Hidden Layers | 8 × 768 units |
+| Activation | SELU |
+| Dropout | AlphaDropout (0.1) |
+| Ensemble | Top-10 from 40-fold CV |
+| Parameters | ~19M total |
+
+## Molecular Features (11,377 total)
+
+| Feature | Dimensions | Description |
+|---------|------------|-------------|
+| ECFP6 | 8,192 | Morgan fingerprints (radius 3) |
+| MACCS | 167 | Structural keys |
+| RDKit | 208 | Physicochemical descriptors |
+| Toxicophores | 1,868 | Toxicity structural alerts |
+| Filters | 815 | PAINS, BRENK, NIH, ZINC |
+| Similarity | 127 | Target ligand similarity |
+
+## Tox21 Endpoints
+
+### Nuclear Receptor Panel
+- **NR-AR**: Androgen Receptor
+- **NR-AR-LBD**: AR Ligand Binding Domain
+- **NR-AhR**: Aryl Hydrocarbon Receptor
+- **NR-Aromatase**: CYP19A1 Enzyme
+- **NR-ER**: Estrogen Receptor
+- **NR-ER-LBD**: ER Ligand Binding Domain
+- **NR-PPAR-gamma**: Peroxisome Proliferator-Activated Receptor
+
+### Stress Response Panel
+- **SR-ARE**: Antioxidant Response Element
+- **SR-ATAD5**: DNA Damage Response
+- **SR-HSE**: Heat Shock Element
+- **SR-MMP**: Mitochondrial Membrane Potential
+- **SR-p53**: Tumor Suppressor p53
+
+## Risk Interpretation
+
+| Score | Risk Level |
+|-------|------------|
+| < 20% | Minimal |
+| 20-40% | Low |
+| 40-70% | Moderate |
+| ≥ 70% | High |
+
+---
+
+Built by [Rasayan Labs](https://rasayan.ai)
+            """)
+
+    gr.Markdown("""
+---
+**API Endpoints**: `/predict` (POST), `/metadata` (GET), `/health` (GET)
+    """)
+
+app = FastAPI()
+
+
 class PredictRequest(BaseModel):
     smiles: List[str] = Field(..., min_length=1, max_length=1000)
 
@@ -113,6 +356,8 @@ def health():
     return {"status": "ok"}
 
 
+app = gr.mount_gradio_app(app, demo, path="/")
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

@@ -5,3 +5,5 @@ numpy>=1.24.0
 torch>=2.0.0
 rdkit>=2023.3.1
 scikit-learn>=1.3.0
+gradio>=4.0.0
+pandas>=2.0.0