Train_test_split du montant de fonction détenu par dict

Chose que tu veux faire

Le montant de la fonction détenue dans le type de dictionnaire comme indiqué ci-dessous

{'field1': array([0, 1, 2, 3, 4, 5]), 'field2': array([5, 4, 3, 2, 1, 0]), 'label': array([1, 0, 1, 0, 0, 0])}

Je veux diviser tout en conservant le type de dictionnaire comme indiqué ci-dessous

{'field1': array([4, 0, 3]), 'field2': array([1, 5, 2]), 'label': array([0, 1, 0])}
{'field1': array([5, 1, 2]), 'field2': array([0, 4, 3]), 'label': array([0, 0, 1])}

manière

Il est assez difficile d'implémenter des options telles que l'extraction compte tenu de la distribution de classe par vous-même, alors créez un processus qui encapsule model_selection.train_test_split de scikit-learn. Cela était possible en créant un tableau d'index aléatoire, en le divisant avec le tableau d'étiquettes avec train_test_split et en extrayant de chaque tableau en spécifiant l'index.

class Splitter:
    def __init__(self, train_size, label_col: str):
        self.train_size = train_size
        self.label_col = label_col
        self.train_indices = None
        self.test_indices = None

    def set_split_indices(self, field_to_values):
        total_length = len(field_to_values[self.label_col])
        split_indices = np.array(range(total_length))
        labels = field_to_values[self.label_col]
        self.train_indices, self.test_indices, _, _ = train_test_split(
            split_indices, labels, train_size=self.train_size,stratify=labels)

    def split(self, field_to_values):
        train = {field: values[self.train_indices] for field, values in field_to_values.items()}
        test = {field: values[self.test_indices] for field, values in field_to_values.items()}
        return train, test

Résultat d'exécution

>>> field_to_values = {"field1": np.array([0, 1, 2, 3, 4, 5]), "field2": np.array([5, 4, 3, 2, 1, 0]), "label": np.array([1, 0, 1, 0, 0, 0])}
>>> splitter = Splitter(train_size=0.5, label_col="label")
>>> splitter.set_split_indices(field_to_values)
>>> splitter.split(field_to_values)
({'field1': array([4, 0, 3]), 'field2': array([1, 5, 2]), 'label': array([0, 1, 0])}, {'field1': array([5, 1, 2]), 'field2': array([0, 4, 3]), 'label': array([0, 0, 1])})

Cette fois, je voulais diviser plusieurs dictionnaires par le même index, donc j'en ai fait une classe, mais si je l'écris comme une fonction, ce sera comme suit

def train_test_split_dict(field_to_values, train_size, label_col: str)
    total_length = len(field_to_values[label_col])
    split_indices = np.array(range(total_length))
    labels = field_to_values[label_col]
    train_indices, test_indices, _, _ = train_test_split(
        split_indices, labels, train_size=train_size, stratify=labels)
    train = {field: values[train_indices] for field, values in field_to_values.items()}
    test = {field: values[test_indices] for field, values in field_to_values.items()}
    return train, test