Blame for app.py

import os

os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import pandas as pd

try:
    import spaces
except ImportError:
    spaces = None

import csv
import html
import json
import logging
import math
import shutil
import tempfile
import time
from pathlib import Path
from threading import Lock

import gradio as gr
import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np
import plotly.graph_objects as go
from rdkit import Chem
from rdkit.Chem import Draw


ROOT = Path(__file__).resolve().parent
ASSETS = Path(os.getenv("SMART_ASSETS", str(ROOT / "assets"))).resolve()
EXPORT_ROOT = Path(tempfile.mkdtemp(prefix="smart_exports_"))
ZERO_GPU = os.getenv("SPACES_ZERO_GPU") == "1" or os.getenv("SMART_ZEROGPU") == "1"
DEVICE = os.getenv("SMART_DEVICE", "cuda" if ZERO_GPU else "cpu").lower()
LOCK = Lock()
EXPORT_LOCK = Lock()
PREDICTOR = None
LOAD_ERROR = None
LOG = logging.getLogger("smart")

TEXT = {'intro': 'Enter one molecular SMILES to predict its electron ionization (EI) mass spectrum for GC–MS.',
 'input': 'Molecular SMILES',
 'placeholder': 'Enter one molecular SMILES or choose an example below',
 'predict': 'Predict spectrum',
 'clear': 'Clear',
 'examples': 'Example molecules',
 'caffeine': 'Caffeine',
 'ketamine': 'Ketamine',
 'lidocaine': 'Lidocaine',
 'structure': 'Molecular structure',
 'spectrum': 'Predicted EI spectrum',
 'intensity': 'Relative intensity (%)',
 'mass': 'Monoisotopic mass',
 'formula': 'Molecular formula',
 'reference': 'Usable reference spectra',
 'csv': 'Download CSV',
 'msp': 'Download MSP',
 'png': 'Download spectrum PNG',
 'ready': 'Model ready · Single-molecule prediction',
 'missing': 'Model assets are not ready. The maintainer must upload the packaged assets folder and restart '
            'this Space.',
 'empty': 'Please enter a SMILES first.',
 'long': 'Please enter one molecular SMILES with no more than 1000 characters.',
 'invalid': 'This SMILES could not be parsed. Please check the structure expression.',
 'large': 'This molecule exceeds the demo limit of 200 heavy atoms.',
 'components': 'Enter one neutral molecule. Dot-separated salts and multicomponent structures are '
               'unsupported.',
 'charged': 'Enter the neutral molecular structure; charged structures are unsupported.',
 'isotope': 'Isotope-labeled structures are unsupported by this model.',
 'wildcard': 'Use explicit element symbols; wildcard atoms are unsupported.',
 'mapped': 'Remove atom-mapping numbers before predicting.',
 'failed': 'Prediction could not be completed. Check the input or ask the maintainer to inspect the runtime '
           'logs.',
 'blank': 'Enter a SMILES to display its predicted spectrum here',
 'note': 'Displays a model-predicted EI spectrum; chromatographic retention time is not predicted. Hover to '
         'inspect peaks, drag to zoom, and double-click to reset.',
 'peak': 'Peak',
 'predicted': 'SMART prediction',
 'stereochemistry_ignored': 'The current model does not distinguish stereoisomers.',
 'no_eligible_reference': 'No eligible reference spectrum was found; the base model prediction is used.',
 'reference_transfer_unavailable': 'Reference transfer was unavailable; the base model prediction is used.',
 'limited_mass_range': 'The output is limited to the model range of m/z 0–750.'}

EXAMPLES = [
    ("caffeine", "Cn1c(=O)c2c(ncn2C)n(C)c1=O"),
    ("ketamine", "CNC1(CCCCC1=O)c1ccccc1Cl"),
    ("lidocaine", "CCN(CC)CC(=O)Nc1c(C)cccc1C"),
]

CSS = """
.gradio-container {max-width: 1260px !important; margin: auto; font-family: Arial, sans-serif;}
#smart-header h1 {font-size: 38px; letter-spacing: 2px; margin-bottom: 2px; color: #17324d;}
#smart-header p {color: #536779; max-width: 830px;}
#smart-status {font-size: 13px; color: #536779; margin: 4px 0 18px;}
#smart-input textarea {font-family: ui-monospace, monospace; font-size: 15px;}
.smart-facts {display: flex; flex-wrap: wrap; gap: 16px 38px; padding: 8px 0 14px;}
.smart-fact small {display: block; color: #637789; font-size: 12px; margin-bottom: 5px;}
.smart-fact strong {font-size: 18px; color: #17324d; font-weight: 600;}
.smart-warning {color: #6d571d; background: #faf7ee; padding: 10px 14px; border-radius: 7px; line-height: 1.6; margin-top: 8px;}
.smart-smiles {overflow-wrap: anywhere; font-family: ui-monospace, monospace; color: #637789; font-size: 12px;}
footer {display: none !important;}
"""


