Spaces:
Configuration error
Configuration error
| #!/usr/bin/env python | |
| # -*- coding: utf-8 -*- | |
| # | |
| # Copyright (c) 2022 Intel Corporation | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # | |
| # SPDX-License-Identifier: Apache-2.0 | |
| # | |
| import os | |
| from datasets import concatenate_datasets | |
| from datasets.arrow_dataset import Dataset | |
| from tlt import TLT_BASE_DIR | |
| from tlt.utils.file_utils import read_json_file | |
| from tlt.datasets.hf_dataset import HFDataset | |
| from tlt.datasets.text_classification.text_classification_dataset import TextClassificationDataset | |
| from downloader.datasets import DataDownloader | |
| DATASET_CONFIG_DIR = os.path.join(TLT_BASE_DIR, "datasets/configs") | |
| class HFTextClassificationDataset(TextClassificationDataset, HFDataset): | |
| """ | |
| A text classification dataset from the Hugging Face datasets catalog | |
| """ | |
| def __init__(self, dataset_dir, dataset_name, split=['train'], num_workers=0, shuffle_files=True, | |
| distributed=False): | |
| if not isinstance(split, list): | |
| raise ValueError("Value of split argument must be a list.") | |
| TextClassificationDataset.__init__(self, dataset_dir, dataset_name, "huggingface") | |
| self._preprocessed = {} | |
| self._split = split | |
| self._data_loader = None | |
| self._train_loader = None | |
| self._test_loader = None | |
| self._validation_loader = None | |
| self._train_subset = None | |
| self._test_subset = None | |
| self._validation_subset = None | |
| self._num_workers = num_workers | |
| self._shuffle = shuffle_files | |
| self._distributed = distributed | |
| self._info = { | |
| 'name': dataset_name | |
| } | |
| if len(split) == 1: | |
| self._validation_type = None # Train & evaluate on the whole dataset | |
| # If only one split is given use it as the main dataset object | |
| self._dataset = self.load_hf_dataset(dataset_name, split=split[0]) | |
| else: | |
| self._validation_type = 'defined_split' # Defined by user or Hugging Face | |
| if 'train' in split: | |
| self._dataset = self.load_hf_dataset(dataset_name, split='train') | |
| self._train_indices = range(len(self._dataset)) | |
| self._train_subset = self.train_subset | |
| if 'test' in split: | |
| test_dataset = self.load_hf_dataset(dataset_name, split='test') | |
| test_length = len(test_dataset) | |
| if self._dataset: | |
| current_length = len(self._dataset) | |
| self._dataset = concatenate_datasets([self._dataset, test_dataset]) | |
| self._test_indices = range(current_length, current_length + test_length) | |
| else: | |
| self._dataset = test_dataset | |
| self._test_indices = range(test_length) | |
| self._test_subset = self.test_subset | |
| if 'validation' in split: | |
| validation_dataset = self.load_hf_dataset(dataset_name, split='validation') | |
| validation_length = len(validation_dataset) | |
| if self._dataset: | |
| current_length = len(self._dataset) | |
| self._dataset = concatenate_datasets([self._dataset, validation_dataset]) | |
| self._validation_indices = range(current_length, current_length + validation_length) | |
| else: | |
| self._dataset = validation_dataset | |
| self._validation_indices = range(validation_length) | |
| self._validation_subset = self.validation_subset | |
| if 'unsupervised' in split: | |
| unsupervised_dataset = self.load_hf_dataset(dataset_name, split='unsupervised') | |
| if self._dataset: | |
| self._dataset = concatenate_datasets([self._dataset, unsupervised_dataset]) | |
| else: | |
| self._dataset = unsupervised_dataset | |
| def load_hf_dataset(self, dataset_name: str, split: str) -> Dataset: | |
| """ | |
| Helper function to load the dataset from hugging face catalog | |
| """ | |
| main_dataset = dataset_name | |
| subset = None | |
| config_file = os.path.join(DATASET_CONFIG_DIR, "hf_text_classification_datasets.json") | |
| config_dict = read_json_file(config_file) | |
| available_datasets = list(config_dict.keys()) | |
| if dataset_name not in available_datasets: | |
| raise ValueError("Dataset is not supported. Choose from: {}".format(available_datasets)) | |
| # We separate the dataset_name by checking whether it has the format of "dataset/subset" | |
| if '/' in dataset_name: | |
| main_dataset = dataset_name.split('/')[0] | |
| subset = dataset_name.split('/')[1] | |
| if subset is not None: | |
| downloader = DataDownloader(main_dataset, self._dataset_dir, catalog='hugging_face', subset=subset) | |
| else: | |
| downloader = DataDownloader(main_dataset, self._dataset_dir, catalog='hugging_face') | |
| return downloader.download(split=split) | |
| def dataset(self) -> Dataset: | |
| """ | |
| Returns datasets.arrow_dataset.Dataset object | |
| """ | |
| return self._dataset | |
| def class_names(self) -> list: | |
| """ | |
| Returns a list of class labels | |
| """ | |
| try: | |
| names = self.dataset.features['label'].names | |
| except KeyError: | |
| names = self.dataset.features['labels'].names | |
| return names | |
| def info(self): | |
| """ | |
| Returns a dictionary of information about the dataset | |
| """ | |
| return {'dataset_info': self._info, 'preprocessing_info': self._preprocessed} | |
| def __len__(self): | |
| return len(self._dataset) | |