from pt_variety_identifier.src.data import Data as DataBase


class Data(DataBase):
    def __init__(self, dataset_name, test_set_list) -> None:
        self._DOMAINS = ['journalistic', 'literature',
                         'legal', 'politics', 'web', 'social_media']

        self.dataset_name = dataset_name
        self.test_set_list = test_set_list

    def load_test_set(self, filter_label_2=False):
        dataset_return = super().load_test_set(filter_label_2)

        validation_dataset_dict = self.load_validation_set()

        for key in validation_dataset_dict:
            dataset_return[key] = validation_dataset_dict[key]

        return dataset_return