diff --git a/flair/datasets/sequence_labeling.py b/flair/datasets/sequence_labeling.py index 0c0900dc2f..5779a66c5c 100644 --- a/flair/datasets/sequence_labeling.py +++ b/flair/datasets/sequence_labeling.py @@ -19,6 +19,7 @@ Sentence, Span, Token, + _iter_dataset, get_spans_from_bio, ) from flair.datasets.base import find_train_dev_test_files @@ -476,7 +477,7 @@ def write_dataset_to_file( (2) every sentence is separated from the previous one by an empty line """ with open(file_path, mode="w") as output_file: - for sentence in dataset: + for sentence in _iter_dataset(dataset): texts = [token.text for token in sentence.tokens] texts_and_labels = [texts] for label_type, level in label_type_tuples: @@ -515,7 +516,7 @@ def get_level_of_label(self, label_type: str) -> Union[Type[Token], Type[Span]]: """ for dataset in [self.train, self.dev, self.test]: if dataset: - for sentence in dataset: + for sentence in _iter_dataset(dataset): for label in sentence.get_labels(label_type): if isinstance(label.data_point, Token): return Token