Minor fix again

huggingface · Oct 3, 2022 · 45b61c1 · 45b61c1 · github-actions · Oct 3, 2022
1 parent bf783f2
commit 45b61c1
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/src/datasets/packaged_modules/folder_based_builder/folder_based_builder.py b/src/datasets/packaged_modules/folder_based_builder/folder_based_builder.py
@@ -2,14 +2,15 @@
 import itertools
 import os
 from dataclasses import dataclass
-from typing import Any, List, Optional, Tuple
+from typing import List, Optional, Tuple
 
 import pandas as pd
 import pyarrow as pa
 import pyarrow.compute as pc
 import pyarrow.json as paj
 
 import datasets
+from datasets.features.features import FeatureType
 from datasets.tasks.base import TaskTemplate
 
 
@@ -66,7 +67,7 @@ class FolderBasedBuilder(datasets.GeneratorBasedBuilder):
         CLASSIFICATION_TASK: classification task to use if labels are obtained from the folder structure
     """
 
-    BASE_FEATURE: Any
+    BASE_FEATURE: FeatureType
     BASE_COLUMN_NAME: str
     BUILDER_CONFIG_CLASS: FolderBasedBuilderConfig
     EXTENSIONS: List[str]

diff --git a/tests/packaged_modules/test_folder_based_builder.py b/tests/packaged_modules/test_folder_based_builder.py
@@ -11,14 +11,15 @@
     FolderBasedBuilder,
     FolderBasedBuilderConfig,
 )
+from datasets.tasks import TextClassification
 
 
 class DummyFolderBasedBuilder(FolderBasedBuilder):
     BASE_FEATURE = None
     BASE_COLUMN_NAME = "base"
     BUILDER_CONFIG_CLASS = FolderBasedBuilderConfig
     EXTENSIONS = [".txt"]
-    CLASSIFICATION_TASK = None
+    CLASSIFICATION_TASK = TextClassification(text_column="base", label_column="label")
 
 
 @pytest.fixture