How to Customize a Dataset
============

By using :func:`load_dataset`, :class:`MapDataset`, and :class:`IterDataset` provided by PaddleNLP, anyone can easily define their own dataset.

Creating Dataset from Local Files
-------------------

When creating a dataset from local files, we **recommend** providing a reader function based on the local data format and passing it to :func:`load_dataset` to create the dataset.

.. code-block::

    from paddlenlp.datasets import load_dataset

    def read(data_path):
        with open(data_path, 'r', encoding='utf-8') as f:
            # Skip header
            next(f)
            for line in f:
                words, labels = line.strip('\n').split('\t')
                words = words.split('\002')
                labels = labels.split('\002')
                yield {'tokens': words, 'labels': labels}

    # data_path is the parameter for read()
    map_ds = load_dataset(read, data_path='train.txt', lazy=False)
    iter_ds = load_dataset(read, data_path='train.txt', lazy=True)

We recommend writing data reading code as a generator, which better constructs both :class:`MapDataset` and :class:`IterDataset`. We also suggest formatting each data sample as a dictionary, which facilitates data flow monitoring.

In fact, :class:`MapDataset` can meet requirements in most cases. Generally, :class:`IterDataset` is only considered when the dataset is too large to be loaded into memory at once.

.. note::

    Note that only PaddleNLP's built-in datasets support automatic label-to-id conversion (detailed conditions refer to :doc:`Creating DatasetBuilder <../community/contribute_datasets/how_to_write_a_DatasetBuilder>`).
    
    For custom datasets like the above example, you need to manually add label-to-id functionality in your custom feature conversion method.

    Parameters from custom data reader functions can be directly passed to :func:`load_dataset` as keyword arguments. For custom datasets, the :attr:`lazy` parameter is **required**.

Creating from :class:`paddle.io.Dataset/IterableDataset`
Creating a Dataset
-------------------

While PaddlePaddle's built-in :class:`Dataset` and :class:`IterableDataset` can be directly used with :class:`DataLoader` for model training, sometimes we need more convenient ways to perform data processing (e.g., convert to features, data cleaning, data augmentation). PaddleNLP's built-in :class:`MapDataset` and :class:`IterDataset` provide APIs to implement these functions.

If you're accustomed to using :class:`paddle.io.Dataset/IterableDataset` to create datasets, you can simply wrap them with :class:`MapDataset` or :class:`IterDataset` to convert them into PaddleNLP dataset objects.

Here's a simple example. The usage of :class:`IterDataset` is similar.

.. code-block::

    from paddle.io import Dataset
    from paddlenlp.datasets import MapDataset

    class MyDataset(Dataset):
        def __init__(self, path):

            def load_data_from_source(path):
                ...
                ...
                return data

            self.data = load_data_from_source(path)

        def __getitem__(self, idx):
            return self.data[idx]

        def __len__(self):
            return len(self.data)
    
    ds = MyDataset(data_path)  # paddle.io.Dataset
    new_ds = MapDataset(ds)    # paddlenlp.datasets.MapDataset

Creating Dataset from Other Python Objects
-------------------

In theory, we can use any Python object containing :func:`__getitem__` and :func:`__len__` methods to create a :class:`MapDataset`, including :class:`List`, :class:`Tuple`, :class:`DataFrame`, etc. Simply pass the qualified Python object as initialization parameter to :class:`MapDataset`.
.. code-block::

    from paddlenlp.datasets import MapDataset

    data_source_1 = [1,2,3,4,5]
    data_source_2 = ('a', 'b', 'c', 'd')

    list_ds = MapDataset(data_source_1)
    tuple_ds = MapDataset(data_source_2)

    print(list_ds[0])  # 1
    print(tuple_ds[0]) # a

Similarly, we can create :class:`IterDataset` using Python objects containing :func:`__iter__` method, such as :class:`List` and :class:`Generator`. The creation method is the same as :class:`MapDataset`.

.. code-block::

    from paddlenlp.datasets import IterDataset

    data_source_1 = ['a', 'b', 'c', 'd']
    data_source_2 = (i for i in range(5))

    list_ds = IterDataset(data_source_1)
    gen_ds = IterDataset(data_source_2)

    print([data for data in list_ds]) # ['a', 'b', 'c', 'd']
    print([data for data in gen_ds])  # [0, 1, 2, 3, 4]

.. note::

    Note that when passing a **generator** object to :class:`IterDataset` as shown in the above example, the dataset will automatically reset the generator when iterating over it for the second time. To reuse the generated data, you need to convert the generator to a list first.
The dataset generated by `IterDataset` can only be iterated **once**.

Like regular Python objects, as long as the above conditions are met, we can use the same method to create PaddleNLP datasets from third-party datasets.

For example, HuggingFace Dataset:

.. code-block::

    from paddlenlp.datasets import MapDataset
    from datasets import load_dataset
    
    hf_train_ds = load_dataset('msra_ner', split='train')
    print(type(train_ds)) # <class 'datasets.arrow_dataset.Dataset'>

    train_ds = MapDataset(train_ds)
    print(type(train_ds)) # <class 'paddlenlp.datasets.dataset.MapDataset'>

    print(train_ds[2]) # {'id': '2', 
                       #  'ner_tags': [0, 0, 0, 5, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                       #               0, 0, 0, 0, 0, 0, 5, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0], 
                       #  'tokens': ['因', '有', '关', '日', '寇', '在', '京', '掠', '夺', '文', '物',
                       #             '详', '情', '，', '藏', '界', '较', '为', '重', '视', '，', '也', 
                       #             '是', '我', '们', '收', '藏', '北', '京', '史', '料', '中', '的',
                       #             '要', '件', '之', '一', '。']}

    hf_train_ds = load_dataset('cmrc2018', split='train')
    train_ds = MapDataset(hf_train_ds)
    print(train_ds[1818]) # {'answers': {'answer_start': [9], 'text': ['字仲可']}, 
                          #  'context': '徐珂（），原名昌，字仲可，浙江杭县（今属杭州市）人。光绪举人。
                          #              后任商务印书馆编辑。参加南社。1901年在上海担任了《外交报》、
                          #              《东方杂志》的编辑，1911年，接管《东方杂志》的“杂纂部”。与潘仕成、
                          #              王晋卿、王辑塘、冒鹤亭等友好。编有《清稗类钞》、《历代白话诗选》、
                          #              《古今词选集评》等。光绪十五年（1889年）举人。后任商务印书馆编辑。
                          #              参加南社。曾担任袁世凯在天津小站练兵时的幕僚，不久离去。', 
                          #  'id': 'TRAIN_113_QUERY_0', 
                          #  'question': '徐珂字什么？'}
    
    hf_train_ds = load_dataset('glue', 'sst2', split='train')
    train_ds = MapDataset(hf_train_ds)
    print(train_ds[0]) # {'idx': 0, 'label': 0, 'sentence': 'hide new secretions from the parental units '}

    hf_train_ds = load_dataset('ptb_text_only', split='train')
    train_ds = MapDataset(hf_train_ds)
    print(train_ds[1]) # {'sentence': 'pierre <unk> N years old will join the board as a nonexecutive director nov. N'}