From 5cfc0091a521c30928757f9d1a1d089700f09aa4 Mon Sep 17 00:00:00 2001 From: Arne Binder Date: Tue, 7 Nov 2023 20:23:52 +0100 Subject: [PATCH] use GeneratorBasedBuilder from pie_datasets --- dataset_builders/pie/brat/brat.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/dataset_builders/pie/brat/brat.py b/dataset_builders/pie/brat/brat.py index 56b55db6..752de772 100644 --- a/dataset_builders/pie/brat/brat.py +++ b/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,