Spaces:
Running
Running
| import json | |
| from collections import defaultdict, Counter | |
| import matplotlib.pyplot as plt | |
| import gradio as gr | |
| import pandas as pd | |
| from transformers import pipeline | |
| plt.switch_backend("Agg") | |
| examples = [] | |
| with open("examples.json", "r") as f: | |
| content = json.load(f) | |
| examples = [f"{x['label']}: {x['text']}" for x in content] | |
| pipe = pipeline( | |
| "ner", | |
| model="Clinical-AI-Apollo/Medical-NER", | |
| aggregation_strategy="simple", | |
| ) | |
| def plot_to_figure(grouped): | |
| fig = plt.figure() | |
| plt.bar(x=list(grouped.keys()), height=list(grouped.values())) | |
| plt.margins(0.2) | |
| plt.subplots_adjust(bottom=0.4) | |
| plt.xticks(rotation=90) | |
| return fig | |
| def run_ner(text): | |
| raw = pipe(text) | |
| ner_content = { | |
| "text": text, | |
| "entities": [ | |
| { | |
| "entity": x["entity_group"], | |
| "word": x["word"], | |
| "score": x["score"], | |
| "start": x["start"], | |
| "end": x["end"], | |
| } | |
| for x in raw | |
| ], | |
| } | |
| 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") | |
| submit = gr.Button("Submit") | |
| # with gr.Accordion("Examples", open=False): | |
| example_dropdown = gr.Dropdown(label="Examples", choices=examples) | |
| example_dropdown.change(lambda x: x, inputs=example_dropdown, outputs=note) | |
| highlight = gr.HighlightedText(label="NER", combine_adjacent=True) | |
| table = gr.Dataframe(headers=["Entity", "Count"]) | |
| plot = gr.Plot(label="Bar") | |
| submit.click(run_ner, [note], [highlight, table, plot]) | |
| note.submit(run_ner, [note], [highlight, table, plot]) | |
| demo.launch() | |