WIP

app.py

@@ -1,23 +1,33 @@
 import json
+from collections import defaultdict
 
 import gradio as gr
+import pandas as pd
 from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassification
 
 tokenizer = AutoTokenizer.from_pretrained("d4data/biomedical-ner-all")
 model = AutoModelForTokenClassification.from_pretrained("d4data/biomedical-ner-all")
 
-EXAMPLE_TEXTS = []
+
+EXAMPLE_MAP = {}
 with open("examples.json", "r") as f:
     example_json = json.load(f)
-    EXAMPLE_TEXTS = [x["text"] for x in example_json]
-
+    EXAMPLE_MAP = {x["text"]: x["label"] for x in example_json}
 
 pipe = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
 
 
+def group_by_entity(raw):
+    out = defaultdict(int)
+    for ent in raw:
+        out[ent["entity_group"]] += 1
+    out["total"] = sum(out.values())
+    return out
+
+
 def ner(text):
     raw = pipe(text)
-    result = {
+    ner_content = {
         "text": text,
         "entities": [
             {
@@ -30,14 +40,22 @@ def ner(text):
             for x in raw
         ],
     }
-    return result, {}
+    grouped = group_by_entity(raw)
+    df = pd.DataFrame({"Entity": grouped.keys(), "Count": grouped.values()})
+    label = EXAMPLE_MAP.get(text, None)
+    return (ner_content, grouped, label, df.hist())
 
 
 interface = gr.Interface(
     ner,
-    inputs=gr.Textbox(label="Input", value=""),
-    outputs=[gr.HighlightedText(combine_adjacent=True), "json"],
-    examples=EXAMPLE_TEXTS,
+    inputs=gr.Textbox(label="Note text", value=""),
+    outputs=[
+        gr.HighlightedText(label="NER", combine_adjacent=True),
+        gr.JSON(label="Entity Counts"),
+        gr.Label(label="Rating"),
+        "plot",
+    ],
+    examples=list(EXAMPLE_MAP.keys()),
 )
 
 interface.launch()

requirements.txt

@@ -60,6 +60,7 @@ sniffio==1.3.0
 starlette==0.20.4
 tokenizers==0.12.1
 tomli==2.0.1
+torch==1.12.1
 tqdm==4.64.1
 transformers==4.22.2
 typing_extensions==4.4.0