Update app.py

app.py

@@ -1,23 +1,23 @@
+import gradio as gr
+import pandas as pd
 import json
 from collections import defaultdict
 
-import matplotlib.pyplot as plt
-import gradio as gr
-import pandas as pd
+# Create tokenizer for biomed model
 from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassification
-
 tokenizer = AutoTokenizer.from_pretrained("d4data/biomedical-ner-all")
 model = AutoModelForTokenClassification.from_pretrained("d4data/biomedical-ner-all")
+pipe = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
 
+# Matplotlib for entity graph
+import matplotlib.pyplot as plt
 plt.switch_backend("Agg")
 
-EXAMPLE_MAP = {}
+# Load examples from JSON
+EXAMPLES = {}
 with open("examples.json", "r") as f:
     example_json = json.load(f)
-    EXAMPLE_MAP = {x["text"]: x["label"] for x in example_json}
-
-pipe = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
-
+    EXAMPLES = {x["text"]: x["label"] for x in example_json}
 
 def group_by_entity(raw):
     out = defaultdict(int)
@@ -51,9 +51,10 @@ def ner(text):
             for x in raw
         ],
     }
+    
     grouped = group_by_entity(raw)
     figure = plot_to_figure(grouped)
-    label = EXAMPLE_MAP.get(text, "Unknown")
+    label = EXAMPLES.get(text, "Unknown")
 
     meta = {
         "entity_counts": grouped,
@@ -73,7 +74,7 @@ interface = gr.Interface(
         gr.Label(label="Rating"),
         gr.Plot(label="Bar"),
     ],
-    examples=list(EXAMPLE_MAP.keys()),
+    examples=list(EXAMPLES.keys()),
     allow_flagging="never",
 )