Upload 3 files

README.md

@@ -1,13 +1,29 @@
 ---
-title: SNOMED NER
-emoji: 👁
-colorFrom: red
-colorTo: yellow
+title: Snomed NER
+emoji: 🏥
+colorFrom: blue
+colorTo: indigo
 sdk: gradio
-sdk_version: 6.16.0
-python_version: '3.13'
 app_file: app.py
 pinned: false
+license: apache-2.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SNOMED NER — Clinical Disease Entity Extractor
+
+A Named Entity Recognition model for identifying disease entities in clinical free-text,
+fine-tuned on the MedMentions corpus using Bio_ClinicalBERT.
+
+## Model
+
+- **Hub**: [`acebirim/snomed-ner-model`](https://huggingface.co/acebirim/snomed-ner-model)
+- **Base model**: `emilyalsentzer/Bio_ClinicalBERT`
+- **Task**: Token classification (NER) — `O`, `B-DISEASE`, `I-DISEASE`
+- **Training data**: MedMentions (132K PubMed abstracts with UMLS annotations)
+
+## How it works
+
+Paste any clinical free-text into the form. The model identifies disease entity spans
+using BIO tagging and highlights them inline. Results are also shown as a summary table.
+
+*Built as part of an NHS/OneLondon Secure Data Environment capstone project.*

app.py

@@ -1,7 +1,154 @@
 import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModelForTokenClassification
+import warnings
+warnings.filterwarnings("ignore")
 
-def greet(name):
-    return "Hello " + name + "!!"
+MODEL_ID = "acebirim/snomed-ner-model"
+LABEL_LIST = ["O", "B-DISEASE", "I-DISEASE"]
+ID2LABEL = {i: label for i, label in enumerate(LABEL_LIST)}
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+print(f"Loading model from {MODEL_ID}...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForTokenClassification.from_pretrained(MODEL_ID)
+model.eval()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+print(f"Model loaded on {device}")
+
+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'
+}
+
+def predict(text: str):
+    if not text.strip():
+        return [], "No text provided."
+
+    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()})
+        predictions = outputs.logits.argmax(dim=-1)[0]
+
+    word_ids = tokenized.word_ids()
+    predicted_labels = []
+    previous_word_idx = None
+    for word_idx, pred_id in zip(word_ids, predictions):
+        if word_idx is None:
+            continue
+        if word_idx != previous_word_idx:
+            predicted_labels.append(ID2LABEL[pred_id.item()])
+        previous_word_idx = word_idx
+
+    entities = []
+    current_entity_tokens = []
+    current_label = None
+
+    for token, label in zip(tokens, predicted_labels):
+        if label.startswith("B-"):
+            if current_entity_tokens:
+                entity_text = " ".join(current_entity_tokens)
+                if entity_text.lower() not in STOPWORD_BLOCKLIST:
+                    entities.append({"text": entity_text, "label": current_label})
+            current_entity_tokens = [token]
+            current_label = label.replace("B-", "")
+        elif label.startswith("I-") and current_entity_tokens:
+            current_entity_tokens.append(token)
+        else:
+            if current_entity_tokens:
+                entity_text = " ".join(current_entity_tokens)
+                if entity_text.lower() not in STOPWORD_BLOCKLIST:
+                    entities.append({"text": entity_text, "label": current_label})
+            current_entity_tokens = []
+            current_label = None
+
+    if current_entity_tokens:
+        entity_text = " ".join(current_entity_tokens)
+        if entity_text.lower() not in STOPWORD_BLOCKLIST:
+            entities.append({"text": entity_text, "label": current_label})
+
+    # Highlighted text
+    entity_set = {e["text"]: e["label"] for e in entities}
+    highlighted = []
+    i = 0
+    while i < len(tokens):
+        matched = False
+        for length in range(min(10, len(tokens) - i), 0, -1):
+            span = " ".join(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
+    if entities:
+        md = "| Entity | Type |\n|--------|------|\n"
+        for e in entities:
+            md += f"| {e['text']} | {e['label']} |\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
+    **Model**: [`acebirim/snomed-ner-model`](https://huggingface.co/acebirim/snomed-ner-model)
+    &nbsp;|&nbsp; **Base**: Bio_ClinicalBERT &nbsp;|&nbsp; **Trained on**: MedMentions (132K examples)
+
+    Paste any clinical free-text below. The model will identify **disease entities** within it.
+    """)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            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, outputs=[highlighted_output, markdown_output])
+    text_input.submit(fn=predict, inputs=text_input, outputs=[highlighted_output, markdown_output])
+
+    gr.Markdown("""
+    ---
+    Fine-tuned on [MedMentions](https://github.com/chanzuckerberg/MedMentions) using `emilyalsentzer/Bio_ClinicalBERT`.
+    *Built as part of an NHS/OneLondon Secure Data Environment capstone project.*
+    """)
+
+demo.launch()

requirements.txt

@@ -0,0 +1,2 @@
+torch>=2.0.0
+transformers>=4.35.0