<commit_before>from collections import defaultdict, namedtuple, Mapping

import h5py
import torch
from torch.utils.data import Dataset
import numpy as np
from torch.autograd import Variable


class DataTransform:
    def initialize(self, dataset):
        pass

    def __repr__(self):
        return self.__class__.__name__


class SubsampleNeurons(DataTransform):
    def __init__(self, datakey, idx, axis):
        super().__init__()
        self.idx = idx
        self.datakey = datakey
        self._subsamp = None
        self.axis = axis

    def initialize(self, dataset):
        self._subsamp = []
        for d in dataset.data_keys:
            if d == self.datakey:
                self._subsamp.append([slice(None) for _ in range(self.axis - 1)] + [self.idx, ...])
            else:
                self._subsamp.append(...)

    def __call__(self, item):
        return tuple(it[sub] for sub, it in zip(self._subsamp, item))


class TransformFromFuncs(DataTransform):
    def __init__(self):
        super().__init__()

    def initialize(self, dataset):
        self._transforms = []
        self.transformees = []
        for d in dataset.data_keys:
            if hasattr(dataset, d + '_transform'):
                self._transforms.append(getattr(dataset, d + '_transform'))
                self.transformees.append(d)
            else:
                self._transforms.append(lambda x: x)

    def __call__(self, item):
        return tuple(tr(it) for tr, it in zip(self._transforms, item))<commit_msg>modulator with LM<commit_after>from collections import defaultdict, namedtuple, Mapping

import h5py
import torch
from torch.utils.data import Dataset
import numpy as np
from torch.autograd import Variable


class DataTransform:
    def initialize(self, dataset):
        pass

    def __repr__(self):
        return self.__class__.__name__


class SubsampleNeurons(DataTransform):
    def __init__(self, datakey, idx, axis):
        super().__init__()
        self.idx = idx
        self.datakey = datakey
        self._subsamp = None
        self.axis = axis

    def initialize(self, dataset):
        self._subsamp = []
        for d in dataset.data_keys:
            if d == self.datakey:
                self._subsamp.append([slice(None) for _ in range(self.axis - 1)] + [self.idx, ...])
            else:
                self._subsamp.append(...)

    def __call__(self, item):
        return tuple(it[sub] for sub, it in zip(self._subsamp, item))

class Neurons2Behavior(DataTransform):
    def __init__(self, idx):
        super().__init__()
        self.idx = idx

    def __call__(self, item):
        return tuple((item[0], np.hstack((item[1], item[3][~self.idx])), item[2], item[3][self.idx]))



class TransformFromFuncs(DataTransform):
    def __init__(self):
        super().__init__()

    def initialize(self, dataset):
        self._transforms = []
        self.transformees = []
        for d in dataset.data_keys:
            if hasattr(dataset, d + '_transform'):
                self._transforms.append(getattr(dataset, d + '_transform'))
                self.transformees.append(d)
            else:
                self._transforms.append(lambda x: x)

    def __call__(self, item):
        return tuple(tr(it) for tr, it in zip(self._transforms, item))
{'fraction_non_alphanumeric': 0.11154895422855411, 'fraction_numerical': 0.003031221582297666, 'mean_word_length': 2.3570701932858595, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3053084', 'n_tokens_mistral': 994, 'n_tokens_neox': 966, 'n_words': 248}