Class BloomRawDataset

linear_moe/data/bloom.py:21–97

A class for processing a Bloom text dataset

from linear_moe.tokenizer import get_tokenizer

class BloomRawDataset(torch.utils.data.Dataset):
    """A class for processing a Bloom text dataset"""
    def __init__(self, datapaths, max_seq_length):
        """
        Initializes the dataset.
        Args:
            path(str): The path of the dataset file.
            tokenizer(object): The tokenizer object.
            max_seq_length(int): The maximum length of sequences.
        """
        self.tokenizer = get_tokenizer()
        self.max_seq_length = max_seq_length
        self.prompt = ''
        self.samples = []
        for datapath in datapaths:
            self.samples.extend(
                self.process_samples_from_single_path(datapath))
        print('  >> total number of samples: {}'.format(len(self.samples)))

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        raw_sample = self.samples[idx]
        return self.gpt_convert_example_to_feature(raw_sample, self.tokenizer,
                                                   self.max_seq_length)

    def truncate(self, tokenizer, array, max_length):
        """
        Truncates an array to a maximum length or pads it with zeros if its length is less than `max_length`.
        Args:
            tokenizer: The tokenizer used to encode the input.
            array: The numpy array to truncate or pad.
            max_length: The maximum length of the array.
        Returns:
            A numpy array of length `max_length` containing the contents of `array`, truncated if necessary or padded with zeros.
        """

        if len(array) < max_length:
            return np.pad(array, (0, max_length - len(array)),
                          constant_values=tokenizer.eod)
        else:
            return array[:max_length]

    def process_samples_from_single_path(self, filename):
        """
        Process a single file containing prompt-answer pairs and return a list of samples.
        """

        print(' > Processing {} ...'.format(filename))
        samples = []
        total = 0
        with open(filename, encoding='utf-8-sig') as f:
            for example in f:
                text = json.loads(example)['text']
                sample = {
                    'prompt':
                    text + '</s>' if not text.endswith('</s>') else text,