我们了解下如何使用 Datasets 库来加载、处理数据集。安装命令如下：

pip install datasets

1. 加载数据集

Datasets 库可以加载在线数据集：https://huggingface.co/datasets，也可以加载本地数据集。加载 local 数据集，请参考下图：

下面为代码演示：

from datasets import load_dataset
import pandas as pd
from sklearn.model_selection import train_test_split


def create_local_data():

    # 读取数据文件
    data = pd.read_csv('data/ChnSentiCorp_htl_all.csv')
    # 将数据文件拆分训练集和测试集
    x_train, x_valid, y_train, y_valid = \
        train_test_split(data['review'], data['label'], test_size=0.2, stratify=data['label'])


    train = pd.concat([y_train, x_train], axis=1)
    valid = pd.concat([y_valid, x_valid], axis=1)

    train.to_csv('data/senti-train.csv')
    valid.to_csv('data/senti-valid.csv')


if __name__ == '__main__':

    # 1. 加载在线数据集
    train_data = load_dataset("Datatang/mandarin_chinese", split='train')
    print('在线加载:\n', train_data)

    # 2. 加载本地数据集
    create_local_data()
    data_files = {'train': 'data/senti-train.csv', 'valid': 'data/senti-valid.csv'}
    train_data = load_dataset('csv', data_files=data_files)
    print('本地加载:\n', train_data)

程序运行结果：

在线加载:
 Dataset({
    features: ['text'],
    num_rows: 2
})

本地加载:
 DatasetDict({
    train: Dataset({
        features: ['Unnamed: 0', 'label', 'review'],
        num_rows: 6212
    })
    valid: Dataset({
        features: ['Unnamed: 0', 'label', 'review'],
        num_rows: 1554
    })
})

我们发现数据都被封装成了 Dataset 类型的对象。

2. Dataset 基本使用

import pandas as pd
from datasets import load_dataset
from datasets import Dataset


# 1. 基本使用
def test():

    train_data = load_dataset('csv',
                              data_files='data/senti-train.csv',
                              split='train')

    # 重命名列名
    train_data = train_data.rename_column('Unnamed: 0', 'id')
    print('重命名列名:', train_data)
    print('-' * 50)

    # 打印数据集大小
    print('数据集大小:', train_data.dataset_size)
    print('-' * 50)

    # 打乱数据集
    print('打乱前索引:', train_data['id'][:10])
    train_data_shuffle = train_data.shuffle()
    print('打乱后索引:', train_data_shuffle['id'][:10])
    print('-' * 50)

    # 数据集选择
    print('选择指定数据:', train_data.select([1, 4, 6, 9]))
    print('-' * 50)


if __name__ == '__main__':
    test()

程度输出结果：

重命名列名: Dataset({
    features: ['id', 'label', 'review'],
    num_rows: 6212
})
--------------------------------------------------
数据集大小: 2447355
--------------------------------------------------
打乱前索引: [3499, 1940, 5210, 3585, 7326, 7485, 6182, 7416, 3767, 4126]
打乱后索引: [1998, 6766, 7182, 1710, 2869, 4385, 1132, 7023, 2276, 5867]
--------------------------------------------------
选择指定数据: Dataset({
    features: ['id', 'label', 'review'],
    num_rows: 4
})
--------------------------------------------------

3. Dataset 数据处理

import pandas as pd
from datasets import load_dataset
from datasets import Dataset


def test():

    train_data = load_dataset('csv',
                              data_files='data/senti-train.csv',
                              split='train')

    # 1. 数据过滤
    def filter_hanlder(data):
        return (data['review'] is not None) and (len(data['review']) < 10)

    # 默认向回调函数传递一条数据，当设置 batched 时，会一次传输人 batch_size 个数据
    # train_data_filer = train_data.filter(filter_hanlder, batched=True, batch_size=10)
    # 过滤出不为 None 并且长度小于 10 的数据
    print(train_data.num_rows)  # 6212
    train_data_filer = train_data.filter(filter_hanlder)
    print(train_data_filer.num_rows)  # 8

    # 2. 数据处理
    def map_handler(data):
        return {'id': data['Unnamed: 0'],
                'labels': [label + 10 for label in data['label']],
                'review': data['review']}

    train_data_map = train_data.map(map_handler, batched=True, batch_size=10, num_proc=10)
    print(train_data_map['labels'][:5])  # [11, 11, 11, 11, 10]


if __name__ == '__main__':
    test()

4. Dataset 类型转换

import pandas as pd
from datasets import load_dataset
from datasets import Dataset


def test():

    train_data = load_dataset('csv',
                              data_files='data/senti-train.csv',
                              split='train')

    # 在我们的数据集中包含了 None 缺失数据，我们使用 pandas 对其进行处理
    # 1. 将 Dataset 后端存储从字典转换为 pandas DataFrame 类型
    # 使用 [:] 语法可以获得所有数据
    print(type(train_data[:]))  #  <class 'dict'>
    train_data.set_format('pandas')
    print(type(train_data[:]))  #  <class 'pandas.core.frame.DataFrame'>

    # 去除包含缺失值的行
    print(train_data)
    new_train_data = train_data[:].dropna()
    new_train_data = Dataset.from_pandas(new_train_data)
    print(new_train_data)

    # 2. 将 DataFrame 后端存储转换为字典
    train_data.reset_format()
    print(type(train_data[:]))  #  <class 'dict'>


if __name__ == '__main__':
    test()

5. Dataset 数据集分割

import pandas as pd
from datasets import load_dataset
from datasets import Dataset


def test():

    train_data = load_dataset('csv',
                              data_files='data/senti-train.csv',
                              split='train')

    new_dataset = train_data.train_test_split(test_size=0.2)
    """
    DatasetDict
    ({
        train: Dataset({
            features: ['Unnamed: 0', 'label', 'review'],
            num_rows: 4969
        })
        test: Dataset({
            features: ['Unnamed: 0', 'label', 'review'],
            num_rows: 1243
        })
    })
    """
    # Dataset 对象多了一个 test 数据集，该函数没有 stratified 参数
    print(new_dataset)


if __name__ == '__main__':
    test()

6. Dataset 大数据

如果我们加载的数据集非常大，全部加载到内存会带来不小的压力，我们可以是流的方式来加载数据。datasets 库的 load_dataset 函数提供了 steaming 参数，只需要将该参数设置为 True，load_dataset 函数会返回一个可迭代对象，该对象会建立物理设备到程序的内存映射，会一次性加载部分数据到内存，当内存中数据不足时，再从磁盘进行加载。

数据存储的话，我们可以使用下面的方法：

Dataset.save_to_disk
Dataset.to_csv
Dataset.to_json