better postprocessing

app.py

@@ -1,8 +1,8 @@
 # -*- coding: utf-8 -*-
 import os
 
-os.system("pip3 install torch==1.10.1+cpu torchvision==0.11.2+cpu torchaudio==0.10.1+cpu -f "
-          "https://download.pytorch.org/whl/cpu/torch_stable.html")
+# os.system("pip3 install torch==1.10.1+cpu torchvision==0.11.2+cpu torchaudio==0.10.1+cpu -f "
+#           "https://download.pytorch.org/whl/cpu/torch_stable.html")
 
 import gradio as gr
 from transformers import pipeline
@@ -29,9 +29,7 @@ ner_pipeline = pipeline(task='ner', model="crabz/slovakbert-ner")
 nlp = spacy.blank("sk")
 
 
-def apply_ner(sentence: str):
-    classifications = ner_pipeline(sentence)
-
+def postprocess(classifications):
     entities = []
     for i in range(len(classifications)):
         if classifications[i]['entity'] != 0:
@@ -41,13 +39,37 @@ def apply_ner(sentence: str):
                     j += 1
                 entities.append((ner_map[classifications[i]['entity']].split('-')[1], classifications[i]['start'],
                                  classifications[j - 1]['end']))
-    doc = nlp(sentence)
+    to_remove = []
+    merged_entities = []
+    for i in range(len(entities)):
+        for j in range(i + 1, len(entities)):
+            if entities[i] != entities[j] and entities[i][0] == entities[j][0] and (entities[i][2] == entities[j][1] or
+                                                                                    entities[i][1] == entities[j][2]):
+                to_remove.append(entities[i])
+                to_remove.append(entities[j])
+
+                new_start = min(entities[i][1], entities[j][1])
+                new_end = max(entities[i][2], entities[j][2])
+                merged_entities.append((entities[i][0], new_start, new_end))
+    for ent in to_remove:
+        entities.remove(ent)
+    entities += merged_entities
+    return entities
 
+
+def set_entities(sentence, entities):
+    doc = nlp(sentence)
     ents = []
     for ee in entities:
         ents.append(doc.char_span(ee[1], ee[2], ee[0]))
     doc.ents = ents
+    return doc
+
 
+def apply_ner(sentence: str):
+    classifications = ner_pipeline(sentence)
+    entities = postprocess(classifications)
+    doc = set_entities(sentence, entities)
     displacy_html = displacy.render(doc, style="ent", options=options)
     return displacy_html
 
@@ -57,9 +79,14 @@ intf = gr.Interface(fn=apply_ner, inputs="text", outputs="html", title='Slovak N
                     examples=[["Laboratóriá Úradu verejného zdravotníctva sekvenovaním potvrdili výskyt ďalších "
                                "štyroch prípadov variantu omikron na Slovensku."],
                               ["Čaputová opakovane tvrdí, že \"spravodlivosť na Slovensku neplatí vždy pre všetkých "
-                               "rovnako\"."]],
-                    description="Named Entity Recognition (NER) labels persons, organizations and locations in the "
-                                "given sentence. You can try out one of the examples below or type your own sentence. "
-                                "Don't forget to use double quotes (\") instead of curved quotes („“)",
+                               "rovnako\"."],
+                              ["Minister financií a líder mandátovo najsilnejšieho hnutia OĽaNO Igor Matovič "
+                               "upozorňuje, že následky tretej vlny budú na Slovensku veľmi veľké."],
+                              ["Začiatkom roka sa objavili nezhody medzi Richardom Sulíkom a šéfom hnutia OĽANO "
+                               "Igorom Matovičom, ktoré v istej miere pretrvávajú aj dodnes."]],
+                    description="Named-entity recognition (NER) labels named-entities in unstructured text. This "
+                                "implementation supports three labels: person (OSOBA), organization (ORGANIZÁCIA) and "
+                                "location (LOKALITA). You can try out one of the examples below or type your own "
+                                "sentence. Don't forget to use double quotes (\" \") instead of curved quotes („ “)",
                     article="")
 intf.launch()