openmind/docs/zh/basic_tutorial/load_dataset.md

数据加载

目前openMind Library集成适配了datasets，可用于加载魔乐社区上的数据集。openmind.omdatasets.OmDataset类基于datasets抽取了下载相关代码，依托于patch的方式适配了openmind_hub，支持从魔乐社区下载数据集。采用OmDataset.load_dataset('repo_name/dataset_name')的方式可以下载数据集，返回特定的datasets.dataset_dict.DatasetDict类。

目前该类仅支持在PyTorch框架下使用，暂不支持MindSpore框架。

数据集下载

OmDataset.load_dataset()方法目前支持下载的数据集格式如下：

• parquet
• json或者jsonl
• tar.gz
• csv
• txt
• 下载python脚本加载魔乐社区数据集
• 下载python脚本加载三方站点数据集

具体示例如下。

导入对应的类

openmind已暴露Omdataset接口，可直接从openmind中导入。

from openmind import OmDataset

如果有其他开发需求，也可基于相对路径导入。

from openmind.omdatasets import OmDataset

下载parquet格式数据集

OmDataset支持下载parquet格式数据集，以下为用例：

data = OmDataset.load_dataset("modeltesting/test_parquet")
print(data)

'''
DatasetDict({
   test: Dataset({
       features: ['input', 'output'],
       num_rows: 2
   })
})
'''

print(data['test']['input'])

'''
['input_label1', 'input2_label1']
'''

下载json或jsonl格式数据集

OmDataset下载数据集支持json或jsonl格式，以下为用例：

data = OmDataset.load_dataset("modeltesting/test_json")
print(data)

'''
DatasetDict({
   test: Dataset({
       features: ['input', 'output', 'instruction'],
       num_rows: 2
   })
})
'''

print(data['test']['instruction'])

'''
['Create an array of length 2 which .....
'''

下载targ.gz格式数据集

OmDataset下载数据集支持tar.gz格式，以下为用例：

data = OmDataset.load_dataset("modeltesting/test_targz")
print(data)

'''
DatasetDict({
   test: Dataset({
       features: ['__key__', '__url__', 'jsonl'],
       num_rows: 1
   })
})
'''

print(data['test']['jsonl'])

'''
[b'{"chosen": "\\n\\nHuman: Do you know why turkeys became the official food of thanksgiving?....
'''

下载csv格式数据集

OmDataset下载数据集支持csv格式，以下为用例：

data = OmDataset.load_dataset("modeltesting/test_csv")
print(data)

'''
DatasetDict({
   test: Dataset({
       features: ['input', 'output'],
       num_rows: 2
   })
})
'''

print(data['test']['input'])

'''
['input1', 'input2']
'''

下载python脚本加载魔乐社区数据集

支持从魔乐社区数据集仓库下载python脚本，并根据python脚本进一步下载数据集文件，适用于对数据集进行更细化的处理操作。

python脚本需要带有魔乐社区对应文件的下载链接，以test_py_image的test.py脚本为例：

def _split_generators(self, dl_manager):
    DL_URLS = [
        f"https://modelers.cn/api/v1/file/modeltesting/test_py_image/main/media/{name}"
        for name in ["test.jpg"]
    ]
    archive_path = dl_manager.download_and_extract(DL_URLS)
    return [
        datasets.SplitGenerator(
            name=datasets.Split.TEST,
            gen_kwargs={"archive_path": archive_path},
        ),
    ]

从相关代码中可见，涉及下载对应文件时，URL指定了魔乐社区文件下载接口路径，代码里https://modelers.cn/api/v1/file/modeltesting/test_py_image/main/media/XXX 指定了文件下载的接口路径，然后使用该接口路径可以自定义下载。该接口路径可以通过om_hub_url拼接获取。

from openmind_hub import om_hub_url
om_hub_url(repo_id="repo_name/datasets_name", filename="test.json")

'''
https://modelers.cn/api/v1/file/repo_name/datasets_name/main/media/test.json
'''

可尝试下载以下用例，请注意使用python脚本下载数据集文件时需要添加参数trust_remote_code=True。

data = OmDataset.load_dataset("modeltesting/test_py_image")
print(data)

'''
DatasetDict({
   test: Dataset({
       features: ['image'],
       num_rows: 1
   })
})
'''

print(data['test']['image'])

'''
[<PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=1788x688 at 0x7FCB56B6A3D0>]
'''

下载python脚本加载三方站点数据集

