"""CliniGuard Laboratory NER — interactive demo for clinical laboratory result entity extraction. Built on Bio_ClinicalBERT fine-tuned for token classification (BIO tagging, 21 labels, 10 laboratory result entity categories). Runs on ZeroGPU. """ import spaces # MUST be first — before torch / transformers import torch from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline import gradio as gr MODEL_ID = "genzeonplatform/cliniguard-laboratory-ner" # Entity type -> (legend label, description, hex color). Colors are picked to be # visually distinct across the 10 lab categories when rendered by HighlightedText. ENTITY_META = { "LAB_TEST_NAME": ("Test name", "#1f77b4"), "LAB_VALUE": ("Value", "#ff7f0e"), "LAB_UNIT": ("Unit", "#2ca02c"), "REFERENCE_RANGE": ("Reference range", "#d62728"), "ABNORMALITY_FLAG": ("Abnormality flag", "#9467bd"), "SPECIMEN_TYPE": ("Specimen", "#8c564b"), "LAB_DATE": ("Date", "#e377c2"), "LOINC_CODE": ("LOINC code", "#17becf"), "TEST_PANEL": ("Panel", "#bcbd22"), "LAB_STATUS": ("Status", "#7f7f7f"), } DESCRIPTIONS = { "LAB_TEST_NAME": "Name of the laboratory test (e.g. hemoglobin, glucose, troponin I)", "LAB_VALUE": "Numeric or qualitative result (e.g. 7.2, 145, positive, trace)", "LAB_UNIT": "Unit of measurement (e.g. mg/dL, g/dL, mEq/L, x10^3/uL)", "REFERENCE_RANGE": "Normal reference range (e.g. 12.0-17.5, < 200, 3.5-5.0)", "ABNORMALITY_FLAG": "Abnormality indicator (e.g. H, L, Critical High, Normal)", "SPECIMEN_TYPE": "Type of specimen (e.g. blood, serum, urine, CSF)", "LAB_DATE": "Collection date/time (e.g. 03/15/2024, this morning, hospital day 3)", "LOINC_CODE": "LOINC identifier code (e.g. 718-7, 2345-7, 2160-0)", "TEST_PANEL": "Panel/order set name (e.g. CBC, BMP, CMP, lipid panel)", "LAB_STATUS": "Result status (e.g. final, preliminary, pending)", } # Load model at module scope, eager .to("cuda") — ZeroGPU intercepts & packs weights. tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) model = AutoModelForTokenClassification.from_pretrained(MODEL_ID).to("cuda") nlp = pipeline( "token-classification", model=model, tokenizer=tokenizer, aggregation_strategy="simple", device=0, # pipeline follows model device ) def _entities_to_highlight(text: str, entities): """Convert HF pipeline entities into the (text, label) list HighlightedText wants, filling the gaps with plain text spans.""" if not entities: return [(text, None)] parts = [] cursor = 0 for ent in entities: start = ent["start"] end = ent["end"] label = ent["entity_group"] if start > cursor: parts.append((text[cursor:start], None)) parts.append((text[start:end], label)) cursor = end if cursor < len(text): parts.append((text[cursor:], None)) return parts @spaces.GPU(duration=30) def extract_entities(text: str): """Extract laboratory result entities (test names, values, units, ranges, flags, …) from clinical text. Args: text: clinical text containing laboratory results (lab reports, progress notes, discharge summaries). """ text = (text or "").strip() if not text: return [], "Please enter some clinical text to analyze." entities = nlp(text) # Summary table for the side panel rows = [] for ent in entities: label = ent["entity_group"] legend_label = ENTITY_META.get(label, (label, None))[0] rows.append([legend_label, ent["word"], f"{ent['score']:.3f}", ent["start"], ent["end"]]) highlight = _entities_to_highlight(text, entities) n = len(rows) summary = f"Found **{n}** entit{'y' if n == 1 else 'ies'} across {len({r[0] for r in rows}) or 0} categor{'y' if n == 1 else 'ies'}." if n else "No entities found." return highlight, rows, summary # Build the legend options for HighlightedText from the model's labels. # HighlightedText accepts `combine_entities=True` to merge adjacent same-label spans. HIGHLIGHT_LABELS = {label: meta[1] for label, meta in ENTITY_META.items()} CSS = """ #col-container { max-width: 1100px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } .legend-chip { display: inline-block; padding: 2px 8px; margin: 2px; border-radius: 6px; color: #fff; font-size: 0.85em; font-weight: 600; } """ EXAMPLES = [ ["Lab Results: Hemoglobin: 7.2 g/dL (Reference: 12.0-17.5) [Critical Low]. Specimen: blood. Status: final. CBC panel. LOINC: 718-7."], ["Sodium 145 mEq/L (ref 135-145) [H]. Potassium 3.2 mEq/L (ref 3.5-5.0) [L]. BMP panel drawn this morning, blood specimen, preliminary."], ["Glucose: 210 mg/dL (ref 70-99) [Critical High]. HbA1c 9.8% (ref < 5.7). Lipid panel: LDL 160 mg/dL [H]. Status: final. LOINC: 2345-7."], ["Troponin I 0.05 ng/mL (ref < 0.04) [H]. CK-MB 8.2 ng/mL (ref < 6.3) [H]. Cardiac panel, serum specimen, hospital day 3, pending."], ] with gr.Blocks(title="CliniGuard Laboratory NER") as demo: gr.Markdown( "# 🧪 CliniGuard Laboratory NER\n" "Extract laboratory result entities — test names, values, units, reference ranges, " "abnormality flags, specimen types, dates, LOINC codes, panels, and status — from " "unstructured clinical text. Built on **Bio_ClinicalBERT** fine-tuned by Genzeon Platforms " "for token classification across 10 laboratory entity categories." ) legend_html = "
" + "".join( f'{label}' for label, (_, color) in ENTITY_META.items() ) + "
" gr.HTML(legend_html) with gr.Row(): text_in = gr.Textbox( label="Clinical text", placeholder="Paste clinical text containing laboratory results…", lines=8, scale=3, ) run = gr.Button("Extract entities", variant="primary", scale=1) highlight_out = gr.HighlightedText( label="Highlighted entities", combine_adjacent=True, show_legend=True, color_map={label: meta[1] for label, meta in ENTITY_META.items()}, ) summary_md = gr.Markdown() table_out = gr.Dataframe( headers=["Category", "Entity", "Score", "Start", "End"], label="Extracted entities", interactive=False, wrap=True, ) run.click( fn=extract_entities, inputs=text_in, outputs=[highlight_out, table_out, summary_md], api_name="extract_entities", ) text_in.submit( fn=extract_entities, inputs=text_in, outputs=[highlight_out, table_out, summary_md], api_name="extract_entities_submit", ) gr.Examples( examples=EXAMPLES, inputs=text_in, outputs=[highlight_out, table_out, summary_md], fn=extract_entities, cache_examples=True, cache_mode="lazy", ) with gr.Accordion("Entity types", open=False): gr.Markdown("\n".join( f"- **{label}** — {DESCRIPTIONS[label]}" for label in ENTITY_META )) if __name__ == "__main__": demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)