WIP

app.py

@@ -1,6 +1,7 @@
 import json
 from collections import defaultdict
 
+import matplotlib.pyplot as plt
 import gradio as gr
 import pandas as pd
 from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassification
@@ -8,6 +9,7 @@ from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassificatio
 tokenizer = AutoTokenizer.from_pretrained("d4data/biomedical-ner-all")
 model = AutoModelForTokenClassification.from_pretrained("d4data/biomedical-ner-all")
 
+plt.switch_backend("Agg")
 
 EXAMPLE_MAP = {}
 with open("examples.json", "r") as f:
@@ -21,10 +23,17 @@ def group_by_entity(raw):
     out = defaultdict(int)
     for ent in raw:
         out[ent["entity_group"]] += 1
-    out["total"] = sum(out.values())
+    # 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()))
+    plt.xticks(rotation=90)
+    return fig
+
+
 def ner(text):
     raw = pipe(text)
     ner_content = {
@@ -41,9 +50,16 @@ def ner(text):
         ],
     }
     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())
+    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(
@@ -53,9 +69,10 @@ interface = gr.Interface(
         gr.HighlightedText(label="NER", combine_adjacent=True),
         gr.JSON(label="Entity Counts"),
         gr.Label(label="Rating"),
-        "plot",
+        gr.Plot(label="Bar"),
     ],
     examples=list(EXAMPLE_MAP.keys()),
+    allow_flagging="never",
 )
 
 interface.launch()