Transformer를 위한 Custom Dataset 만들기
1 min read

Transformer를 위한 Custom Dataset 만들기

Background

  • Transformer에서 제공하는 Pipeline을 이용해서 Sentiment Classifier를 하는데 기존  Dataframe을 그대로 이용하려니 GPU를 활용하더라도 상당히 느려서 Efficiency를 높이기 위해서 Dataset으로 변환을 하기로 결정하였다.

Usage

1.한개의 csv 파일로 만들 때

from datasets import load_dataset
dataset = load_dataset('csv', data_files='my_file.csv')

2.두 개이상의 csv 파일로 만들 때

dataset = load_dataset('csv', data_files=['my_file_1.csv', 'my_file_2.csv', 'my_file_3.csv'])

3.Train, Test 나눠서 만들고자 할 때

dataset = load_dataset('csv', data_files={'train': ['my_train_file_1.csv', 'my_train_file_2.csv'], 'test': 'my_test_file.csv'})

4.Feature를 구체화해서 만들 때

아래과 같이 작성해서 실행하면 다음과 같은 결과를 얻을 수 있다.

from datasets import Features, Value, ClassLabel
from datasets import load_dataset

class_names = ["sadness", "joy", "love", "anger", "fear", "surprise"]
emotion_features = Features({'text': Value('string'), 'label': ClassLabel(names=class_names)})

dataset = load_dataset('csv', data_files=file_dict, delimiter=';', column_names=['text', 'label'], features=emotion_features)
dataset['train'].features
{'text': Value(dtype='string', id=None),
'label': ClassLabel(num_classes=6, names=['sadness', 'joy', 'love', 'anger', 'fear', 'surprise'], names_file=None, id=None)}