def initialize_predictor():
    global PREDICTOR, LOAD_ERROR
    try:
        if DEVICE not in {"cpu", "cuda"}:
            raise ValueError("SMART_DEVICE must be cpu or cuda")
        if ZERO_GPU and spaces is None:
            raise RuntimeError("The spaces package is required for ZeroGPU")
        from inference import SMARTPredictor

        PREDICTOR = SMARTPredictor(ASSETS)
        PREDICTOR.to(DEVICE)
    except Exception as exc:
        PREDICTOR = None
        LOAD_ERROR = type(exc).__name__
        LOG.exception("SMART model initialization failed")


def validate_smiles(smiles):
    smiles = (smiles or "").strip()
    if not smiles:
        raise ValueError("empty")
    if len(smiles) > 1000:
        raise ValueError("long")
    if "\n" in smiles or "\r" in smiles or any(c.isspace() for c in smiles):
        raise ValueError("invalid")
    mol = Chem.MolFromSmiles(smiles)
    if mol is None or mol.GetNumAtoms() == 0:
        raise ValueError("invalid")
    if mol.GetNumHeavyAtoms() > 200:
        raise ValueError("large")
    if len(Chem.GetMolFrags(mol)) != 1:
        raise ValueError("components")
    if Chem.GetFormalCharge(mol) != 0:
        raise ValueError("charged")
    if any(atom.GetIsotope() for atom in mol.GetAtoms()):
        raise ValueError("isotope")
    if any(atom.GetAtomicNum() == 0 for atom in mol.GetAtoms()):
        raise ValueError("wildcard")
    if any(atom.GetAtomMapNum() for atom in mol.GetAtoms()):
        raise ValueError("mapped")
    return smiles


def run_model(smiles):
    return PREDICTOR.predict(smiles, device=DEVICE)


def predict(smiles):
    try:
        smiles = validate_smiles(smiles)
    except ValueError as exc:
        return {"error": str(exc)}
    if PREDICTOR is None:
        return {"error": "missing"}
    try:
        with LOCK:
            result = dict(run_model(smiles))
        spectrum = np.asarray(result["spectrum"], dtype=float)
        if spectrum.shape != (751,) or not np.isfinite(spectrum).all() or np.any(spectrum < 0) or spectrum.max() <= 0:
            raise ValueError("Invalid spectrum returned by model")
        result["spectrum"] = (spectrum / spectrum.max() * 100).tolist()
        result["input_smiles"] = smiles
        result["export_dir"] = tempfile.mkdtemp(prefix="result_", dir=EXPORT_ROOT)
        return result
    except Exception:
        LOG.exception("SMART prediction failed")
        return {"error": "failed"}


if spaces is not None:
    predict = spaces.GPU(duration=60)(predict)


def _smart_api_error(code, message):
    return {"ok": False, "api_version": 1, "error": {"code": code, "message": message}}


def _smart_api_plain(value):
    """Convert selected NumPy values without exposing internal result fields."""
    if isinstance(value, np.ndarray):
        return _smart_api_plain(value.tolist())
    if isinstance(value, np.generic):
        return _smart_api_plain(value.item())
    if value is None or isinstance(value, (str, bool, int, float)):
        return value
    if isinstance(value, (list, tuple)):
        return [_smart_api_plain(item) for item in value]
    if isinstance(value, dict) and all(isinstance(key, str) for key in value):
        return {key: _smart_api_plain(item) for key, item in value.items()}
    raise ValueError("Unsupported prediction value")


