Update Inference code

README.md

@@ -68,20 +68,18 @@ ner_results = nlp(example)
 
 modified_results = []
 for ent in ner_results:
-  ent["entity_group"] = int(ent["entity_group"].lower().replace("label_",""))
-  ent["entity_group"] = labels_list[ent["entity_group"]]
-
-  if ent["entity_group"] != "O":
-    if "B-" in ent["entity_group"]:
-      ent["entity_group"] = ent["entity_group"].replace("B-","")
-      modified_results.append(ent)
-    elif "I-" in ent["entity_group"]:
-      ## check related entity-group
-      label = ent["entity_group"].replace("I-","")
-      if len(modified_results) > 0 and label == modified_results[-1]["entity_group"]:
-        modified_results[-1]["word"] += f" {ent['word']}"
-        modified_results[-1]["score"] = sum([modified_results[-1]["score"], ent["score"]])/2
-        modified_results[-1]["end"] = ent["end"]
+  if ent["entity_group"].lower().replace("label_","").isnumeric():
+      ent["entity_group"] = int(ent["entity_group"].lower().replace("label_",""))
+      ent["entity_group"] = labels_list[ent["entity_group"]]
+
+  if len(modified_results) > 0 and ent["start"] == modified_results[-1]["end"]:
+    modified_results[-1]["word"] += f"{ent['word']}".replace("▁"," ").strip()
+    modified_results[-1]["word"] = modified_results[-1]["word"].replace("▁"," ").strip()
+    modified_results[-1]["score"] = sum([modified_results[-1]["score"], ent["score"]])/2
+    modified_results[-1]["end"] = ent["end"]
+  else:
+    modified_results.append(ent)
+
 
 for res in modified_results:
   print(res["word"], "==>", res["entity_group"])