import gradio as gr
from transformers import pipeline

model_checkpoint = 'zinoubm/bert-finetuned-ner'
model = pipeline(
    "token-classification", model=model_checkpoint,
)


def concat_prediction(prediction):
    entity = prediction[0]['entity'][2:]
    start = prediction[0]['start']
    end = prediction[-1]['end']
    return {
        'entity': entity,
        'start': start,
        'end': end}


def concat_predictions(predictions):
    concatenated_predictions = []
    for_concat = []
    for i in range(len(predictions)):
        if predictions[i]['entity'].startswith('B'):
            for_concat.append(predictions[i])
            j = i+1
            while j < len(predictions) and predictions[j]['entity'].startswith('I'):
                for_concat.append(predictions[j])
                j += 1

            concatenated_predictions.append(concat_prediction(for_concat))
            for_concat = []

    return concatenated_predictions


title = 'Extended Name Entity Recognition'

examples = [
    "Does Chicago have any stores and does Joe live here?",
    "My name is Sylvain and I work at Hugging Face in Brooklyn."
]


article = '''
# How to use this interface
Here's the [notebook](https://colab.research.google.com/drive/1xAsW1YNC38NEHXn0XMZVL_W-VPNYD-TO?usp=sharing) used to train the model used in this app.
Using the interface is very easy, just type some text that and the model will give the names of entities in one of these categories:

- **org** : organization
- **per** : person
- **geo** : location
- **tim** :  dates and times
- **gpe** : Geopolitical Entity
- **art**
- **nat**
- **eve**

just hit **Submit** to see the results.You can also try some of the provided examples.
'''


def predict(text):
    output = model(text)
    return {"text": text, "entities": concat_predictions(output)}


demo = gr.Interface(predict,
                    gr.Textbox(placeholder="Enter sentence here..."),
                    gr.HighlightedText(),
                    title=title,
                    examples=examples,
                    article=article)

demo.launch()