def _smart_api_payload(result, input_smiles):
    """Validate the existing predictor's result; do not recalculate its spectrum."""
    fields = (
        "model_version", "input_smiles", "canonical_smiles", "formula", "exact_mass",
        "mz", "intensity_scale", "reference_count", "selected_reference_count",
        "reference_covered", "training_reference_match", "warning_codes", "warnings",
    )
    payload = {key: _smart_api_plain(result[key]) for key in fields}
    # predict() has already normalized this spectrum for the existing interface.
    payload["intensity"] = _smart_api_plain(result["spectrum"])
    references = result["references"]
    if not isinstance(references, (list, tuple)):
        raise ValueError("Invalid references")
    reference_fields = ("group_key", "smiles", "similarity", "usable")
    payload["references"] = [
        {key: _smart_api_plain(reference[key]) for key in reference_fields}
        for reference in references
    ]

    def text_value(value):
        return isinstance(value, str) and bool(value.strip())

    def number(value):
        return type(value) in (int, float) and math.isfinite(value)

    for key in ("model_version", "input_smiles", "canonical_smiles", "formula", "intensity_scale"):
        if not text_value(payload[key]):
            raise ValueError("Missing prediction metadata")
    if payload["input_smiles"] != input_smiles.strip():
        raise ValueError("Prediction input mismatch")
    if not number(payload["exact_mass"]) or payload["exact_mass"] <= 0:
        raise ValueError("Invalid exact mass")
    mz = payload["mz"]
    if not isinstance(mz, list) or len(mz) != 751 or any(type(v) is not int or v != i for i, v in enumerate(mz)):
        raise ValueError("Invalid m/z bins")
    intensity = payload["intensity"]
    if not isinstance(intensity, list) or len(intensity) != 751:
        raise ValueError("Invalid intensity shape")
    if any(not number(v) or v < 0 or v > 100 for v in intensity):
        raise ValueError("Invalid intensity values")
    if not math.isclose(max(intensity), 100.0, rel_tol=0, abs_tol=1e-5):
        raise ValueError("Invalid base-peak normalization")
    for key in ("reference_count", "selected_reference_count"):
        if type(payload[key]) is not int or payload[key] < 0:
            raise ValueError("Invalid reference count")
    for key in ("reference_covered", "training_reference_match"):
        if type(payload[key]) is not bool:
            raise ValueError("Invalid prediction flag")
    for reference in payload["references"]:
        if not text_value(reference["group_key"]) or not text_value(reference["smiles"]):
            raise ValueError("Invalid reference identity")
        if not number(reference["similarity"]) or not 0 <= reference["similarity"] <= 1:
            raise ValueError("Invalid reference similarity")
        if type(reference["usable"]) is not bool:
            raise ValueError("Invalid reference usability")
    if payload["selected_reference_count"] != len(payload["references"]):
        raise ValueError("Inconsistent selected reference count")
    if payload["reference_count"] != sum(ref["usable"] for ref in payload["references"]):
        raise ValueError("Inconsistent usable reference count")
    for key in ("warning_codes", "warnings"):
        if not isinstance(payload[key], list) or any(not text_value(v) for v in payload[key]):
            raise ValueError("Invalid prediction warnings")
    if len(payload["warning_codes"]) != len(payload["warnings"]):
        raise ValueError("Inconsistent prediction warnings")
    payload.update(ok=True, api_version=1, model="SMART", ionization="EI")
    # Fail closed on any non-JSON value or malformed Unicode in model metadata.
    json.dumps(payload, ensure_ascii=False, allow_nan=False).encode("utf-8")
    return payload


def predict_spectrum(smiles):
    """Return an EI spectrum as JSON using the existing GPU-wrapped predictor."""
    if not isinstance(smiles, str):
        return _smart_api_error("invalid", TEXT["invalid"])
    try:
        smiles.encode("utf-8")
    except UnicodeError:
        return _smart_api_error("invalid", TEXT["invalid"])
    try:
        result = predict(smiles)
    except Exception:
        LOG.exception("SMART API prediction failed")
        return _smart_api_error("failed", TEXT["failed"])
    if not isinstance(result, dict):
        return _smart_api_error("invalid_result", "SMART returned an invalid prediction result.")
    if "error" in result:
        known_codes = {"empty", "long", "invalid", "large", "components", "charged",
                       "isotope", "wildcard", "mapped", "missing", "failed"}
        code = result["error"]
        if not isinstance(code, str) or code not in known_codes:
            code = "failed"
        return _smart_api_error(code, TEXT[code])
    try:
        return _smart_api_payload(result, smiles)
    except Exception:
        LOG.exception("SMART API result validation failed")
        return _smart_api_error("invalid_result", "SMART returned an invalid prediction result.")


