use GeneratorBasedBuilder from pie_datasets

dataset_builders/pie/brat/brat.py

@@ -4,16 +4,12 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import datasets
-import pytorch_ie
-from pytorch_ie.annotations import (
-    BinaryRelation,
-    LabeledMultiSpan,
-    LabeledSpan,
-    _post_init_single_label,
-)
+from pytorch_ie.annotations import BinaryRelation, LabeledMultiSpan, LabeledSpan
 from pytorch_ie.core import Annotation, AnnotationList, annotation_field
 from pytorch_ie.documents import TextBasedDocument
 
+from pie_datasets import GeneratorBasedBuilder
+
 logger = logging.getLogger(__name__)
 
 
@@ -288,7 +284,7 @@ def __init__(self, merge_fragmented_spans: bool = False, **kwargs):
         self.merge_fragmented_spans = merge_fragmented_spans
 
 
-class BratDatasetLoader(pytorch_ie.data.builder.GeneratorBasedBuilder):
+class BratDatasetLoader(GeneratorBasedBuilder):
     # this requires https://github.com/ChristophAlt/pytorch-ie/pull/288
     DOCUMENT_TYPES = {
         "default": BratDocument,