removed split dataset

megatron/data/split_dataset.py

-# coding=utf-8
-# Copyright (c) 2019, NVIDIA CORPORATION.  All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""dataset to split one large one into multiple smaller datasets"""
-
-import torch
-import numpy as np
-
-def get_train_valid_test_split(splits_string, size):
-    """ Get dataset splits from comma or '/' separated string list."""
-
-    splits = []
-    if splits_string.find(',') != -1:
-        splits = [float(s) for s in splits_string.split(',')]
-    elif splits_string.find('/') != -1:
-        splits = [float(s) for s in splits_string.split('/')]
-    else:
-        splits = [float(splits_string)]
-    while len(splits) < 3:
-        splits.append(0.)
-    splits = splits[:3]
-    splits_sum = sum(splits)
-    assert splits_sum > 0.0
-    splits = [split/splits_sum for split in splits]
-    splits_index = [0]
-    for index, split in enumerate(splits):
-        splits_index.append(splits_index[index] +
-                            int(round(split * float(size))))
-    diff = splits_index[-1] - size
-    for index in range(1, len(splits_index)):
-        splits_index[index] -= diff
-    return splits_index
-
-class SplitDataset(torch.utils.data.Dataset):
-    """
-    Dataset wrapper to access a subset of another dataset.
-    Purpose: useful to index into existing datasets, possibly
-    large-scale datasets as the subindexing operation is done in an
-    on-the-fly manner.
-    Arguments:
-        ds (Dataset or array-like): List of datasets to be subindexed
-        split_inds (1D array-like): List of indices part of subset
-    """
-    def __init__(self, ds, split_inds, **kwargs):
-        self.split_inds = list(split_inds)
-        self.wrapped_data = ds
-
-    def __len__(self):
-        return len(self.split_inds)
-
-    def __getitem__(self, index):
-        return self.wrapped_data[self.split_inds[index]]
-
-    def num_tokens(self):
-        return self.wrapped_data.num_tokens()
-
-    def __iter__(self):
-        for idx in self.split_inds:
-            yield self.wrapped_data[idx]
-
-def split_ds(ds, split=[.8,.2,.0], shuffle=True):
-    """
-    Split a dataset into subsets given proportions of how
-    much to allocate per split. If a split is 0% returns None for that split.
-    Purpose: Useful for creating train/val/test splits
-    Arguments:
-        ds (Dataset or array-like): Data to be split.
-        split (1D array-like): proportions to split `ds`. `sum(splits) != 0`
-        shuffle (boolean): Randomly split dataset. Default: True
-    """
-    split_sum = sum(split)
-    if split_sum == 0:
-        raise Exception('Split cannot sum to 0.')
-    split = np.array(split)
-    split /= split_sum
-    ds_len = len(ds)
-    inds = np.arange(ds_len)
-    if shuffle:
-        np.random.shuffle(inds)
-    start_idx = 0
-    residual_idx = 0
-    rtn_ds = [None]*len(split)
-    for i, f in enumerate(split):
-        if f != 0:
-            proportion = ds_len*split[i]
-            residual_idx += proportion % 1
-            split_ = int(int(proportion) + residual_idx)
-            split_inds = inds[start_idx:start_idx+max(split_, 1)]
-            rtn_ds[i] = SplitDataset(ds, split_inds)
-            start_idx += split_
-            residual_idx %= 1
-    return rtn_ds