Function build_dataset

linear_moe/data/__init__.py:206–222

(index, name)

    # Print stats about the splits.

    def build_dataset(index, name):
        dataset = None
        if splits[index + 1] > splits[index]:
            documents = np.arange(start=splits[index],
                                  stop=splits[index + 1],
                                  step=1,
                                  dtype=np.int32)
            if dataset_type == 'LLama-Pretrain-Idxmap':
                from .llama import LLamaIdxMapDataset
                dataset = LLamaIdxMapDataset(
                    name, data_prefix, documents, indexed_dataset,
                    train_valid_test_num_samples[index],
                    seed, max_padding_length, return_doc_ids)
            else:
                raise RuntimeError("The provided dataset_type is not supported in Pretrain mode. \nChoose from [LLama-Pretrain-Idxmap].")

        return dataset

    train_dataset = build_dataset(0, 'train')
    valid_dataset = build_dataset(1, 'valid')