Cleanup and re-add rating

app.py

@@ -6,8 +6,12 @@ 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")
+MODELS = ["d4data/biomedical-ner-all", "samrawal/bert-base-uncased_clinical-ner"]
+
+current_model = MODELS[1]
+
+tokenizer = AutoTokenizer.from_pretrained(current_model)
+model = AutoModelForTokenClassification.from_pretrained(current_model)
 
 plt.switch_backend("Agg")
 
@@ -43,23 +47,26 @@ def run_ner(text):
             for x in raw
         ],
     }
-
+    label = examples.get(text, None)
     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
+    return label, 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])
+    submit = gr.Button("Submit")
+    # with gr.Accordion("Examples", open=False):
+    example_dropdown = gr.Dropdown(label="Examples", choices=list(examples.keys()))
+    example_dropdown.change(
+        lambda x: gr.Textbox.update(value=x), inputs=example_dropdown, outputs=note
+    )
+    rating = gr.Label(label="Given rating")
+    highlight = gr.HighlightedText(label="NER", combine_adjacent=True)
+    table = gr.Dataframe(headers=["Entity", "Count"])
+    plot = gr.Plot(label="Bar")
+    submit.click(run_ner, [note], [rating, highlight, table, plot])
+    note.submit(run_ner, [note], [rating, highlight, table, plot])
 
 demo.launch()