aakothari/tox21_deepberta

1"""2Gradio interface for the Tox21 DeepBERTa model, on ZeroGPU.3 4The browser UI works. The raw /predict and /metadata endpoints the5leaderboard harness needs are NOT served here -- Gradio's catch-all routing6claims every path. Those require a Docker Space (see the FastAPI app.py).7 8NOTE: `spaces` only exists inside a ZeroGPU Space, so this file will not run9locally. Use predict.py directly for local testing.10"""11 12import json13 14import gradio as gr15import spaces16 17from predict import predict, TASKS18 19EXAMPLES = [20    "CCO",21    "c1ccc2c(c1)ccc1ccccc12",22    "CC(=O)Oc1ccccc1C(=O)O",23]24 25 26@spaces.GPU(duration=120)27def predict_api(smiles_list):28    """Score molecules on all 12 Tox21 endpoints."""29    if isinstance(smiles_list, str):30        raw = smiles_list.replace(",", "\n").split("\n")31        smiles_list = [s.strip() for s in raw if s.strip()]32    if not smiles_list:33        return {}34    return predict(smiles_list)35 36 37def ui_predict(text):38    """Browser-facing wrapper: returns a table plus the raw JSON."""39    preds = predict_api(text)40    if not preds:41        return [], json.dumps({"error": "no SMILES provided"}, indent=2)42    rows = [[smi] + [round(scores[t], 4) for t in TASKS]43            for smi, scores in preds.items()]44    return rows, json.dumps(preds, indent=2)45 46 47with gr.Blocks(title="DeepBERTa Tox21") as demo:48    gr.Markdown(49        "# DeepBERTa-Tox21\n"50        "DeepSMILES RoBERTa fine-tuned on the original Tox21 Challenge "51        "training split. Returns a ranking score in [0, 1] for each of the "52        "12 endpoints.\n\n"53        "**Scores are rankings, not calibrated probabilities.** "54        "Molecules that cannot be converted to DeepSMILES receive 0.5 on all "55        "tasks rather than being dropped."56    )57 58    inp = gr.Textbox(59        label="SMILES (one per line)",60        lines=6,61        placeholder="CCO\nc1ccc2c(c1)ccc1ccccc12",62    )63    btn = gr.Button("Predict", variant="primary")64    table = gr.Dataframe(headers=["SMILES"] + TASKS, label="Predictions", wrap=True)65    raw = gr.Code(label="Raw JSON response", language="json")66 67    btn.click(ui_predict, inputs=inp, outputs=[table, raw])68    gr.Examples(examples=[["\n".join(EXAMPLES)]], inputs=inp)69 70 71demo.launch()