Repository: ArneBinder/pytorch-ie (language: python)

Issue:
implement windowing for text classification based rel
As discussed. Tested on example prediction pipeline, train run ~~is still running, but looks good.~~ EDIT: finished successful.

~~This also adds a parameter `add_negative_examples` that allows to add all relation candidates that are not in the existing relations as relations with non_label, i.e. `no_relation`.~~ EDIT: This was removed again as it does not proved as useful as expected.

This outsources several parts of the `encode_inputs` method to test them separately.

relevant changes:
- implement token based windowing. This is activated if `max_window` is provided. The `max_window` defines the maximum number of token ids in the encoded input (takes added marker and special tokens into account), i.e. it is ready to be set to the max input size of the respective transformer. 
- annotation getters for `Document` take a default argument. By calling them with `default=None` it is possible to know if relations are available and should be taken into account to filter entity pair candidates or not (like for inference). This is relevant when not all documents in the training set have relations. This was previously causing to add all possible entity pair combinations for such documents as training instances.
- outsource `_get_window_around_slice` and add tests
- outsource `_encode_text` to easy testing
- outsource `_enumerate_entity_pairs` and add tests
- add tests for partitioning (since we are already on it)
- add tests for windowing

NOTE: This revealed a bug in `decode()`, see #62. The tests for `decode()` are not complete (but pipeline tests run as expected) and should be fixed in the respective PR.


Discussion / hints:
Overall, this looks good. However, it also seems like the kind of logic that requires some unit tests.

Also, I have the feeling that the logic should not be that tightly integrated into the TaskModule. Windowing seems to be an operation that could be applied to many tasks. The question is if we can find a way to add windowing without touching the general logic of a task module.
This required a bit more effort than expected to add tests for all relevant functionality. It should be finally ready now. @ChristophAlt could you have a look?

Fix patch:
diff --git a/pytorch_ie/data/datasets/brat.py b/pytorch_ie/data/datasets/brat.py
index 278503be..ee68f820 100644
--- a/pytorch_ie/data/datasets/brat.py
+++ b/pytorch_ie/data/datasets/brat.py
@@ -209,17 +209,20 @@ def convert_document_to_brat(
     doc: Document,
     head_argument_name: str = DEFAULT_HEAD_ARGUMENT_NAME,
     tail_argument_name: str = DEFAULT_TAIL_ARGUMENT_NAME,
-    span_annotation_names: List[str] = [DEFAULT_SPAN_ANNOTATION_NAME],
-    relation_annotation_names: List[str] = [DEFAULT_RELATION_ANNOTATION_NAME],
+    span_annotation_names: Optional[List[str]] = None,
+    relation_annotation_names: Optional[List[str]] = None,
     **kwargs,
 ) -> Tuple[Optional[str], str, List[str]]:
     serialized_annotations: Dict[Annotation, str] = {}
-
+    if span_annotation_names is None:
+        span_annotation_names = [DEFAULT_SPAN_ANNOTATION_NAME]
+    if relation_annotation_names is None:
+        relation_annotation_names = [DEFAULT_RELATION_ANNOTATION_NAME]
     for entity_annotation_name in span_annotation_names:
-        for span_ann in doc.span_annotations(entity_annotation_name):
+        for span_ann in doc.span_annotations(entity_annotation_name) or []:
             serialized_annotations[span_ann] = serialize_labeled_span(span_ann, doc, **kwargs)
     for relation_annotation_name in relation_annotation_names:
