一个包搞定中文数据集: datasetstore

工作中，总是要使用各种中文数据集，每次使用数据集都要花费不少的时间进行寻找，写预处理代码，结合不同的模型和框架做出相应的处理。有的时候好不容易找到合适的数据集，但是却因为网络问题，无法下载，下载了很长一段时间，突然弹出 timeout。

既浪费时间，也浪费精力。

所以，就决定自己造个轮子，搞定这个问题。

考虑到这个包要能有很好的多框架兼容性，并且还要有很好的性能和源码的架构。找来找去，最终找到了 Huggingface 的 Datasets 库，这个包有着非常好的框架兼容性，性能和源码架构，是一个非常好的解决方案。但是！它依然存在一个问题，由于它采用的存储后端是国外的 AWS S3 和 github 的 LFS，必然的，导致了它的网络非常的不稳定，经常的出各种网络问题。

既然找到了问题，那么也就操刀解决掉它。于是，就有了 datasetstore，把原有的国外存储后端进行了替换，使用国内的存储，并且也针对一些中文数据集的特有情况，做了一些处理，方便我们使用各种各种各样的数据集。

接下来，让我们看看如何快速的使用 datasetstore。

快速开始

pip install datasetstore

一条命令，安装好所有的依赖，就可以直接使用了。

使用方法

from datasetstore import load_dataset, list_datasets

# 打印支持的数据集
print(list_datasets())

# 加载数据及并打印并第一个样本
hotel_review = load_dataset('hotel-review')
print(hotel_review['train'][0])

# 处理数据集 - 给每个样本增加一个文本长度的特征
hotel_review = hotel_review.map(lambda x: {"length": len(x["text"])})

# 结合 transformers 库，快速使用各种模型处理任务
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')

tokenized_dataset = hotel_review.map(lambda x: tokenizer(x['text']), batched=True)

数据加载

在使用 load_dataset() 接口的时候，datasetstore 会从云端下载所需要的数据集，目前支持的数据集在这里。下载截图如下：

也可以调用 list_datasets() 来查看所有已经支持的数据集。

如果有你想用的数据集是不支持的，你可以联系作者（也就是我） zhangchunyang_pri@126.com，包含数据集附件和数据集的使用方法，我来上传上去。后续也会开放接口，让大家自己上传数据集，目前这个功能先在测试中。

数据集的使用

数据集的使用也非常的简单，下面我对一些常用的方法做些介绍，更多的可以参考 HuggingFace Datasets 文档来查看更多细节的使用方法。那如果你觉得看起来费劲，也可以留言给我，我后面会写一些使用的方法。如果觉得原始包有些地方不符合我们的使用习惯，也可以留言给我，我会在后续的迭代中，更新到 datasetstore 中。

# 通过索引获取样本
dataset[0]

# 通过特证明获取特征
dataset["text"]

# 使用数据集切片功能
dataset[:3]

# 数据集排序
dataset.sort("label")

# 数据集打乱
shuffled_dataset = sorted_dataset.shuffle(seed=42)

# 数据集的选择
small_dataset = dataset.select([0, 10, 20, 30, 40, 50])

# 数据集的筛选
start_with_ar = dataset.filter(lambda example: example["sentence1"].startswith("Ar"))
len(start_with_ar)

# 数据集切分
dataset.train_test_split(test_size=0.1)

# 数据集切片
datasets = load_dataset("imdb", split="train")
dataset.shard(num_shards=4, index=0)

# 数据集重命名特征
dataset = dataset.rename_column("text", "sentenceA")

# 数据集特征移除
dataset.remove_columns("label")

# 数据集 map
def add_prefix(example):
    example["sentence1"] = 'My sentence: '' + example["sentence1"]
    return example

updated_dataset = small_dataset.map(add_prefix)

# 多进程处理
updated_dataset = dataset.map(lambda example, idx: {"sentence2": f"{idx}: " + example["sentence2"]}, num_proc=4)