fix output

app.py

@@ -6,20 +6,19 @@ tokenizer, model, gazetteers_for_matching = load()
 print("Loaded model")
 
 examples = [
-    "Masarykova univerzita se nachází v Brně.",
-    "Barack Obama navštívil Prahu minulý týden.",
-    "Angela Merkelová se setkala s francouzským prezidentem v Paříži.",
-    "Karel Čapek napsal knihu R.U.R., která byla poprvé představena v Praze.",
-    "Nobelova cena za fyziku byla udělena týmu vědců z MIT."
+    "Masarykova univerzita se nachází v Brně .",
+    "Barack Obama navštívil Prahu minulý týden .",
+    "Angela Merkelová se setkala s francouzským prezidentem v Paříži .",
+    "Nobelova cena za fyziku byla udělena týmu vědců z MIT ."
 ]
 
 def ner(text):
     result = run(tokenizer, model, gazetteers_for_matching, text)
-    return result
+    return {"text": text, "entities": result}
 
 demo = gr.Interface(ner,
              gr.Textbox(placeholder="Enter sentence here..."), 
-             gr.HighlightedText(show_legend=True,),
+             gr.HighlightedText(),
              examples=examples)
 
 if __name__ == "__main__":

website_script.py

@@ -24,7 +24,7 @@ def load():
 def run(tokenizer, model, gazetteers_for_matching, text):
 
     tokenized_inputs = tokenizer(
-        text, truncation=True, is_split_into_words=False
+        text, truncation=True, is_split_into_words=False, return_offsets_mapping=True
     )
     matches = gazetteer_matching(text, gazetteers_for_matching)
     new_g = []
@@ -48,12 +48,36 @@ def run(tokenizer, model, gazetteers_for_matching, text):
     softmax = torch.nn.Softmax(dim=2)
     scores = softmax(output).squeeze(0).tolist()
     result = []
+    temp = {
+            "start": 0,
+            "end": 0,
+            "entity": "O",
+            "score": 0,
+            "word": "",
+            "count": 0
+        }
     for pos, entity, score in zip(tokenized_inputs.offset_mapping, predicted_tags[0], scores):
-        result.append({
+        if pos[0] == pos[1] or entity == "O":
+            continue
+        if temp["entity"] == entity[2:]:  # same entity
+            space = " " if pos[0] - temp["end"] >= 1 else ""
+            temp["end"] = pos[1]
+            temp["word"] += space + text[pos[0]:pos[1]]
+            temp["count"] += 1
+            temp["score"] += max(score)
+        else:  # new entity
+            if temp["count"] > 0:
+                temp["score"] /= temp.pop("count")
+                result.append(temp)
+            temp = {
             "start": pos[0],
             "end": pos[1],
-            "entity": entity,
-            "score": max(score),
+            "entity": entity[2:],
+            "score": 0,
             "word": text[pos[0]:pos[1]],
-        })
+            "count": 1
+            }
+    if temp["count"] > 0:
+        temp["score"] /= temp.pop("count")
+        result.append(temp)
     return result