fix entity detection (#127)

deeppavlov · Mar 21, 2022 · 531a34b · 531a34b
1 parent 482a238
commit 531a34b
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/annotators/entity_detection_rus/entity_detection_parser.py b/annotators/entity_detection_rus/entity_detection_parser.py
@@ -54,6 +54,7 @@ def __init__(
         return_entities_with_tags: bool = False,
         add_nouns: bool = False,
         thres_proba: float = 0.8,
+        misc_proba: float = 0.8,
         **kwargs
     ):
         """
@@ -73,6 +74,7 @@ def __init__(
         self.ignore_points = ignore_points
         self.return_entities_with_tags = return_entities_with_tags
         self.thres_proba = thres_proba
+        self.misc_proba = misc_proba
         self.tag_ind_dict = {}
         with open(str(expand_path(tags_file))) as fl:
             tags = [line.split("\t")[0] for line in fl.readlines()]
@@ -206,7 +208,10 @@ def entities_from_tags(self, tokens, tags, tag_probas):
 
                 entity_dict[f_tag].append(tok)
                 entity_positions_dict[f_tag].append(cnt)
-                entity_probas_dict[f_tag].append(probas[self.tags_ind[tag]])
+                if self.tags_ind[tag] < len(probas):
+                    entity_probas_dict[f_tag].append(probas[self.tags_ind[tag]])
+                else:
+                    entity_probas_dict[f_tag].append(self.misc_proba)
 
             elif any(entity_dict.values()):
                 for tag, entity in entity_dict.items():