Cleanup

app.py

@@ -1,5 +1,5 @@
 import json
-from collections import defaultdict
+from collections import defaultdict, Counter
 
 import matplotlib.pyplot as plt
 import gradio as gr
@@ -11,22 +11,14 @@ model = AutoModelForTokenClassification.from_pretrained("d4data/biomedical-ner-a
 
 plt.switch_backend("Agg")
 
-EXAMPLE_MAP = {}
+examples = {}
 with open("examples.json", "r") as f:
-    example_json = json.load(f)
-    EXAMPLE_MAP = {x["text"]: x["label"] for x in example_json}
+    content = json.load(f)
+    examples = {x["text"]: x["label"] for x in content}
 
 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 plot_to_figure(grouped):
     fig = plt.figure()
     plt.bar(x=list(grouped.keys()), height=list(grouped.values()))
@@ -36,7 +28,7 @@ def plot_to_figure(grouped):
     return fig
 
 
-def ner(text):
+def run_ner(text):
     raw = pipe(text)
     ner_content = {
         "text": text,
@@ -51,30 +43,23 @@ def ner(text):
             for x in raw
         ],
     }
-    grouped = group_by_entity(raw)
-    figure = plot_to_figure(grouped)
-    label = EXAMPLE_MAP.get(text, "Unknown")
-
-    meta = {
-        "entity_counts": grouped,
-        "entities": len(set(grouped.keys())),
-        "counts": sum(grouped.values()),
-    }
 
-    return (ner_content, meta, label, figure)
-
-
-interface = gr.Interface(
-    ner,
-    inputs=gr.Textbox(label="Note text", value=""),
-    outputs=[
-        gr.HighlightedText(label="NER", combine_adjacent=True),
-        gr.JSON(label="Entity Counts"),
-        gr.Label(label="Rating"),
-        gr.Plot(label="Bar"),
-    ],
-    examples=list(EXAMPLE_MAP.keys()),
-    allow_flagging="never",
-)
-
-interface.launch()
+    grouped = Counter((x["entity_group"] for x in raw))
+    rows = [[k, v] for k, v in grouped.items()]
+    figure = plot_to_figure(grouped)
+    return ner_content, rows, figure
+
+
+with gr.Blocks() as demo:
+    note = gr.Textbox(label="Note text")
+    with gr.Accordion("Examples", open=False):
+        examples = gr.Examples(examples=list(examples.keys()), inputs=note)
+    with gr.Tab("NER"):
+        highlight = gr.HighlightedText(label="NER", combine_adjacent=True)
+    with gr.Tab("Bar"):
+        plot = gr.Plot(label="Bar")
+    with gr.Tab("Table"):
+        table = gr.Dataframe(headers=["Entity", "Count"])
+    note.submit(run_ner, [note], [highlight, table, plot])
+
+demo.launch()

requirements.txt

@@ -18,7 +18,7 @@ filelock==3.8.0
 fonttools==4.37.4
 frozenlist==1.3.1
 fsspec==2022.8.2
-gradio==3.4.1
+gradio==3.11.0
 h11==0.12.0
 httpcore==0.15.0
 httpx==0.23.0
@@ -60,9 +60,9 @@ sniffio==1.3.0
 starlette==0.20.4
 tokenizers==0.12.1
 tomli==2.0.1
-torch==1.12.1
+torch==1.13.0
 tqdm==4.64.1
-transformers==4.22.2
+transformers==4.24.0
 typing_extensions==4.4.0
 uc-micro-py==1.0.1
 urllib3==1.26.12