David Kagramanyan
updated ner display
37516c9
raw
history blame
2.56 kB
import gradio as gr
import os
import requests
from spacy import displacy
os.system("python -m spacy download en_core_web_md")
import spacy
colors = {
"Observation": "#9bddff",
"Evaluation": "#f08080",
}
nlp = spacy.load("en_core_web_md") #Esto es para usar displacy y renderizar las entidades
nlp.disable_pipes("ner")
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)
entities = response.json()
doc = nlp(input_text_message)
potential_entities = []
for entity in entities:
start = entity["start"]
end = entity["end"]
label = entity["entity"]
if label == "I-Observation" or label == "B-Observation":
label = "Observation"
if label == "I-Evaluation" or label == "B-Evaluation":
label = "Evaluation"
entity["entity"]=label
ent = doc.char_span(start, end, label=label)
if ent != None:
doc.ents += (ent,)
else:
potential_entities.append(entity)
potential_entities.append({"entity": "NONE", "start": -1, "end": -1})
start = potential_entities[0]["start"]
end = potential_entities[0]["end"]
label = potential_entities[0]["entity"]
for item in potential_entities:
if item["entity"] == label and item["start"] == end:
end = item["end"]
continue
else:
if item["start"] != start:
ent = doc.char_span(start, end, label=label)
doc.ents += (ent,)
start = item["start"]
end = item["end"]
label = item["entity"]
options = {"ents": colors.keys(), "colors": colors}
return displacy.render(doc, style="ent", options=options)
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",
label='Check your text for compliance with the NVC rules'),
outputs="html", examples=examples)
iface.launch()