-        for rel_ann in doc.relation_annotations(relation_annotation_name):
+        for rel_ann in doc.relation_annotations(relation_annotation_name) or []:
             serialized_annotations[rel_ann] = serialize_binary_relation(
                 rel_ann,
                 doc,
diff --git a/pytorch_ie/data/document.py b/pytorch_ie/data/document.py
index 94a09e40..33b8527e 100644
--- a/pytorch_ie/data/document.py
+++ b/pytorch_ie/data/document.py
@@ -182,20 +182,20 @@ def add_prediction(self, name: str, prediction: Annotation) -> None:
         self._predictions[name].append(prediction)
 
     # TODO: rework all getters for annotations (e.g. remove type: ignore)
-    def annotations(self, name: str) -> AnnotationLayer:
-        return self._annotations.get(name, [])
+    def annotations(self, name: str) -> Optional[AnnotationLayer]:
+        return self._annotations.get(name, None)
 
-    def span_annotations(self, name: str) -> List[LabeledSpan]:
+    def span_annotations(self, name: str) -> Optional[List[LabeledSpan]]:
         return self.annotations(name=name)  # type: ignore
 
-    def relation_annotations(self, name: str) -> List[BinaryRelation]:
+    def relation_annotations(self, name: str) -> Optional[List[BinaryRelation]]:
         return self.annotations(name=name)  # type: ignore
 
-    def label_annotations(self, name: str) -> List[Label]:
+    def label_annotations(self, name: str) -> Optional[List[Label]]:
         return self.annotations(name=name)  # type: ignore
 
-    def predictions(self, name: str) -> AnnotationLayer:
-        return self._predictions.get(name, [])
+    def predictions(self, name: str) -> Optional[AnnotationLayer]:
+        return self._predictions.get(name, None)
 
     def clear_predictions(self, name: str) -> None:
         if name in self._predictions:
diff --git a/pytorch_ie/taskmodules/transformer_re_text_classification.py b/pytorch_ie/taskmodules/transformer_re_text_classification.py
index b1f244e8..208b3cd2 100644
--- a/pytorch_ie/taskmodules/transformer_re_text_classification.py
+++ b/pytorch_ie/taskmodules/transformer_re_text_classification.py
@@ -282,6 +282,12 @@ def prepare(self, documents: List[Document]) -> None:
         for document in documents:
             entities = document.span_annotations(self.entity_annotation)
             relations = document.relation_annotations(self.relation_annotation)
+            assert (
+                entities is not None
+            ), f"document has no span annotations with name '{self.entity_annotation}'"
+            assert (
+                relations is not None
+            ), f"document has no relation annotations with name '{self.relation_annotation}'"
 
             if self.add_type_to_marker:
                 for entity in entities:
@@ -350,12 +356,19 @@ def encode_input(
 
         for document in documents:
             entities = document.span_annotations(self.entity_annotation)
+            assert (
+                entities is not None
+            ), f"document has no span annotations of name {self.entity_annotation}"
             relations = document.relation_annotations(self.relation_annotation)
             relation_mapping = {(rel.head, rel.tail): rel.label for rel in relations or []}
 
             partitions: Sequence[Optional[LabeledSpan]]
             if self.partition_annotation is not None:
-                partitions = document.span_annotations(self.partition_annotation)
+                partitions_or_none = document.span_annotations(self.partition_annotation)
+                assert (
+                    partitions_or_none is not None
+                ), f"document has no span annotations of name {self.partition_annotation}"
+                partitions = partitions_or_none
             else:
                 # use single dummy partition
                 partitions = [None]
@@ -493,6 +506,9 @@ def encode_target(
             meta = metadata[i]
 
             relations = document.relation_annotations(self.relation_annotation)
+            assert (
+                relations is not None
+            ), f"document has no relation annotations of name '{self.relation_annotation}'"
 
             head_tail_to_labels = {
                 (relation.head, relation.tail): relation.labels for relation in relations
diff --git a/pytorch_ie/taskmodules/transformer_seq2seq.py b/pytorch_ie/taskmodules/transformer_seq2seq.py
index 95c71b5c..989e5313 100644
--- a/pytorch_ie/taskmodules/transformer_seq2seq.py
+++ b/pytorch_ie/taskmodules/transformer_seq2seq.py
@@ -99,6 +99,9 @@ def encode_input(
 
     def document_to_target_string(self, document: Document) -> str:
         relations = document.relation_annotations(self.relation_annotation)
+        assert (
+            relations is not None
+        ), f"document has no relation annotations of name '{self.relation_annotation}'"
 
         head_to_relation: Dict[LabeledSpan, List[BinaryRelation]] = {}
         for relation in relations:
diff --git a/pytorch_ie/taskmodules/transformer_span_classification.py b/pytorch_ie/taskmodules/transformer_span_classification.py
index 9ea0c596..a6c08a5e 100644
--- a/pytorch_ie/taskmodules/transformer_span_classification.py
+++ b/pytorch_ie/taskmodules/transformer_span_classification.py
@@ -92,6 +92,9 @@ def prepare(self, documents: List[Document]) -> None:
         labels = set()
         for document in documents:
             entities = document.span_annotations(self.entity_annotation)
+            assert (
+                entities is not None
+            ), f"document has no span annotations with name '{self.entity_annotation}'"
 
             for entity in entities:
                 # TODO: labels is a set, use update
@@ -114,29 +117,20 @@ def encode_input(
         List[Metadata],
         Optional[List[Document]],
     ]:
-        expanded_documents = None
-        if self.single_sentence:
-            sent: LabeledSpan
-            input_ = [
-                self.tokenizer(
-                    doc.text[sent.start : sent.end],
-                    padding=False,
-                    truncation=self.truncation,
-                    max_length=self.max_length,
-                    is_split_into_words=False,
-                    return_offsets_mapping=True,
-                    return_special_tokens_mask=True,
-                )
-                for doc in documents
-                for sent in doc.span_annotations(self.sentence_annotation)
-            ]
-            expanded_documents = [
-                doc for doc in documents for _ in doc.span_annotations(self.sentence_annotation)
-            ]
-        else:
-            input_ = [
-                self.tokenizer(
-                    doc.text,
+        input_ = []
+        expanded_documents = []
+        for doc in documents:
+            if self.single_sentence:
+                partitions_or_none = doc.span_annotations(self.sentence_annotation)
+                assert (
+                    partitions_or_none is not None
+                ), f"document has no span annotations with name '{self.sentence_annotation}'"
+                partitions = partitions_or_none
+            else:
+                partitions = [LabeledSpan(start=0, end=len(doc.text), label="FULL_DOCUMENT")]
+            for partition in partitions:
+                encoding = self.tokenizer(
+                    doc.text[partition.start : partition.end],
                     padding=False,
                     truncation=self.truncation,
                     max_length=self.max_length,
@@ -144,8 +138,8 @@ def encode_input(
                     return_offsets_mapping=True,
                     return_special_tokens_mask=True,
                 )
-                for doc in documents
-            ]
+                input_.append(encoding)
+                expanded_documents.append(doc)
 
         metadata = [
             {
@@ -159,7 +153,7 @@ def encode_input(
             i = 0
             for document in documents:
                 for sentence_index in range(
-                    len(document.span_annotations(self.sentence_annotation))
+                    len(document.span_annotations(self.sentence_annotation) or [])
                 ):
                     metadata[i]["sentence_index"] = sentence_index
                     i += 1
@@ -176,8 +170,15 @@ def encode_target(
         if self.single_sentence:
             for i, document in enumerate(documents):
                 entities = document.span_annotations(self.entity_annotation)
+                assert (
+                    entities is not None
+                ), f"document has no span annotations with name '{self.entity_annotation}'"
                 sentence_idx = metadata[i]["sentence_index"]
-                sentence = document.span_annotations(self.sentence_annotation)[sentence_idx]
+                partitions = document.span_annotations(self.sentence_annotation)
+                assert (
+                    partitions is not None
+                ), f"document has no span annotations with name '{self.sentence_annotation}'"
+                sentence = partitions[sentence_idx]
 
                 label_ids = []
                 entity: LabeledSpan
@@ -203,6 +204,9 @@ def encode_target(
         else:
             for i, document in enumerate(documents):
                 entities = document.span_annotations(self.entity_annotation)
+                assert (
+                    entities is not None
+                ), f"document has no span annotations with name '{self.entity_annotation}'"
 
                 label_ids = []
                 for entity in entities:
@@ -245,10 +249,11 @@ def create_annotations_from_output(
         if self.single_sentence:
             document = encoding.document
             metadata = encoding.metadata
-
-            sentence = document.span_annotations(self.sentence_annotation)[
-                metadata["sentence_index"]
-            ]
+            partitions = document.span_annotations(self.sentence_annotation)
+            assert (
+                partitions is not None
+            ), f"document has no span annotations with name '{self.sentence_annotation}'"
+            sentence = partitions[metadata["sentence_index"]]
 
             # tag_sequence = [
             #     "O" if stm else tag
diff --git a/pytorch_ie/taskmodules/transformer_text_classification.py b/pytorch_ie/taskmodules/transformer_text_classification.py
index 37b42f8b..ce0c1321 100644
--- a/pytorch_ie/taskmodules/transformer_text_classification.py
+++ b/pytorch_ie/taskmodules/transformer_text_classification.py
@@ -117,6 +117,9 @@ def prepare(self, documents: List[Document]) -> None:
         labels = set()
         for document in documents:
             annotations = document.label_annotations(self.annotation)
+            assert (
+                annotations is not None
+            ), f"document as no label annotations with name '{self.annotation}'"
 
             for annotation in annotations:
                 # TODO: labels is a set...
@@ -171,17 +174,20 @@ def encode_target(
 
         target: List[TransformerTextClassificationTargetEncoding] = []
         for i, document in enumerate(documents):
+            annotations = document.label_annotations(self.annotation)
+            assert (
+                annotations is not None
+            ), f"document as no label annotations with name '{self.annotation}'"
             if self.multi_label:
                 label_ids = [0] * len(self.label_to_id)
-                for annotation in document.label_annotations(self.annotation):
+                for annotation in annotations:
                     for label in annotation.labels:
                         label_id = self.label_to_id[label]
                         label_ids[label_id] = 1
             else:
-                annotation_labels = document.label_annotations(self.annotation)
-                assert len(annotation_labels) == 1 and not annotation_labels[0].is_multilabel
+                assert len(annotations) == 1 and not annotations[0].is_multilabel
 
-                label = annotation_labels[0].label_single
+                label = annotations[0].label_single
                 label_ids = [self.label_to_id[label]]
 
             target.append(label_ids)
diff --git a/pytorch_ie/taskmodules/transformer_token_classification.py b/pytorch_ie/taskmodules/transformer_token_classification.py
index 6982f78d..55ab27f7 100644
--- a/pytorch_ie/taskmodules/transformer_token_classification.py
+++ b/pytorch_ie/taskmodules/transformer_token_classification.py
@@ -90,6 +90,9 @@ def prepare(self, documents: List[Document]) -> None:
         labels = set()
         for document in documents:
             entities = document.span_annotations(self.entity_annotation)
+            assert (
+                entities is not None
+            ), f"document has no span annotations with name '{self.entity_annotation}'"
 
             for entity in entities:
                 # TODO: labels is a set...
@@ -113,29 +116,21 @@ def encode_input(
         List[Metadata],
         Optional[List[Document]],
     ]:
-        expanded_documents = None
-        if self.single_sentence:
-            sent: LabeledSpan
-            input_ = [
-                self.tokenizer(
-                    doc.text[sent.start : sent.end],
-                    padding=False,
-                    truncation=False,
-                    max_length=None,
-                    is_split_into_words=False,
-                    return_offsets_mapping=True,
-                    return_special_tokens_mask=True,
-                )
-                for doc in documents
-                for sent in doc.span_annotations(self.sentence_annotation)
-            ]
-            expanded_documents = [
-                doc for doc in documents for _ in doc.span_annotations(self.sentence_annotation)
-            ]
-        else:
-            input_ = [
-                self.tokenizer(
-                    doc.text,
+        expanded_documents = []
+        input_ = []
+        for doc in documents:
+            if self.single_sentence:
+                partitions_or_none = doc.span_annotations(self.sentence_annotation)
+                assert (
+                    partitions_or_none
+                ), f"document has no span annotations with name '{self.sentence_annotation}'"
+                partitions = partitions_or_none
+            else:
+                partitions = [LabeledSpan(start=0, end=len(doc.text), label="FULL_DOCUMENT")]
+
+            for partition in partitions:
+                encoding = self.tokenizer(
+                    doc.text[partition.start : partition.end],
                     padding=False,
                     truncation=False,
                     max_length=None,
@@ -143,8 +138,8 @@ def encode_input(
                     return_offsets_mapping=True,
                     return_special_tokens_mask=True,
                 )
-                for doc in documents
-            ]
+                input_.append(encoding)
+                expanded_documents.append(doc)
 
         metadata = [
             {
@@ -158,7 +153,7 @@ def encode_input(
             i = 0
             for document in documents:
                 for sentence_index in range(
-                    len(document.span_annotations(self.sentence_annotation))
+                    len(document.span_annotations(self.sentence_annotation) or [])
                 ):
                     metadata[i]["sentence_index"] = sentence_index
                     i += 1
@@ -175,8 +170,15 @@ def encode_target(
         if self.single_sentence:
             for i, document in enumerate(documents):
                 entities = document.span_annotations(self.entity_annotation)
+                assert (
+                    entities
+                ), f"document has no span annotations with name '{self.entity_annotation}'"
                 sentence_idx = metadata[i]["sentence_index"]
-                sentence = document.span_annotations(self.sentence_annotation)[sentence_idx]
+                partitions = document.span_annotations(self.sentence_annotation)
+                assert (
+                    partitions
+                ), f"document has no span annotations with name '{self.sentence_annotation}'"
+                sentence = partitions[sentence_idx]
 
                 word_ids = input_encodings[i].word_ids()
                 label_ids = [
@@ -215,6 +217,9 @@ def encode_target(
                 ]
 
                 entities = document.span_annotations(self.entity_annotation)
+                assert (
+                    entities
+                ), f"document has no span annotations with name '{self.entity_annotation}'"
 
                 for entity in entities:
                     start_idx = input_encodings[i].char_to_token(entity.start)
@@ -244,10 +249,11 @@ def create_annotations_from_output(
         if self.single_sentence:
             document = encoding.document
             metadata = encoding.metadata
-
-            sentence = document.span_annotations(self.sentence_annotation)[
-                metadata["sentence_index"]
-            ]
+            partitions = document.span_annotations(self.sentence_annotation)
+            assert (
+                partitions
+            ), f"document has no span annotations with name '{self.sentence_annotation}'"
+            sentence = partitions[metadata["sentence_index"]]
 
             tag_sequence = [
                 "O" if stm else tag


Test patch:
diff --git a/tests/data/test_document.py b/tests/data/test_document.py
index 0c18fe52..110dd25e 100644
--- a/tests/data/test_document.py
+++ b/tests/data/test_document.py
@@ -135,4 +135,4 @@ def test_document():
 
     document.clear_predictions("prediction")
     prediction2 = document.predictions("prediction")
-    assert len(prediction2) == 0
+    assert prediction2 is None