fix tars typing

flairNLP · Aug 30, 2024 · 642337b · 642337b
1 parent a66a856
commit 642337b
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 12 deletions.
diff --git a/flair/data.py b/flair/data.py
@@ -1024,8 +1024,7 @@ def __getitem__(self, idx: int) -> Token: ...
     def __getitem__(self, s: slice) -> Span: ...
 
     @typing.overload
-    def __getitem__(self, s: typing.Tuple[Span, Span]) -> Relation:
-        ...
+    def __getitem__(self, s: typing.Tuple[Span, Span]) -> Relation: ...
 
     def __getitem__(self, subscript):
         if isinstance(subscript, tuple):

diff --git a/flair/models/relation_classifier_model.py b/flair/models/relation_classifier_model.py
@@ -392,7 +392,9 @@ def _entity_pair_permutations(
 
             # Obtain gold label, if existing
             gold_relation = sentence[head.span, tail.span]
-            gold_label: Optional[str] = gold_relation.get_label(self.label_type, zero_tag_value=None).value
+            gold_label: Optional[str] = gold_relation.get_label(self.label_type, zero_tag_value="O").value
+            if gold_label == "O":
+                gold_label = None
             yield head, tail, gold_label
 
     def _encode_sentence(

diff --git a/flair/models/tars_model.py b/flair/models/tars_model.py
@@ -571,19 +571,16 @@ def predict(
 
                         already_set_indices: List[int] = []
 
-                        sorted_x = sorted(all_detected.items(), key=operator.itemgetter(1))
-                        sorted_x.reverse()
-                        for tuple in sorted_x:
-                            # get the span and its label
-                            label = tuple[0]
-
+                        sorted_x = sorted(all_detected.items(), key=operator.itemgetter(1), reverse=True)
+                        for label, _ in sorted_x:
+                            span = typing.cast(Span, label.data_point)
                             label_length = (
                                 0 if not self.prefix else len(label.value.split(" ")) + len(self.separator.split(" "))
                             )
 
                             # determine whether tokens in this span already have a label
                             tag_this = True
-                            for token in label.data_point:
+                            for token in span:
                                 corresponding_token = sentence.get_token(token.idx - label_length)
                                 if corresponding_token is None:
                                     tag_this = False
@@ -595,8 +592,8 @@ def predict(
                             # only add if all tokens have no label
                             if tag_this:
                                 # make and add a corresponding predicted span
-                                start_pos = label.data_point.data_point[0].idx - label_length - 1
-                                end_pos = label.data_point.data_point[-1].idx - label_length
+                                start_pos = span.tokens[0].idx - label_length - 1
+                                end_pos = span.tokens[-1].idx - label_length
 
                                 predicted_span = sentence[start_pos:end_pos]
                                 predicted_span.add_label(label_name, value=label.value, score=label.score)