# -*- coding: utf-8 -*-
"""Chinese NER Model Demo with Gradio.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1O5yUjNXwH617SPEYFnu-pCNGdZVwb2sQ

Goals:
1. Demo any model with Gradio
2. Make it accessible inside local network
3. Display output nicely (highlight NER instead of raw dict)
"""

#!pip install gradio
#!pip install transformers

import gradio as gr
from gradio.mix import Parallel

from transformers import pipeline

"""#### Load Chinese NER gradio interface"""

# prepend huggingface to model name
# e.g., bert-base-uncased -> huggingface/bert-base-uncased
chinese_ner = gr.Interface.load("huggingface/ckiplab/bert-base-chinese-ner", title="Chinese NER")

"""#### Load Chinese Environment Topic Classification with gradio interface"""

chinese_cat = gr.Interface.load("huggingface/celtics1863/env-bert-topic",inputs="textbox", title="Chinese Environment Topic Classification")

"""#### Putting it together"""

Parallel(chinese_ner, chinese_cat)

Parallel(chinese_ner, chinese_cat).launch(debug=True)

# set visible to local network
# Parallel(chinese_ner, chinese_cat).launch(server_name="0.0.0.0", server_port=7594, debug=True)

"""--- 

##### Noticed that this chinese NER model uses BIES tag and setting `aggregation_strategy="simple"` coundn't fix/output what we want. So we write a simple aggregate function to group them together.
"""

test_chinese_pipeline = pipeline("ner", "ckiplab/bert-base-chinese-ner", aggregation_strategy="simple")

chinese_article = "（早报讯）促进人口健康，有赖于社区伙伴与各机构紧密合作，也须要调整现有的心态、改变照料病患的方式。新推出的“健康SG”计划把重心从疾病护理转向医疗保健，卫生部长王乙康说，若做对了，可在10年内提升人们的生活质量。  王乙康今天中午（5月7日）在加冷综合诊疗所开幕仪式上致辞时说，对比在医院里进行的严重疾病治疗，在促进人口健康方面，所有行动应集中在社区内不同场所，并融入到人们的生活和习惯当中。  因此，关键在于让社区伙伴和各机构无缝连接起来，为居民和病患提供所需支援，而这主要能从几个方面着手。  王乙康说，首先是要确保数据和信息流通，从医院到综合诊疗所、诊所乃至社区保健站，做到信息共享才能达成无缝护理。因此全国电子健康记录至关重要。  其次是确保业者提供的服务不会出现太大差距，居民无论选择哪个业者，都能获得几乎一致的援助。这可从针对常见和慢性疾病拟定详细的护理流程做起"

def custom_agg(output):
    """Custom agg func to group E tag into previous entities. Note that the input
    is the output of aggregation_strategy="simple", which group only B-tag, I-tag
    but missed out E-tag.
    
    Output of aggregation_strategy="none":
    {'end': 77, 'entity': 'B-ORG', 'index': 76, 'score': 0.99999905, 'start': 76, 'word': '卫'},
    {'end': 78, 'entity': 'I-ORG', 'index': 77, 'score': 0.9999982, 'start': 77, 'word': '生'},
    {'end': 79, 'entity': 'I-ORG', 'index': 78, 'score': 0.99999845, 'start': 78,'word': '部'},
    {'end': 80, 'entity': 'E-ORG', 'index': 79, 'score': 0.9999988,  'start': 79, 'word': '长'},

    Output of aggregation_strategy="simple":
    {'end': 79, 'entity_group': 'ORG', 'score': 0.99999857, 'start': 76, 'word': '卫 生 部'},
    {'end': 80, 'entity_group': 'ORG', 'score': 0.9999988, 'start': 79, 'word': '长'},

    Output of custom_agg:
    {'end': 80, 'entity_group': 'ORG', 'score': 0.99999857, 'start': 76, 'word': '卫 生 部长'},

    returns:
      list of tuple(substring, entity) for gradio.HighlightedText component
    """
    res = []
    for current, next in zip(output, output[1:]):
        same_entity_group = (current["entity_group"] == next["entity_group"])
        neighbour_index = (current["end"] == next["start"])
        if same_entity_group and neighbour_index:
            current["word"] += next["word"]
            current["end"] = next["end"]
            res.append(current)

    return [(item["word"], item["entity_group"]) for item in res]

def _predict_ner(text):
    output = test_chinese_pipeline(text)
    output = custom_agg(output)
    return output

demo = gr.Interface(fn=_predict_ner, inputs="text", outputs=gr.outputs.HighlightedText(label="Output"))

_predict_ner(chinese_article)

demo.launch(debug=True)

"""Reference:
1. [Gradio docs](https://gradio.app/docs/)
2. [Gradio publish network](https://github.com/gradio-app/gradio/issues/69)
3. [Gradio x Huggingface Blog](https://huggingface.co/blog/gradio)
"""