# REF: https://gradio.app/named_entity_recognition/
from transformers import pipeline

import gradio as gr


model_name="xlm-roberta-base"
# model_name="roberta-large"  
from transformers import AutoTokenizer, AutoModelForTokenClassification


label_list= ['literal',"metaphoric"]
label_dict_relations={ i : l for i, l in enumerate(label_list) }
PATH = "./saved-models/my_model"
model_metaphor_detection = AutoModelForTokenClassification.from_pretrained(PATH, id2label=label_dict_relations)
tokenizer = AutoTokenizer.from_pretrained(model_name)

pipeline_metaphors=pipeline(
    "ner", 
    model=model_metaphor_detection, 
    tokenizer=tokenizer, 
    aggregation_strategy="none",
    # aggregation_strategy="simple",
    )


examples = [
    "It would change the trajectory of your legal career.",
    "Washington and the media just explodes on you, you just don’t know where you are at the moment", 
    "Those statements are deeply concerning.", 
]

# Demo usage
import pprint
detection_results = pipeline_metaphors("It would change the trajectory of your legal career.")
pp = pprint.PrettyPrinter(indent=4)
pp.pprint(detection_results)

"""Example Output; aggregation_strategy="none"
[   {   'end': 2,
        'entity': 'literal',
        'index': 1,
        'score': 0.99981445,
        'start': 0,
        'word': '▁It'},
    {   'end': 8,
        'entity': 'literal',
        'index': 2,
        'score': 0.9999882,
        'start': 3,
        'word': '▁would'},
    {   'end': 15,
        'entity': 'literal',
        'index': 3,
        'score': 0.6243065,
        'start': 9,
        'word': '▁change'},
    {   'end': 19,
        'entity': 'literal',
        'index': 4,
        'score': 0.9999826,
        'start': 16,
        'word': '▁the'},
    {   'end': 27,
        'entity': 'metaphoric',
        'index': 5,
        'score': 0.99631363,
        'start': 20,
        'word': '▁traject'},
    {   'end': 30,
        'entity': 'metaphoric',
        'index': 6,
        'score': 0.9979997,
        'start': 27,
        'word': 'ory'},
    {   'end': 33,
        'entity': 'literal',
        'index': 7,
        'score': 0.9996278,
        'start': 31,
        'word': '▁of'},
    {   'end': 38,
        'entity': 'literal',
        'index': 8,
        'score': 0.99985147,
        'start': 34,
        'word': '▁your'},
    {   'end': 44,
        'entity': 'literal',
        'index': 9,
        'score': 0.99984956,
        'start': 39,
        'word': '▁legal'},
    {   'end': 51,
        'entity': 'literal',
        'index': 10,
        'score': 0.998919,
        'start': 45,
        'word': '▁career'},
    {   'end': 52,
        'entity': 'literal',
        'index': 11,
        'score': 0.99775606,
        'start': 51,
        'word': '.'}]
"""


"""Example Output; aggregation_strategy="simple"
[   {   'end': 19,
    'entity_group': 'literal',
    'score': 0.9060229,
    'start': 0,
    'word': 'It would change the'},
{   'end': 30,
    'entity_group': 'metaphoric',
    'score': 0.9971567,
    'start': 20,
    'word': 'trajectory'},
{   'end': 52,
    'entity_group': 'literal',
    'score': 0.9992008,
    'start': 31,
    'word': 'of your legal career.'}]

"""
# exit(0)

def ner(text):
    output = pipeline_metaphors(text)
    # # change name 
    for x in output:
        if 'entity_group' in x:
            x['entity'] = x['entity_group']
    return {"text": text, "entities": output}    



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

demo.launch()