"""DINNNER online tool: biomedical named entity recognition in the browser. Runs as a Gradio app on a Hugging Face Space (ZeroGPU or CPU). Locally: python app.py """ import os # On ZeroGPU Spaces, `spaces` must be imported before torch touches CUDA. try: import spaces except ImportError: # not installed locally; everything runs on CPU/GPU as usual spaces = None ON_ZEROGPU = spaces is not None and os.environ.get("SPACES_ZERO_GPU", "").lower() in ("1", "true") import html # noqa: E402 import tempfile # noqa: E402 import time # noqa: E402 import gradio as gr # noqa: E402 import pandas as pd # noqa: E402 import torch # noqa: E402 from dinnner import DEFAULT_MODEL, MODELS, load_model, predict # noqa: E402 from dinnner.config import BATCH_SIZE, MAX_PDF_FILES, MAX_PDF_PAGES, MAX_TEXT_CHARS # noqa: E402 from dinnner.examples import EXAMPLES # noqa: E402 from dinnner.text import conll_text, entities_frame, extract_pdf_text, sentences_frame, split_sentences # noqa: E402 GITHUB_URL = os.environ.get("DINNNER_GITHUB_URL", "") # set once the code is public PAPER_URL = os.environ.get("DINNNER_PAPER_URL", "") SHOW_LIMIT = 300 # sentences rendered in the annotated view (all are in the downloads) MOBILE_BREAKPOINT = 768 # px; narrower screens start with the sidebar collapsed # Colour per entity type: (background, border). Readable in light and dark themes. PALETTE = [ ("rgba(59,130,246,0.18)", "#3b82f6"), # blue ("rgba(16,185,129,0.20)", "#10b981"), # green ("rgba(245,158,11,0.22)", "#f59e0b"), # amber ("rgba(236,72,153,0.18)", "#ec4899"), # pink ("rgba(139,92,246,0.18)", "#8b5cf6"), # violet ("rgba(20,184,166,0.20)", "#14b8a6"), # teal ] CSS = """ .ent {padding: 0.05em 0.3em; border-radius: 0.3em; border-bottom: 2px solid; line-height: 2;} .ent sub {font-size: 0.62em; font-weight: 600; letter-spacing: 0.02em; margin-left: 0.3em; opacity: 0.85; text-transform: uppercase; white-space: nowrap;} .sent {margin: 0 0 0.6em 0; line-height: 1.9;} .sent-id {display: inline-block; min-width: 2.4em; font-size: 0.75em; opacity: 0.55;} .metrics {display: flex; flex-wrap: wrap; gap: 0.75rem; margin: 0.25rem 0 0.5rem;} .metric {flex: 1 1 8.5rem; padding: 0.6rem 0.8rem; border-radius: 0.6rem; border: 1px solid var(--border-color-primary); background: var(--block-background-fill);} .metric .v {font-size: 1.6rem; font-weight: 600; line-height: 1.2;} .metric .k {font-size: 0.8rem; opacity: 0.7;} .metric.typed {border-left-width: 4px;} .legend-row {margin-bottom: 0.5em;} .legend-row small {display: block; opacity: 0.7; margin-top: 0.1em;} .bars .row {display: grid; grid-template-columns: 11rem minmax(0, 32rem) 2.5rem; align-items: center; gap: 0.6rem; margin: 0.35rem 0;} .bars .bar {height: 0.9rem; border-radius: 0 0.25rem 0.25rem 0;} .bars .n {font-variant-numeric: tabular-nums; text-align: right;} .muted {opacity: 0.7; font-size: 0.9em;} """ # -------------------------------------------------------------------------- # Model (loaded once at start-up so visitors do not wait on the first click) # -------------------------------------------------------------------------- _MODELS, _LOAD_ERRORS = {}, {} # On ZeroGPU the model is placed on "cuda" at start-up (an emulated device); # a real GPU is attached only while a @spaces.GPU function runs. DEVICE = torch.device("cuda" if ON_ZEROGPU or torch.cuda.is_available() else "cpu") def get_model(key): if key not in _MODELS: try: _MODELS[key] = load_model(MODELS[key], DEVICE) _LOAD_ERRORS.pop(key, None) except Exception as exc: # weights missing, network down, wrong checkpoint... _LOAD_ERRORS[key] = str(exc) raise gr.Error(f"The model could not be loaded: {exc}", duration=None) return _MODELS[key] def _prepare_nltk(): try: import nltk nltk.download("punkt_tab", quiet=True, raise_on_error=True) except Exception: pass # the built-in regex sentence splitter is used instead # -------------------------------------------------------------------------- # HTML helpers # -------------------------------------------------------------------------- def type_colours(spec): types = list(spec.entity_info) or sorted({l[2:] for l in spec.id2label.values() if l != "O"}) return {t: PALETTE[i % len(PALETTE)] for i, t in enumerate(types)} def pretty(etype): return etype.replace("_", " ") def ent_html(text, etype, colours, tag=True): bg, border = colours.get(etype, PALETTE[-1]) label = f"{html.escape(pretty(etype))}" if tag else "" return f'{html.escape(text)}{label}' def sentence_html(res, colours, keep): out, pos = [], 0 for e in res.entities: if e.type not in keep: continue out += [html.escape(res.text[pos:e.start]), ent_html(e.text, e.type, colours)] pos = e.end out.append(html.escape(res.text[pos:])) return f'

