2023-02-03

Hugging Face Datasets

Machine Learning

NLP

Hugging Face

Python

Hugging Face Datasets

Hugging FaceのDatasetsでは、データセットの一覧ウェブサイトと、データセットを扱うライブラリであるdatasetsが提供されています。

データセット

公開されているデータセットは次のリンクから確認することができます。

https://huggingface.co/datasets

データセットはキーワード検索できるだけでなく、タスクのカテゴリーやデータサイズ、ライセンスなどでフィルタリングすることができます。

datasets API

datasets APIは次のコマンドでインストールすることができます。

$ pip install datasets

利用可能なデータセット一覧

利用可能なデータセットの一覧はlist_datasetsを使って確認することができます。

from datasets import list_datasets, load_dataset, list_metrics, load_metric
from pprint import pprint

pprint(list_datasets())

2023年2月3日時点では20,267件のデータセットがあります。

>> len(list_datasets)

20267

また、データセットは次のリンクから確認することもできます。

https://huggingface.co/datasets

データセットのロード

データセットはload_datasetを使うことでロードすることができます。

from datasets import load_dataset

dataset = load_dataset(
  'wikitext', # dataset identifier
  'wikitext-103-v1' # instance name in the dataset
)

データセットの型

ロードしたデータセットはDatasetDict型となっています。

from datasets import load_dataset

dataset = load_dataset('glue', 'cola')

type(dataset)

datasets.dataset_dict.DatasetDict

データセットは、次のように様々な型に変換することができます。

dataset.set_format(type='python') # (default) Python object
dataset.set_format(type='torch') # PyTorch tensor
dataset.set_format(type='tensorflow') # TensorFlow tensor
dataset.set_format(type='jax') # JAX
dataset.set_format(type='numpy') # NumPy
dataset.set_format(type='pandas') # pandas DataFrame

https://pypi.org/project/datasets/

CSV ファイルから DatasetDict クラスを生成

DatasetDictクラスはローカルファイルやpandas DataFrameなどから作成することができます。

まずはCSVファイルの準備をします。glueというデータセットを取得してCSVとして保存します。

from datasets import load_dataset
from pprint import pprint

dataset = load_dataset('glue', 'cola')
dataset

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
    })
})

to_csvメソッドでデータセットをCSVとしてローカルに保存します。

dataset.set_format(type="pandas") # convert to pandas
dataset["train"][:].to_csv("train.csv", index=None) # save
dataset["validation"][:].to_csv("validation.csv", index=None) # save
dataset["test"][:].to_csv("test.csv", index=None) # save
dataset.reset_format() # reset format

train.csvは次のようなデータになっています。

$ !cat 'train.csv' | head -n 5

sentence,label,idx
"Our friends won't buy this analysis, let alone the next one we propose.",1,0
One more pseudo generalization and I'm giving up.,1,1
One more pseudo generalization or I'm giving up.,1,2
"The more we study verbs, the crazier they get.",1,3

CSV ファイル

ローカルにあるCSVファイルからDatasetDictクラスを作成するには次のようなコードを実行します。

dataset_files = {
    "train": ["train.csv"],
    "validation": ["validation.csv"],
    "test": ["test.csv"],
}
dataset_from_csv = load_dataset("csv", data_files=dataset_files)
print(dataset_from_csv)

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
    })
})

pandas DataFrame

pandasからDatasetDictクラスを作成するには次のようなコードを実行します。

import pandas as pd
from datasets import Dataset, DatasetDict

train_df = pd.read_csv("train.csv")
validation_df = pd.read_csv("validation.csv")
test_df = pd.read_csv("test.csv")

train_ds = Dataset.from_pandas(train_df)
validation_ds = Dataset.from_pandas(validation_df)

dataset_from_pandas = DatasetDict({
    "train": train_ds,
    "validation": validation_ds,
    "test": test_ds,
})
print(dataset_from_csv)

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
    })
})

label カラムを ClassLabel にキャスト

公開データセットからロードしたdatasetのlabelはClassLabel型である一方、dataset_from_pandasのlabelはValue型となっています。

dataset['train'].features

{'sentence': Value(dtype='string', id=None),
 'label': ClassLabel(names=['unacceptable', 'acceptable'], id=None),
 'idx': Value(dtype='int32', id=None)}

dataset_from_pandas['train'].features

{'sentence': Value(dtype='string', id=None),
 'label': Value(dtype='int64', id=None),
 'idx': Value(dtype='int64', id=None)}

次のようにデータセットのカラムの型をキャストすることができます。

from datasets import ClassLabel

# convert "label" to ClassLabel
class_label = ClassLabel(num_classes=2, names=['unacceptable', 'acceptable'])
dataset_from_pandas = dataset_from_pandas.cast_column("label", class_label)

dataset_from_pandas['train'].features

{'sentence': Value(dtype='string', id=None),
 'label': ClassLabel(names=['unacceptable', 'acceptable'], id=None),
 'idx': Value(dtype='int64', id=None)}

データセットの処理

データセットはmapメソッドを使ってデータセット全体の処理をすることができます。

from datasets import load_dataset
from transformers import AutoTokenizer

dataset = load_dataset('glue', 'cola')

# add "length" column that is the length of "sentence"
dataset_with_length = dataset.map(lambda x: {"length": len(x["sentence"])})

# tokenize "sentence" column
tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')
dataset_tokenized = dataset.map(lambda x: tokenizer(x['sentence']), batched=True)