def spectrum_plot(result):
    t = TEXT
    fig = go.Figure()
    if result and "error" not in result:
        y = np.asarray(result["spectrum"], dtype=float)
        masses = np.flatnonzero(y > 0)
        x_line, y_line = [], []
        for mass in masses:
            x_line.extend([int(mass), int(mass), None])
            y_line.extend([0.0, float(y[mass]), None])
        fig.add_trace(go.Scatter(x=x_line, y=y_line, mode="lines", line={"color": "#2864a0", "width": 1.5}, hoverinfo="skip", showlegend=False))
        fig.add_trace(go.Scatter(x=masses.tolist(), y=y[masses].tolist(), mode="markers", marker={"size": 7, "color": "#2864a0", "opacity": 0.35}, name=t["peak"], showlegend=False, hovertemplate="m/z %{x:.0f}<br>" + t["intensity"] + ": %{y:.2f}<extra></extra>"))
        for mass in sorted(masses, key=lambda m: y[m], reverse=True)[:8]:
            if y[mass] >= 5:
                fig.add_annotation(x=int(mass), y=float(y[mass]), text=str(mass), showarrow=False, yshift=11, font={"size": 12, "color": "#17324d"})
        x_max = min(750, max(100, int(masses.max()) + 15))
    else:
        x_max = 400
        fig.add_annotation(x=0.5, y=0.5, xref="paper", yref="paper", text=t.get((result or {}).get("error"), t["blank"]), showarrow=False, font={"color": "#6b7f91", "size": 14})
    fig.update_layout(template="plotly_white", height=440, margin={"l": 62, "r": 25, "t": 30, "b": 58}, paper_bgcolor="white", plot_bgcolor="white", font={"family": "Arial, sans-serif", "color": "#344c62", "size": 13}, dragmode="zoom", hovermode="closest", uirevision="smart-spectrum")
    fig.update_xaxes(title="m/z", range=[0, x_max], showgrid=False, zeroline=False, ticks="outside", linecolor="#9aabba", mirror=False, showline=True)
    fig.update_yaxes(title=t["intensity"], range=[0, 116], dtick=20, gridcolor="#edf1f5", zeroline=False, ticks="outside", linecolor="#9aabba", showline=True)
    return fig


def export_files(result):
    with EXPORT_LOCK:
        return write_exports(result)


def write_exports(result):
    folder = Path(result["export_dir"]).resolve()
    if not folder.is_relative_to(EXPORT_ROOT.resolve()):
        raise ValueError("Invalid export folder")
    folder.mkdir(parents=True, exist_ok=True)
    y = np.asarray(result["spectrum"], dtype=float)
    csv_path = folder / "SMART_spectrum.csv"
    msp_path = folder / "SMART_spectrum.msp"
    png_path = folder / "SMART_spectrum.png"
    with csv_path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["mz", "relative_intensity"])
        writer.writerows((int(m), f"{v:.6f}") for m, v in enumerate(y))
    active = np.flatnonzero(y > 0)
    smiles = result.get("canonical_smiles", result["smiles"])
    lines = ["Name: SMART predicted EI spectrum", f"SMILES: {smiles}", f"Formula: {result.get('formula', '')}", f"ExactMass: {float(result.get('exact_mass', 0)):.6f}", "Ionization: EI", "Comment: Model prediction; not an experimental spectrum", f"Num Peaks: {len(active)}"]
    lines.extend(f"{int(m)} {float(y[m]):.6f}" for m in active)
    msp_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    if not png_path.exists():
        t = TEXT
        fig, ax = plt.subplots(figsize=(10, 4.3), dpi=180)
        ax.vlines(active, 0, y[active], color="#2864a0", linewidth=1.0)
        for mass in sorted(active, key=lambda m: y[m], reverse=True)[:8]:
            if y[mass] >= 5:
                ax.annotate(str(mass), (mass, y[mass]), xytext=(0, 4), textcoords="offset points", ha="center", fontsize=9, color="#17324d")
        ax.set(xlim=(0, min(750, max(100, int(active.max()) + 15))), ylim=(0, 116), xlabel="m/z")
        ax.set_ylabel(t["intensity"])
        ax.set_title(t["spectrum"] + " · SMART", loc="left", fontsize=12, pad=18)
        ax.spines[["top", "right"]].set_visible(False)
        ax.grid(axis="y", color="#edf1f5")
        ax.set_axisbelow(True)
        fig.tight_layout()
        fig.savefig(png_path, bbox_inches="tight", facecolor="white")
        plt.close(fig)
    return str(csv_path), str(msp_path), str(png_path)