支持从魔乐社区数据集仓库下载python脚本，并根据python脚本从第三方站点进一步下载数据集相关文件。以下为python脚本示例相关代码：

_URL = "https://rajpurkar.github.io/SQuAD-explorer/dataset/"
_URLS = {
    "train": _URL + "train-v1.1.json",
    "dev": _URL + "dev-v1.1.json",
}

...

def _split_generators(self, dl_manager):
    downloaded_files = dl_manager.download_and_extract(_URLS)
    return [
        datasets.SplitGenerator(name=datasets.Split.TRAIN, gen_kwargs={"filepath": downloaded_files["train"]}),
        datasets.SplitGenerator(name=datasets.Split.VALIDATION, gen_kwargs={"filepath": downloaded_files["dev"]}),
    ]

从相关代码中可见，涉及下载对应文件时，URL指定了第三方站点。可尝试以下用例进行下载，请注意使用python脚本下载数据集文件时需要添加参数trust_remote_code=True。

data = OmDataset.load_dataset("modeltesting/test_py_github", trust_remote_code=True)
print(data)

'''
DatasetDict({
  train: Dataset({
      features: ['id', 'title', 'context', 'question', 'answers'],
       num_rows: 87599
   })
   validation: Dataset({
       features: ['id', 'title', 'context', 'question', 'answers'],
       num_rows: 10570
   })
})
'''

print(data['train']['question'][0])

'''
To whom did the Virgin Mary allegedly appear in 1858 in Lourdes France?
'''

常用参数

下载数据集常用的参数涉及到数据集的metadata，该部分内容在上传数据集时写在README.md上方，代表用户加载数据集时可以根据特定配置选取加载其中部分数据，OmDataset基于datastes通过patch方式运行，因此metadata沿用了datastes相关配置，更多撰写方式请查看DatasetCard。以下为部分示例。

配置name参数

通过配置name参数，可以下载数据集特定子数据集。该参数与metadata里面的config_name相关，该参数的设置主要方便用户处理以下两种情况：

• 多个文件夹内数据集格式不一致。
• 数据字段情况下，数据集上传用户基于使用目的做了区分。

以modeltesting/test_parquet为例，以下为该数据集的metadata:

license: apache-2.0
configs:
- config_name: label_1
  data_files:
  - split: test
    path: label_1/test*
- config_name: label_2
  data_files:
  - split: test
    path: label_2/test*
- config_name: label_3
  data_files:
  - split: test
    path: label_3/test*

该数据集在仓库内的格式如下：

modeltesting/test_parquet
│
├── label_1
│   ├── test_label1.parquet
│
├── label_2
│   ├── test_label2.parquet
│
├── label_3
│   ├── test_label3.parquet
│
│── README.md

用户可以设置下载文件夹label_2内的数据集：

data = OmDataset.load_dataset(path="modeltesting/test_parquet", name="label_2")
print(data)

'''
DatasetDict({
   test: Dataset({
       features: ['input', 'output'],
       num_rows: 2
   })
})
'''

print(data['test']['input'])

'''
['input1_label2', 'input2_label2']
'''

需要注意的是，当数据集有多个文件夹且配置了metadata里的configs时，如果不设置name参数，默认下载第一个文件夹，如OmDataset.load_dataset(path="modeltesting/test_parquet")下载的数据集是label_1文件夹内的内容。

配置split参数

可设置下载后加载文件夹中对应的文件，一般作用是方便训练过程中区分train,valid, test文件。该参数的写法可查看前文metadata中split示例。是否设置split参数后返回的内容区别如下：

data = OmDataset.load_dataset(path="modeltesting/test_parquet", name="label_2")
print(data)

'''
DatasetDict({
   test: Dataset({
       features: ['input', 'output'],
       num_rows: 2
   })
})
'''

data_split = OmDataset.load_dataset(path="modeltesting/test_parquet", name="label_2", split="test")
print(data_split)

'''
Dataset({
   features: ['input', 'output'],
   num_rows: 2
})
'''

基于代码示例及输出结果可见，设置了split参数后，会直接返回下载的数据中的特定文件。使用社区用户贡献的数据集AI_Connect/glue可以更明确查看区别。

data = OmDataset.load_dataset(path="AI_Connect/glue", name="cola")
print(data)

'''
DatasetDict({
   train: Dataset({
       features: ['sentence', 'label', 'idx'],
       num_rows: 8551
   })
   validation: Dataset({
       features: ['sentence', 'label', 'idx'],
       num_rows: 1043
   })
   test: Dataset({
       features: ['sentence', 'label', 'idx'],
       num_rows: 1063
   })
})
'''

