app.py

import gradio as gr
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForTokenClassification
import warnings
warnings.filterwarnings("ignore")

MODEL_OPTIONS = {
    "Production — BioBERT (NCBI Disease)": "ugaray96/biobert_ncbi_disease_ner",
    "Custom — Bio_ClinicalBERT (MedMentions / SNOMED)": "acebirim/snomed-ner-model",
}
DEFAULT_MODEL_LABEL = "Production — BioBERT (NCBI Disease)"

LABEL_LIST = ["O", "B-DISEASE", "I-DISEASE"]
ID2LABEL = {i: label for i, label in enumerate(LABEL_LIST)}

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Loaded models are cached here so switching back to a model already
# used in this session doesn't require reloading it from the Hub.
loaded_models = {}


def get_model(model_label: str):
    model_id = MODEL_OPTIONS[model_label]
    if model_id not in loaded_models:
        print(f"Loading model from {model_id}...")
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        model = AutoModelForTokenClassification.from_pretrained(model_id)
        model.eval()
        model.to(device)
        loaded_models[model_id] = (tokenizer, model)
        print(f"Model loaded on {device}")
    return loaded_models[model_id]


# Pre-load the default model so the first request isn't slow
get_model(DEFAULT_MODEL_LABEL)

STOPWORD_BLOCKLIST = {
    'patient', 'subject', 'presents', 'with', 'of', 'the', 'a', 'an', 'and',
    'or', 'but', 'in', 'on', 'at', 'to', 'for', 'from', 'by', 'as', 'was',
    'were', 'is', 'are', 'been', 'being', 'have', 'has', 'had', 'do', 'does',
    'did', 'will', 'would', 'could', 'should', 'may', 'might', 'must', 'can',
    'diagnosed', 'biopsy', 'admitted', 'he', 'history', 'referred', 'presented',
    'she', 'staging', 'systemic', 'acute', 'chronic', 'severe', 'bilateral',
    'complications', 'complication', 'imaging', 'progressive'
}


def predict(text: str, model_label: str):
    if not text.strip():
        return [], "No text provided."

    tokenizer, model = get_model(model_label)

    tokens = text.split()
    tokenized = tokenizer(
        tokens,
        is_split_into_words=True,
        truncation=True,
        padding="max_length",
        max_length=128,
        return_tensors="pt"
    )

    with torch.no_grad():
        outputs = model(**{k: v.to(device) for k, v in tokenized.items()})
        logits = outputs.logits[0]
        probs = F.softmax(logits, dim=-1)
        predictions = logits.argmax(dim=-1)

    word_ids = tokenized.word_ids()
    predicted_labels = []
    predicted_probs = []
    previous_word_idx = None
    for word_idx, pred_id, prob in zip(word_ids, predictions, probs):
        if word_idx is None:
            continue
        if word_idx != previous_word_idx:
            predicted_labels.append(ID2LABEL[pred_id.item()])
            predicted_probs.append(prob[pred_id.item()].item())
        previous_word_idx = word_idx

    entities = []
    current_entity_tokens = []
    current_label = None
    current_probs = []

    for token, label, prob in zip(tokens, predicted_labels, predicted_probs):
        if label.startswith("B-"):
            if current_entity_tokens:
                entity_text = " ".join(current_entity_tokens)
                avg_conf = sum(current_probs) / len(current_probs)
                if entity_text.lower() not in STOPWORD_BLOCKLIST:
                    entities.append({"text": entity_text, "label": current_label, "confidence": avg_conf})
            current_entity_tokens = [token]
            current_label = label.replace("B-", "")
            current_probs = [prob]
        elif label.startswith("I-") and current_entity_tokens:
            current_entity_tokens.append(token)
            current_probs.append(prob)
        else:
            if current_entity_tokens:
                entity_text = " ".join(current_entity_tokens)
                avg_conf = sum(current_probs) / len(current_probs)
                if entity_text.lower() not in STOPWORD_BLOCKLIST:
                    entities.append({"text": entity_text, "label": current_label, "confidence": avg_conf})
            current_entity_tokens = []
            current_label = None
            current_probs = []

    if current_entity_tokens:
        entity_text = " ".join(current_entity_tokens)
        avg_conf = sum(current_probs) / len(current_probs)
        if entity_text.lower() not in STOPWORD_BLOCKLIST:
            entities.append({"text": entity_text, "label": current_label, "confidence": avg_conf})

    # Strip trailing punctuation
    entities = [{"text": e["text"].rstrip(".,;:!?"), "label": e["label"], "confidence": e["confidence"]} for e in entities]
    entities = [e for e in entities if e["text"].lower() not in STOPWORD_BLOCKLIST]

    # Highlighted text
    clean_tokens = [t.rstrip(".,;:!?") for t in tokens]
    entity_set = {e["text"]: e["label"] for e in entities}
    highlighted = []
    i = 0
    while i < len(clean_tokens):
        matched = False
        for length in range(min(10, len(clean_tokens) - i), 0, -1):
            span = " ".join(clean_tokens[i:i+length])
            if span in entity_set:
                highlighted.append((span, entity_set[span]))
                i += length
                matched = True
                break
        if not matched:
            highlighted.append((tokens[i] + " ", None))
            i += 1

    # Markdown table with confidence
    if entities:
        md = "| Entity | Type | Confidence |\n|--------|------|------------|\n"
        for e in entities:
            conf_pct = f"{e['confidence']*100:.1f}%"
            md += f"| {e['text']} | {e['label']} | {conf_pct} |\n"
    else:
        md = "No disease entities detected."

    return highlighted, md