def render(result):
    t = TEXT
    plot = spectrum_plot(result)
    if not result or "error" in result:
        error = t.get((result or {}).get("error"), "")
        details = f'<div class="smart-warning">{html.escape(error)}</div>' if error else ""
        return None, plot, details, gr.update(value=None, visible=False), gr.update(value=None, visible=False), gr.update(value=None, visible=False)
    smiles = result.get("canonical_smiles", result["smiles"])
    mol = Chem.MolFromSmiles(smiles)
    molecule_image = Draw.MolToImage(mol, size=(500, 310)) if mol is not None else None
    facts = [(t["formula"], result.get("formula", "—")), (t["mass"], f"{float(result.get('exact_mass', 0)):.4f} Da"), (t["reference"], str(result.get("reference_count", 0)))]
    details = '<div class="smart-facts">' + "".join(f'<div class="smart-fact"><small>{html.escape(label)}</small><strong>{html.escape(str(value))}</strong></div>' for label, value in facts) + "</div>"
    details += '<div class="smart-smiles">' + html.escape(smiles) + "</div>"
    codes = result.get("warning_codes", [])
    warnings = [t[code] for code in codes if code in t]
    if warnings:
        details += '<div class="smart-warning">' + "<br>".join(html.escape(s) for s in warnings) + "</div>"
    paths = export_files(result)
    return molecule_image, plot, details, *(gr.update(value=path, visible=True) for path in paths)


def cleanup_exports():
    cutoff = time.time() - 24 * 3600
    for path in EXPORT_ROOT.iterdir():
        if path.is_dir() and path.stat().st_mtime < cutoff:
            shutil.rmtree(path, ignore_errors=True)


def build_app():
    t = TEXT
    with gr.Blocks(title="SMART · EI Spectrum Prediction", delete_cache=(3600, 86400), analytics_enabled=False) as demo:
        state = gr.State(None, time_to_live=86400)
        gr.Markdown("# SMART\n\n" + t["intro"], elem_id="smart-header")
        gr.Markdown(t["ready"] if PREDICTOR is not None else t["missing"], elem_id="smart-status")
        with gr.Row():
            with gr.Column(scale=1, min_width=290):
                smiles = gr.Textbox(label=t["input"], placeholder=t["placeholder"], lines=4, max_lines=6, elem_id="smart-input")
                with gr.Row():
                    run = gr.Button(t["predict"], variant="primary")
                    clear = gr.Button(t["clear"], variant="secondary")
                gr.Markdown(t["examples"])
                with gr.Row():
                    example_buttons = [gr.Button(t[name], size="sm") for name, value in EXAMPLES]
                structure = gr.Image(label=t["structure"], type="pil", height=260, interactive=False, buttons=[])
            with gr.Column(scale=2, min_width=400):
                plot = gr.Plot(value=spectrum_plot(None), label=t["spectrum"])
                details = gr.HTML()
                with gr.Row():
                    download_csv = gr.DownloadButton(t["csv"], visible=False)
                    download_msp = gr.DownloadButton(t["msp"], visible=False)
                    download_png = gr.DownloadButton(t["png"], visible=False)
        gr.Markdown(t["note"])
        result_outputs = [structure, plot, details, download_csv, download_msp, download_png]

        run.click(predict, smiles, state, concurrency_limit=1, concurrency_id="smart-predict", api_name="predict").then(render, state, result_outputs, api_visibility="private")
        smiles.submit(predict, smiles, state, concurrency_limit=1, concurrency_id="smart-predict", api_visibility="private").then(render, state, result_outputs, api_visibility="private")
        clear.click(lambda: ("", None), outputs=[smiles, state], queue=False, api_visibility="private").then(render, state, result_outputs, queue=False, api_visibility="private")
        for button, (_, sample_smiles) in zip(example_buttons, EXAMPLES):
            button.click(lambda value=sample_smiles: value, outputs=smiles, queue=False, api_visibility="private")
        timer = gr.Timer(3600)
        timer.tick(cleanup_exports, queue=False, api_visibility="private")
        api_result = gr.JSON(label="SMART spectrum API", visible=False)
        api_trigger = gr.Button("SMART spectrum API", visible=False)
        api_trigger.click(
            predict_spectrum, inputs=smiles, outputs=api_result,
            concurrency_limit=1, concurrency_id="smart-predict",
            api_name="predict_spectrum", api_visibility="public",
        )
    return demo


initialize_predictor()
demo = build_app()

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=16).launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
        ssr_mode=False,
        theme=gr.themes.Soft(primary_hue="blue", neutral_hue="slate"),
        css=CSS,
        allowed_paths=[str(EXPORT_ROOT)],
        blocked_paths=[str(ASSETS)],
        show_error=False,
        max_file_size="1mb",
    )