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',
            '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()