EXAMPLES = [
    ["Patient presents with hypertension and type 2 diabetes mellitus."],
    ["History includes breast cancer treated with chemotherapy."],
    ["Subject presents with symptoms of asthma and chronic obstructive pulmonary disease."],
    ["The patient was diagnosed with pneumonia and required hospitalization."],
    ["Patient suffers from depression and anxiety disorder following myocardial infarction."],
]

with gr.Blocks(title="Clinical Disease Entity Extractor") as demo:

    gr.Markdown("""
    # 🏥 Clinical Disease Entity Extractor

    Paste any clinical free-text below and pick a model. The model will identify **disease entities** and return a confidence score for each.

    - **Production — BioBERT (NCBI Disease)**: [`ugaray96/biobert_ncbi_disease_ner`](https://huggingface.co/ugaray96/biobert_ncbi_disease_ner)
    - **Custom — Bio_ClinicalBERT (MedMentions / SNOMED)**: [`acebirim/snomed-ner-model`](https://huggingface.co/acebirim/snomed-ner-model)
    """)

    with gr.Row():
        with gr.Column(scale=1):
            model_dropdown = gr.Dropdown(
                choices=list(MODEL_OPTIONS.keys()),
                value=DEFAULT_MODEL_LABEL,
                label="Model",
            )
            text_input = gr.Textbox(
                label="Clinical Text",
                placeholder="e.g. Patient presents with hypertension and type 2 diabetes mellitus.",
                lines=6,
            )
            run_btn = gr.Button("Extract Entities", variant="primary")
            gr.Examples(examples=EXAMPLES, inputs=text_input, label="Try an example")

        with gr.Column(scale=1):
            highlighted_output = gr.HighlightedText(
                label="Highlighted Entities",
                combine_adjacent=True,
                show_legend=True,
            )
            markdown_output = gr.Markdown()

    run_btn.click(fn=predict, inputs=[text_input, model_dropdown], outputs=[highlighted_output, markdown_output])
    text_input.submit(fn=predict, inputs=[text_input, model_dropdown], outputs=[highlighted_output, markdown_output])

    gr.Markdown("""
    ---
    **Confidence Score**: Average softmax probability across entity tokens. Scores below 50% indicate the model is not plurality-confident in its DISEASE prediction — a natural decision boundary for a 3-class problem — a useful signal for monitoring model reliability in production.

    *Built as a capstone project for the Advanced ML course from https://ml.electricsheep.africa/grade2/.*
    """)

demo.launch()