import spaces # MUST be first — monkey-patches torch.cuda before anything else touches CUDA import torch import gradio as gr from transformers import pipeline MODEL_ID = "genzeonplatform/cliniguard-medication-ner" # Load the token-classification pipeline at module scope. The model is a # ~110M-param Bio_ClinicalBERT fine-tuned for medication NER (25 BIO labels, # 12 medication entity types). `aggregation_strategy="simple"` is the model # card's documented setting; we then post-merge adjacent same-group spans # whose character offsets touch (e.g. "Met"+"##formin" -> "Metformin"), since # some transformers versions leave WordPiece subwords as separate entities. # Eager `.to("cuda")` on the underlying model so ZeroGPU packs weights to # disk at startup and streams them into VRAM on the first @spaces.GPU call. # (A transformers `Pipeline` has no `.to()`; move `pipeline.model` instead.) _nlp = pipeline( "token-classification", model=MODEL_ID, aggregation_strategy="simple", ) _nlp.model = _nlp.model.to("cuda") # Stable label -> display label map for the legend / colorblind grouping. # `entity_group` from the pipeline strips the B-/I- prefix (e.g. "DRUG_NAME"). ENTITY_TYPES = [ "DRUG_NAME", "DOSAGE", "STRENGTH", "ROUTE", "FREQUENCY", "DURATION", "FORM", "DRUG_STATUS", "REASON", "ADVERSE_REACTION", "NDC_CODE", "RxNorm_CODE", ] # Gradio HighlightedText color palette keyed by entity group. A fixed palette # makes the highlighted text and legend visually consistent across calls. ENTITY_COLORS = { "DRUG_NAME": "#2563eb", # blue "DOSAGE": "#0891b2", # cyan "STRENGTH": "#0d9488", # teal "ROUTE": "#16a34a", # green "FREQUENCY": "#65a30d", # lime "DURATION": "#ca8a04", # amber "FORM": "#d97706", # orange "DRUG_STATUS": "#dc2626", # red "REASON": "#9333ea", # purple "ADVERSE_REACTION": "#db2777", # pink "NDC_CODE": "#475569", # slate "RxNorm_CODE": "#6b7280", # gray } def _merge_entities(text: str, entities): """Post-merge adjacent same-group entity spans whose offsets touch. `aggregation_strategy="simple"` groups tokens by entity type but, on some transformers versions, leaves WordPiece `##` subwords as separate spans (e.g. "met" at [0:3] and "##formin" at [3:10] both tagged DRUG_NAME). We merge any two consecutive entities that share the same `entity_group` and whose character offsets are contiguous (end == start, allowing zero gap) — reconstructing "Metformin" from the pieces. Scores are averaged. """ if not entities: return entities merged = [dict(entities[0])] for ent in entities[1:]: prev = merged[-1] if (ent["entity_group"] == prev["entity_group"] and ent["start"] <= prev["end"]): # Contiguous same-group spans -> extend the previous one. prev["end"] = ent["end"] prev["word"] = text[prev["start"]:prev["end"]] prev["score"] = (prev["score"] + ent["score"]) / 2 else: merged.append(dict(ent)) return merged def _entities_to_highlights(text: str, entities): """Convert merged entities to Gradio HighlightedText tuples. HighlightedText expects a list of (entity_label, token_text) pairs plus a trailing (None, remaining_text) for unannotated tail. We rebuild the text from the entity spans so that any inter-entity whitespace is emitted as a None-labeled run, keeping the highlighted output readable. """ out = [] cursor = 0 for ent in entities: start = ent["start"] end = ent["end"] group = ent["entity_group"] if start > cursor: out.append((None, text[cursor:start])) out.append((group, text[start:end])) cursor = end if cursor < len(text): out.append((None, text[cursor:])) return out @spaces.GPU(duration=20) def extract_medication_entities(text: str): """Extract medication and drug entities from clinical text. Recognizes 12 medication entity types — DRUG_NAME, DOSAGE, STRENGTH, ROUTE, FREQUENCY, DURATION, FORM, DRUG_STATUS, REASON, ADVERSE_REACTION, NDC_CODE, RxNorm_CODE — from free-text clinical narratives such as discharge summaries, progress notes, and medication orders. Args: text: Clinical narrative to annotate. Returns: A tuple of (highlighted_text, entities_table, summary_markdown) suitable for the three Gradio outputs. """ text = (text or "").strip() if not text: return [], [], "Please enter some clinical text to annotate." raw_entities = _nlp(text) entities = _merge_entities(text, raw_entities) highlights = _entities_to_highlights(text, entities) table_rows = [ [ent["entity_group"], text[ent["start"]:ent["end"]], f"{ent['score']:.3f}"] for ent in entities ] counts = {} for ent in entities: counts[ent["entity_group"]] = counts.get(ent["entity_group"], 0) + 1 if counts: summary_lines = [f"**{count}** × `{name}`" for name, count in sorted(counts.items(), key=lambda kv: -kv[1])] summary_md = "### Extraction summary\n\n" + "\n".join(summary_lines) summary_md += f"\n\n**Total entities found:** {len(entities)}" else: summary_md = "### Extraction summary\n\nNo medication entities detected in the provided text." return highlights, table_rows, summary_md CSS = """ #col-container { max-width: 1100px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } .legend { display: flex; flex-wrap: wrap; gap: 8px; margin-top: 4px; } .legend-chip { display: inline-flex; align-items: center; gap: 4px; font-size: 12px; padding: 2px 8px; border-radius: 999px; color: #fff; } """ # Build the legend HTML once from ENTITY_COLORS. legend_html = '
' + "".join( f'{name}' for name, color in ENTITY_COLORS.items() ) + "
" with gr.Blocks() as demo: gr.Markdown( "# CliniGuard Medication NER\n" "Extract medication and drug entities (drug names, dosages, strengths, " "routes, frequencies, durations, forms, status, reasons, adverse " "reactions, NDC/RxNorm codes) from unstructured clinical text. " "Built on [genzeonplatform/cliniguard-medication-ner]" f"(https://huggingface.co/{MODEL_ID}) — a Bio_ClinicalBERT " "fine-tuned for 12 medication entity types." ) with gr.Column(elem_id="col-container"): with gr.Row(): text_in = gr.Textbox( label="Clinical text", placeholder="Paste a discharge summary, medication order, or clinical note…", lines=8, scale=4, ) run_btn = gr.Button("Extract entities", variant="primary", scale=1) highlighted = gr.HighlightedText( label="Annotated text", combine_adjacent=True, color_map=ENTITY_COLORS, show_legend=True, ) gr.HTML(legend_html) with gr.Row(): with gr.Column(scale=2): summary_md = gr.Markdown(label="Summary") with gr.Column(scale=3): entities_table = gr.Dataframe( headers=["Entity type", "Text", "Confidence"], datatype=["str", "str", "str"], label="Extracted entities", wrap=True, ) gr.Examples( examples=[ ["Discharge medications: Continue Metformin 500 mg tablet by mouth twice daily for diabetes. New: Amoxicillin 500 mg capsule PO TID for 7 days for sinusitis. Discontinue Lisinopril due to persistent cough."], ["Start Atorvastatin 40 mg tablet PO at bedtime for high cholesterol."], ["ADR: Patient developed rash after Penicillin IV. Drug discontinued."], ["Albuterol 90 mcg inhaler 2 puffs inhaled Q4-6H PRN for bronchospasm."], ["MAR: Administered Vancomycin 1 g IV Q12H. NDC: 00409-6509-01."], ], inputs=text_in, outputs=[highlighted, entities_table, summary_md], fn=extract_medication_entities, cache_examples=True, cache_mode="lazy", ) run_btn.click( fn=extract_medication_entities, inputs=text_in, outputs=[highlighted, entities_table, summary_md], api_name="extract", ) if __name__ == "__main__": demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)