# coding=utf-8
# author: xusong <xusong28@jd.com>
# time: 2022/8/25 16:57


"""

## ner demo
- https://gradio.app/named_entity_recognition/
- https://huggingface.co/dslim/bert-base-NER?text=My+name+is+Wolfgang+and+I+live+in+Berlin

"""

from transformers import pipeline

import gradio as gr

ner_pipeline = pipeline("ner")

examples = [
    "Does Chicago have any stores and does Joe live here?",
]

import json

def ner(text):
    output = ner_pipeline(text)
    for ent in output:
        ent["score"] = float(ent["score"])
    aa = {"text": text, "entities": output}
    return aa, output


demo = gr.Interface(
    ner,
    inputs=gr.Textbox(placeholder="Enter sentence here..."),
    outputs=
    [
        gr.HighlightedText(
            label="NER",
            show_legend=True,
        ),
        gr.JSON(),
    ],
    examples=examples)

demo.launch()