debeir.evaluation.cross_validation

  1from enum import Enum
  2from typing import Dict, List, Union
  3
  4import numpy as np
  5from debeir.datasets.types import DatasetTypes, InputExample
  6from sklearn.model_selection import KFold, StratifiedKFold
  7
  8import datasets
  9
 10
 11def split_k_fold(n_fold, data_files):
 12    percentage = 100 // n_fold
 13
 14    vals_ds = datasets.load_dataset('csv', split=[
 15        f'train[{k}%:{k + percentage}%]' for k in range(0, 100, percentage)
 16    ], data_files=data_files)
 17
 18    trains_ds = datasets.load_dataset('csv', split=[
 19        f'train[:{k}%]+train[{k + percentage}%:]' for k in range(0, 100, percentage)
 20    ], data_files=data_files)
 21
 22    return trains_ds, vals_ds
 23
 24
 25class CrossValidatorTypes(Enum):
 26    """
 27    Cross Validator Strategies for separating the dataset
 28    """
 29    Stratified = "StratifiedKFold"
 30    KFold = "KFold"
 31
 32
 33str_to_fn = {
 34    "StratifiedKFold": StratifiedKFold,
 35    "KFold": KFold
 36}
 37
 38
 39class CrossValidator:
 40    """
 41    Cross Validator Class for different types of data_sets
 42
 43    E.g. List -> [[Data], label]
 44         List[Dict] -> {"data": Data, "label": label}
 45         Huggingface Dataset Object -> Data(set="train", label = "label").select(idx)
 46    """
 47
 48    def __init__(self, dataset: Union[List, List[Dict], datasets.Dataset],
 49                 x_idx_label_or_attr: Union[str, int], y_idx_label_or_attr: Union[str, int],
 50                 cross_validator_type: [str, CrossValidatorTypes] = CrossValidatorTypes.Stratified,
 51                 seed=42, n_splits=5):
 52        # self.evaluator = evaluator
 53        self.cross_vali_fn = str_to_fn[cross_validator_type](n_splits=n_splits,
 54                                                             shuffle=True,
 55                                                             random_state=seed)
 56        self.dataset = dataset
 57        self.splits = []
 58
 59        self.x_label = x_idx_label_or_attr
 60        self.y_label = y_idx_label_or_attr
 61
 62        if self.dataset_type is None:
 63            self._determine_dataset_type()
 64            x, y = self.split_fn(x_idx_label_or_attr, y_idx_label_or_attr)
 65            self.splits = self.cross_vali_fn.split(x, y)
 66
 67    def _determine_dataset_type(self):
 68        if isinstance(self.dataset, list):
 69            if isinstance(self.dataset[0], dict):
 70                self.dataset_type = DatasetTypes.ListDict
 71                self.split_fn = self._split_dict
 72            elif isinstance(self.dataset[0], InputExample):
 73                self.dataset_type = DatasetTypes.ListInputExample
 74                self.split_fn = self._split_list
 75            else:
 76                self.dataset_type = DatasetTypes.List
 77                self.split_fn = self._split_list
 78        elif isinstance(self.dataset, datasets.Dataset):
 79            self.dataset_type = DatasetTypes.HuggingfaceDataset
 80            self.split_fn = self._split_dataset
 81        else:
 82            raise NotImplementedError("Unknown Dataset format")
 83
 84    def _split_list(self, *args, **kwargs):
 85        X = np.zeros(len(list(map(lambda k: k[self.x_label], self.dataset))))
 86        Y = map(lambda k: k[self.y_label], self.dataset)
 87
 88        return X, Y
 89
 90    def _split_dict(self, *args, **kwargs):
 91        X = np.zeros(len(list(map(lambda k: k[self.x_label], self.dataset))))
 92        Y = map(lambda k: k[self.y_label], self.dataset)
 93
 94        return X, Y
 95
 96    def _split_dataset(self, *args, **kwargs):
 97        # Rows data doesn't matter
 98        X = np.zeros(self.dataset.num_rows)
 99        Y = self.dataset[self.y_label]
100
101        return X, Y
102
103    def get_fold(self, fold_num: int):
104        """
105
106        :param fold_num: Which fold to pick
107        :return:
108        """
109
110        split = self.splits[fold_num]
111
112        return {
113            "train_idxs": split[0],
114            "val_idxs": split[1]
115        }