import gradio as gr
import os

import requests

from spacy import displacy


def compute_ner(input_text_message):
    endpoint_url = 'https://on1m82uknekghqeh.us-east-1.aws.endpoints.huggingface.cloud'

    headers = {
        'Authorization': 'Bearer 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',
        'Content-Type': 'application/json',
    }

    json_data = {
        'inputs': input_text_message,
    }

    response = requests.post(endpoint_url, headers=headers, json=json_data)

    tokens = response.json()

    entities = []

    for token in tokens:
        label = token["entity"]

        if label == "I-Observation" or label == "B-Observation":
            label = "Observation"
            token["label"] = label
            entities.append(token)

        if label == "I-Evaluation" or label == "B-Evaluation":
            label = "Evaluation"
            token["label"] = label
            entities.append(token)

    params = [{"text": input_text_message,
               "ents": entities,
               "title": None}]

    return displacy.render(params, style="ent", manual=True, options={
        "colors": {
            "Observation": "#9bddff",
            "Evaluation": "#f08080",
        },
    })


examples = ['You are dick',
            'My dad is an asshole and took his anger out on my mom by verbally abusing her and when she left he eventually moved on to my brother']

iface = gr.Interface(fn=compute_ner, inputs=gr.inputs.Textbox(lines=5, placeholder="Enter your text here"),
                     outputs="html", examples=examples)
iface.launch()