{res.index + 1}{"".join(out)}

' def legend_html(spec): colours = type_colours(spec) return "".join( f'
{ent_html(pretty(t), t, colours, tag=False)}{html.escape(d)}
' for t, d in spec.entity_info.items() ) # -------------------------------------------------------------------------- # Processing # -------------------------------------------------------------------------- def _tag_docs(docs, model_key, progress_cb=None): """Run DINNNER over {source: [sentences]}; returns {source: [SentenceResult]}.""" spec = MODELS[model_key] model, tokenizer = _MODELS[model_key] total = sum(len(s) for s in docs.values()) results, done = {}, 0 for src, sents in docs.items(): def cb(d, n, base=done, k=len(sents), src=src): if progress_cb: progress_cb((base + k * d / max(n, 1)) / total, src) results[src] = predict(sents, model, tokenizer, spec.id2label, DEVICE, max_length=spec.max_length, batch_size=BATCH_SIZE, progress=cb) done += len(sents) return results def _gpu_seconds(docs, model_key): """GPU time to reserve on ZeroGPU: short requests get better queue priority.""" n = sum(len(s) for s in docs.values()) return int(min(120, 15 + n / 25)) def _tag_docs_gpu(docs, model_key): return _tag_docs(docs, model_key) if ON_ZEROGPU: _tag_docs_gpu = spaces.GPU(duration=_gpu_seconds)(_tag_docs_gpu) def _run(docs, model_key, notes, progress): if not any(docs.values()): raise gr.Error("No sentences were found in the input.") spec = MODELS[model_key] progress(0, desc="Loading model...") get_model(model_key) # surfaces loading errors before any GPU is requested t0 = time.time() if ON_ZEROGPU: progress(0.05, desc="Tagging entities on GPU...") results = _tag_docs_gpu(docs, model_key) else: results = _tag_docs(docs, model_key, lambda f, src: progress(f, desc=f"Tagging entities in {src}")) # Write the downloads once per run (they always contain everything). out_dir = tempfile.mkdtemp(prefix="dinnner_") stem = "dinnner_" + (list(results)[0].rsplit(".", 1)[0][:40].replace(" ", "_") if len(results) == 1 else "results") paths = [os.path.join(out_dir, f"{stem}_{s}") for s in ("entities.csv", "sentences.csv", "iob.tsv")] entities_frame(results).to_csv(paths[0], index=False) sentences_frame(results).to_csv(paths[1], index=False) with open(paths[2], "w", encoding="utf-8") as f: f.write(conll_text(results)) state = {"results": results, "model_key": model_key, "seconds": time.time() - t0, "notes": notes} types = list(type_colours(spec)) sources = list(results) doc = sources[0] only_hits = len(results[doc]) > 40 return ( state, gr.update(visible=True), # results column summary_html(state), gr.update(choices=[(pretty(t), t) for t in types], value=types), # type filter gr.update(choices=sources, value=doc, visible=len(sources) > 1), # document picker gr.update(value=only_hits), *render_views(state, types, doc, only_hits), *[gr.update(value=p, visible=True) for p in paths], ) def analyse_text(text, model_key, progress=gr.Progress()): text = (text or "").strip() if not text: raise gr.Error("Please paste some text first, or pick an example.") notes = [] if len(text) > MAX_TEXT_CHARS: text = text[:MAX_TEXT_CHARS] notes.append(f"Only the first {MAX_TEXT_CHARS:,} characters were processed.") return _run({"Pasted text": split_sentences(text)}, model_key, notes, progress) def analyse_pdfs(files, body_only, model_key, progress=gr.Progress()): if not files: raise gr.Error("Please upload at least one PDF.") docs, notes = {}, [] if len(files) > MAX_PDF_FILES: notes.append(f"Only the first {MAX_PDF_FILES} files were processed.") for path in files[:MAX_PDF_FILES]: path = getattr(path, "name", path) name = os.path.basename(path) try: with open(path, "rb") as f: text, read, total = extract_pdf_text(f.read(), MAX_PDF_PAGES, body_only) except Exception as exc: notes.append(f"Could not read {name}: {exc}") continue if not text: notes.append(f"No extractable text in {name} (it may be a scanned image).") continue if read < total: notes.append(f"{name}: only the first {read} of {total} pages were processed.") docs[name] = split_sentences(text) if not docs: raise gr.Error(" ".join(notes) or "No text could be read from the PDFs.") return _run(docs, model_key, notes, progress) # -------------------------------------------------------------------------- # Views (re-rendered from stored results when filters change, no re-run) # -------------------------------------------------------------------------- def summary_html(state): spec = MODELS[state["model_key"]] colours = type_colours(spec) ents = entities_frame(state["results"]) n_sent = sum(len(r) for r in state["results"].values()) cards = [("Sentences", n_sent, None), ("Entities", len(ents), None)] cards += [(pretty(t), int((ents["type"] == t).sum()), colours[t][1]) for t in colours] body = "".join( f'
' f'
{v:,}
{html.escape(k)}
' for k, v, c in cards ) notes = "".join(f"
ℹ️ {html.escape(n)}
" for n in state["notes"]) return (f'
{body}
' f'
{html.escape(spec.display_name)} · processed in {state["seconds"]:.1f} s
{notes}') def render_views(state, keep_types, doc, only_hits): if not state: return "", pd.DataFrame(), pd.DataFrame(), "" spec = MODELS[state["model_key"]] colours = type_colours(spec) keep = set(keep_types or []) results = state["results"] doc = doc if doc in results else list(results)[0] # Annotated text shown = [r for r in results[doc] if not only_hits or any(e.type in keep for e in r.entities)] annotated = "".join(sentence_html(r, colours, keep) for r in shown[:SHOW_LIMIT]) if not shown: annotated = "