data_split = OmDataset.load_dataset(path="AI_Connect/glue", name="cola", split="test")
print(data_split)

'''
Dataset({
    features: ['sentence', 'label', 'idx'],
    num_rows: 1063
})
'''

配置data_dir参数

该参数可以指定下载数据集仓库内主目录下的某个文件夹内数据集相关文件。以社区用户贡献的AI_Connect/hh-rlhf为例，该数据集README.md中未配置name参数相关的metadata。所以实际运行的时候会报错：

data = OmDataset.load_dataset(path="AI_Connect/hh-rlhf", name="helpful-base")

'''
ValueError: BuilderConfig 'helpful-base' not found. Available: ['default']
'''

但是我们可以通过data_dir参数, 指定下载特定文件夹：

data = OmDataset.load_dataset(path="AI_Connect/hh-rlhf", data_dir="helpful-base")
print(data)

'''
DatasetDict({
    train: Dataset({
        features: ['chosen', 'rejected'],
        num_rows: 43835
    })
    test: Dataset({
        features: ['chosen', 'rejected'],
        num_rows: 2354
    })
})
'''

配置data_files参数

该参数配置可下载特定文件，以社区用户贡献的AI_Connect/hh-rlhf为例。

data = OmDataset.load_dataset(path="AI_Connect/hh-rlhf", data_dir="helpful-base")
print(data)

'''
DatasetDict({
    train: Dataset({
        features: ['chosen', 'rejected'],
        num_rows: 43835
    })
    test: Dataset({
        features: ['chosen', 'rejected'],
        num_rows: 2354
    })
})
'''

data = OmDataset.load_dataset(path="AI_Connect/hh-rlhf", data_dir="helpful-base", data_files="train.jsonl.gz")
print(data)

'''
DatasetDict({
    train: Dataset({
        features: ['chosen', 'rejected'],
        num_rows: 43835
    })
})
'''

从示例中可见，设置data_files后可以指定返回特定数据文件。

配置dataset_info_only参数

配置该参数可以只返回相关数据信息，而不实际下载数据。

data = OmDataset.load_dataset(path="modeltesting/test_parquet", dataset_info_only=True)
print(data)

'''
{'label_1': ['test'], 'label_2': ['test'], 'label_3': ['test']}
'''

配置trust_remote_code参数

如果下载python文件进行数据集下载时，脚本中涉及复杂的处理逻辑或远程自定义代码，需要设置该参数。如果python脚本仅仅加载了对应的json等文件时（不进一步处理），可以忽略该参数。

OmDataset.load_dataset(path="repo_name/dataset_name", trust_remote_code=True)

配置token参数

可下载私仓数据集。

OmDataset.load_dataset(path="repo_name/dataset_name", token="xx")

配置cache_dir参数

可以指定缓存路径。

OmDataset.load_dataset(path="repo_name/dataset_name", cache_dir="/xx/xx")

配置streaming参数

如果需要使用数据集流, 只需要将streaming=True参数传递给load_dataset()函数。streaming=True返回的对象是一个IterableDataset。

data = OmDataset.load_dataset(path="modeltesting/test_parquet", streaming=True)
print(data)

'''
IterableDatasetDict({
    test: IterableDataset({
        features: ['input', 'output'],
        n_shards: 1
    })
})
'''

IterableDataset类可被迭代式获取使用。

next(iter(data['test']))

'''
{'input': 'input_label1', 'output': 'output1_label1'}
'''

数据集使用

OmDataset基于datasets开发，下载数据集返回的值可使用datastes相关函数。以下为示例。

• 下载数据集

from openmind import OmDataset
data = OmDataset.load_dataset(path="modeltesting/test_parquet", name="label_2", split="test")
print(data)

'''
Dataset({
    features: ['input', 'output'],
    num_rows: 2
})
'''

• 加载Tokenizer

from openmind import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("PyTorch-NPU/bert_base_uncased")

• 处理数据集

def tokenization(example):
    return tokenizer(example["input"])

# OmDataset里已经运行了import datasets，所以可以直接使用map方法
data_res = data.map(tokenization, batched=True)
print(data_res)

'''
Dataset({
    features: ['input', 'output', 'input_ids', 'token_type_ids', 'attention_mask'],
    num_rows: 2
})
'''