updated ner display

app.py

@@ -5,6 +5,17 @@ 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'
@@ -20,33 +31,53 @@ def compute_ner(input_text_message):
 
     response = requests.post(endpoint_url, headers=headers, json=json_data)
 
-    tokens = response.json()
+    entities = response.json()
 
-    entities = []
+    doc = nlp(input_text_message)
 
-    for token in tokens:
-        label = token["entity"]
+    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"
-            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",
-        },
-    })
+
+        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',