No entities of the selected types were found.

" elif len(shown) > SHOW_LIMIT: annotated += f"

Showing the first {SHOW_LIMIT} of {len(shown)} sentences. The downloads contain all of them.

" # Entity table ents = entities_frame(results) ents = ents[ents["type"].isin(keep)] table = pd.DataFrame({ "Source": ents["source"], "Sent.": ents["sentence_id"], "Entity": ents["entity"], "Type": ents["type"].map(pretty), "p(entity)": ents["p_entity"].round(3), "p(type)": ents["p_type"].round(3), "Sentence": ents["sentence"], }) if len(results) == 1: table = table.drop(columns="Source") # Per-entity summary and bar chart if ents.empty: summary, bars = pd.DataFrame(columns=["Entity", "Type", "Mentions", "Sentences", "Mean p(entity)"]), "" else: summary = ( ents.assign(key=ents["entity"].str.lower()) .groupby(["key", "type"], as_index=False) .agg(Entity=("entity", "first"), Mentions=("entity", "size"), Sentences=("sentence_id", "nunique"), mean=("p_entity", "mean")) .sort_values(["Mentions", "Entity"], ascending=[False, True]) ) summary = pd.DataFrame({ "Entity": summary["Entity"], "Type": summary["type"].map(pretty), "Mentions": summary["Mentions"], "Sentences": summary["Sentences"], "Mean p(entity)": summary["mean"].round(3), }) counts = ents["type"].value_counts() top = counts.max() bars = "
" + "".join( f"
{html.escape(pretty(t))}" f"
" f"{n}
" for t, n in counts.items() ) + "
" return annotated, table, summary, bars # -------------------------------------------------------------------------- # Interface # -------------------------------------------------------------------------- spec0 = MODELS[DEFAULT_MODEL] links = [f"[{name}]({url})" for name, url in (("Source code", GITHUB_URL), ("Paper", PAPER_URL)) if url] with gr.Blocks(title="DINNNER | Biomedical NER") as demo: state = gr.State(None) gr.Markdown( "# 🧬 DINNNER: Find Entities in Biomedical Text\n" "Paste biomedical text or upload research articles (PDF). DINNNER tags **genes and gene products, " "simple chemicals, complexes and cellular components**. Results appear below and can be downloaded as CSV." ) # Starts collapsed (on phones an open sidebar covers the whole page); opened on wide screens below. with gr.Sidebar(open=False, width=310, position="left") as sidebar: gr.Markdown("## DINNNER\nDistribution-Informed Neural Network for Named Entity Recognition") if len(MODELS) > 1: model_key = gr.Dropdown( choices=[(m.display_name, k) for k, m in MODELS.items()], value=DEFAULT_MODEL, label="Model", ) else: model_key = gr.State(DEFAULT_MODEL) gr.Markdown(f"**Model** \n{spec0.display_name}") gr.Markdown("**Entity types**") gr.HTML(legend_html(spec0)) gr.Markdown( "---\n" "Every entity comes with two confidence scores: **p(entity)**, how sure the model is that " "the words are an entity, and **p(type)**, how sure it is of the entity type.\n\n" "The paper describing DINNNER is under review. Details of the method will be added here " "once it is published.\n\n" + " · ".join(links) + f"\n\n{spec0.description}" ) with gr.Tabs(): with gr.Tab("✍️ Paste text"): text_in = gr.Textbox( lines=8, max_lines=20, max_length=MAX_TEXT_CHARS, show_label=False, placeholder="Paste an abstract, a paragraph or a full article here...", ) gr.Examples(examples=[[t] for t in EXAMPLES.values()], inputs=[text_in], example_labels=list(EXAMPLES), label="Try an example") run_text = gr.Button("Recognise entities", variant="primary") with gr.Tab("📄 Upload PDFs"): pdf_in = gr.File( file_count="multiple", file_types=[".pdf"], type="filepath", label=f"Up to {MAX_PDF_FILES} PDFs (first {MAX_PDF_PAGES} pages of each are processed)", ) body_only = gr.Checkbox( value=True, label="Process the article body only (from the Abstract or Introduction up to the References)", ) run_pdf = gr.Button("Recognise entities", variant="primary") with gr.Column(visible=False) as results_col: gr.Markdown("## Results") summary_out = gr.HTML() type_filter = gr.CheckboxGroup(label="Entity types to show", choices=[]) with gr.Tabs(): with gr.Tab("Annotated text"): with gr.Row(): doc_pick = gr.Dropdown(label="Document", choices=[], visible=False, scale=2) only_hits = gr.Checkbox(label="Only show sentences that contain entities", scale=3) annotated_out = gr.HTML() with gr.Tab("Entity table"): table_out = gr.Dataframe(wrap=True, max_height=520, show_search="search", interactive=False) with gr.Tab("Entity summary"): bars_out = gr.HTML() entity_summary_out = gr.Dataframe(max_height=520, interactive=False) with gr.Tab("⬇️ Download"): gr.Markdown( "Full results (all entity types, all sentences).\n\n" "* **Entities (CSV):** one row per entity with type, character offsets, both probabilities and the sentence.\n" "* **Annotated sentences (CSV):** one row per sentence, entities marked inline as `text`.\n" "* **Token labels (TSV):** CoNLL-style word and IOB label per line." ) with gr.Row(): dl_ents = gr.DownloadButton("Entities (CSV)", visible=False) dl_sents = gr.DownloadButton("Annotated sentences (CSV)", visible=False) dl_iob = gr.DownloadButton("Token labels (TSV)", visible=False) view_outputs = [annotated_out, table_out, entity_summary_out, bars_out] run_outputs = [state, results_col, summary_out, type_filter, doc_pick, only_hits, *view_outputs, dl_ents, dl_sents, dl_iob] run_text.click(analyse_text, [text_in, model_key], run_outputs, api_name="tag_text") run_pdf.click(analyse_pdfs, [pdf_in, body_only, model_key], run_outputs, api_name="tag_pdfs") for comp in (type_filter, doc_pick, only_hits): comp.input(render_views, [state, type_filter, doc_pick, only_hits], view_outputs, api_visibility="private") screen_width = gr.Number(visible=False) demo.load(lambda w: gr.Sidebar(open=(w or 0) >= MOBILE_BREAKPOINT), screen_width, sidebar, js="() => [window.innerWidth]", concurrency_limit=None, api_visibility="private") if __name__ == "__main__": _prepare_nltk() try: get_model(DEFAULT_MODEL) # warm up so the first visitor does not wait print(f"DINNNER loaded on {DEVICE}") except gr.Error: print(f"WARNING: model not loaded at start-up: {_LOAD_ERRORS.get(DEFAULT_MODEL)}") demo.queue(default_concurrency_limit=2).launch( theme=gr.themes.Soft(primary_hue="teal", secondary_hue="teal", font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"]), css=CSS, server_name=os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0"), )