HuggingFace에 여러 태스크로 이루어진 데이터셋 만들기

Hugging Face에는 딥러닝 모델뿐만 아니라 수많은 데이터셋이 있습니다.
허깅페이스 라이브러리를 통해서 이것들을 다운 받을 수도 있고, 새로 만들어서 올릴 수도 있습니다.
꼭 웹사이트에 올리지 않고 로컬에서 작업할 때에도 허깅페이스의 datasets 라이브러리는 유용합니다.
모델 학습 라이브러리와도 잘 붙고 분석하기도 간편합니다.

JSON이나 CSV 파일 있을 때 이를 허깅페이스 datasets 포맷으로 바꾸는 것은 매우 간단합니다. 디렉토리를 만들고 파일을 갖다두면 끝입니다.
예를 들어 mmlu라는 데이터셋이 있고 test.jsonl이란 파일이 있다면 mmlu 디렉토리 아래에 test.jsonl을 옮겨두면 됩니다.

from datasets import load_dataset
dataset = load_dataset("mmlu")

이렇게 간단하게 읽을 수 있습니다. test.jsonl이라는 파일명은 알 필요가 없습니다. 만약 mmlu 디렉토리에 파일이 여러 개라면 모두 읽어들입니다.

잠시 dataset의 내용을 볼까요?

>>> dataset
DatasetDict({
    test: Dataset({
        features: ['text', 'choices', 'answer'],
        num_rows: 500
    })
})

기본적으로 DatasetDict라는 일종의 딕셔너리로 읽어옵니다. 아래와 같이 내용을 볼 수도 있습니다.

>>> dataset['test'][42]
{'text': '생물 다양성을 보전해야 하는 까닭으로 옳지 않은 것은?', 'choices': ['생태계를 안정적으로 유지하기 위해서이다.', '생물 다양성은 그 자체로 중요하기 때문이다.', '사람에게 유용한 생물만 남겨 두기 위해서이다.', '살아가는 데 필수적인 자원을 얻기 위해서이다.', '휴식과 여가 활동을 위한 공간을 제공받기 위해서이다.'], 'answer': 3}

문제집에서 추출한 데이터인가봅니다.

데이터셋은 대개 한 가지 태스크로 되어 있습니다. 그런데 GLUE나 MMLU처럼 여러 가지 태스크를 묶어서 만들어진 커다란 데이터셋도 있습니다.
이것들을 카테고리별로 따로따로 데이터셋으로 만들 수도 있겠지만, 하나의 이름 아래 다룰 좋은 방법이 있을까요?
물론 데이터셋 안에 필드를 추가해서 구분하는 방법도 있지만, 데이터셋의 내용을 건드리지 않고 하고 싶습니다.

이를 위해 허깅페이스에서는 데이터셋 빌더를 제공합니다. 설명을 더 하기 전에 데이터셋 빌더가 있으면 무엇이 가능한지를 먼저 설명해야 할 것 같습니다.
데이터셋을 읽을 때 아래와 같이 태스크명을 제공해서 그 태스크만 읽을 수 있습니다

dataset = load_dataset("mmlu", name="college_mathematics")

이것이 가능하게 하려면 우선 데이터셋 디렉토리 아래에 그 디렉토리명과 똑같은 이름으로 .py 파일을 만듭니다.
예를 들어 데이터셋 이름이 sample이라면 sample/sample.py를 만드는 것입니다. 그리고 내용은 아래와 같이 작성합니다:

import os

import datasets
import json


_CITATION = """\
    논문 등의 인용 서지 정보
"""

_DESCRIPTION = """\
    데이터셋 설명
"""

_HOMEPAGE = "URL"

_LICENSE = "GPL 등의 라이센스 정보"

task_list = [
    "태스크 A",
    "태스크 B",
    "태스크 C",
]


class SampleConfig(datasets.BuilderConfig):
    def __init__(self, **kwargs):
        super().__init__(version=datasets.Version("1.0.0"), **kwargs)


class Sample(datasets.GeneratorBasedBuilder):
    BUILDER_CONFIGS = [
        SampleConfig(
            name=task_name,
        )
        for task_name in task_list
    ]

    def _info(self):
        features = datasets.Features(
            {
                "text": datasets.Value("string"),
                "choices": datasets.Sequence(datasets.Value("string")),
                "answer": datasets.Value("int8"),
            }
        )
        return datasets.DatasetInfo(
            description=_DESCRIPTION,
            features=features,
            homepage=_HOMEPAGE,
            license=_LICENSE,
            citation=_CITATION,
        )

    def _split_generators(self, dl_manager):
        data_dir = "data/sample"
        task_name = self.config.name
        return [
            datasets.SplitGenerator(
                name=datasets.Split.TRAIN,
                gen_kwargs={
                    "filepath": os.path.join(data_dir, "data", f"{task_name}.jsonl"),
                },
            ),
            datasets.SplitGenerator(
                name=datasets.Split.VALIDATION,
                gen_kwargs={
                    "filepath": os.path.join(data_dir, "data", f"{task_name}.jsonl"),
                },
            ),
            datasets.SplitGenerator(
                name=datasets.Split.TEST,
                gen_kwargs={
                    "filepath": os.path.join(data_dir, "data", f"{task_name}.jsonl"),
                },
            ),
        ]

    def _generate_examples(self, filepath):
        with open(filepath, encoding="utf-8") as f:
            for key, row in enumerate(f):
                data = json.loads(row)
                yield key, data

여기서 주의해야 할 점이 몇 가지 있는데요, Features가 정수 타입인 경우 int를 쓰면 에러가 납니다.
꼭 int8, int16과 같이 비트수를 명시해 주어야 합니다.
그리고 리스트 타입인 경우에는 위와 같이 datasets.Sequence(datasets.Value("string")) 이렇게 Sequence를 써야 합니다.
이 경우에는 문자열의 리스트라서 안에 datasets.Value("string")을 인자로 주었습니다.

이 파일을 갖다놓고 코드에 있는대로 data_dir 위치에 실제 데이터를 갖다놓으면 잘 작동합니다.
dl_manager를 사용하고 URL을 제공하면 인터넷으로부터 다운로드할 수도 있습니다.

훌륭한 데이터셋이 만들어진 경우 push_to_hub 기능을 이용해서 허깅페이스에 데이터셋을 올릴 수도 있습니다.