File size: 6,189 Bytes
a01ef8c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
#!/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)

    @property
    def dataset(self) -> Dataset:
        """
        Returns datasets.arrow_dataset.Dataset object
        """
        return self._dataset

    @property
    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

    @property
    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)