Pytorch Dataset 클래스(상속) 파악하기!!

from torch.utils.data import Dataset
import gluonnlp as nlp

class BERT_Dataset(Dataset):


    # __init__ 함수는 Dataset 객체가 생성(instantiate)될 때 한 번만 실행됩니다. 
    # 여기서는 dataset, sent_idx, label_idx, bert_tokenizer, max_len, pad, pair)를 불러옵니다.
    # 함수를 호출하게되면 nlp.data.BERTSentenceTransform을 하고, sentences와 labels의 리스트를 만들게 됩니다.
    def __init__(self, dataset, sent_idx, label_idx, bert_tokenizer, max_len,
                 pad, pair):
        transform = nlp.data.BERTSentenceTransform(
            bert_tokenizer, max_seq_length = max_len, pad = pad, pair = pair)

        self.sentences = [transform([i[sent_idx]]) for i in dataset]
        self.labels = [np.int32(i[label_idx]) for i in dataset]


    # 인덱스를 기반으로 문장과 라벨을 반환합니다. return에 슬라이싱 "[]"이 필요함.
    def __getitem__(self, i):
        return (self.sentences[i] + (self.labels[i], ))


    # __len__ 함수는 데이터셋의 레이블 개수를 반환합니다.
    def __len__(self):
        return (len(self.labels))

######################################################################

#사용하기
tokenizer = get_tokenizer()
token = nlp.data.BERTSPTokenizer(tokenizer, vocab, lower=False)

# 이런식으로 만든 함수를 사용합니다.
data_train = BERTDataset(dataset = train_set_data, sent_idx = 0, label_idx = 1, bert_tokenizer = token, max_len = 64, pad = True, pair = False)
# 아니면 아래와 같은 형식으로 '=' 기호를 빼고 사용하셔도 결과는 같습니다.
data_train = BERTDataset(train_set_data, 0, 1, token, max_len, True, False)

from kobert.utils import get_tokenizer
import gluonnlp as nlp
tokenizer = get_tokenizer()


bertmodel, vocab = get_pytorch_kobert_model()
token = nlp.data.BERTSPTokenizer(tokenizer, vocab, lower=False)

data_train = BERTDataset(dataset = train_set_data, sent_idx = 0, label_idx = 1, bert_tokenizer = token, max_len = 64, pad = True, pair = False)

# data를 리스트[0]에는 문장을 넣고, 리스트[1]에는 감정(labels)으로 넣어서 리스트로 만들어줌.
train_set_data = [[i, str(j)] for i, j in zip(train_set['data'], train_set['label'])]


print(train_set_data[0])
# 출력 결과 => ['큰아들이 결혼하는데 집을 사달라고 해서 화가 나.', '4']

print(data_train[0])
# 출력 결과 =>
(array([   2, 4688, 6797, 5940,  950, 7795, 4384, 7088, 2573, 5794, 5439,
        5007, 5112, 5330, 1370,  517,   54,    3,    1,    1,    1,    1,
           1,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1,
           1,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1,
           1,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1,
           1,    1,    1,    1,    1,    1,    1,    1,    1], dtype=int32),
 array(18, dtype=int32),
 array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       dtype=int32),
 4)
# array의 첫 번째는 패딩된 시퀀스, 두 번째는 길이와 타입에 대한 내용, 세 번재는 어텐션 마스크 시퀀스이다.

def get_kobert_model(model_path, vocab_file, ctx="cpu"):
    bertmodel = BertModel.from_pretrained(model_path)
    device = torch.device(ctx)
    bertmodel.to(device)
    bertmodel.eval()
    vocab_b_obj = nlp.vocab.BERTVocab.from_sentencepiece(vocab_file,
                                                         padding_token='[PAD]')
    return bertmodel, vocab_b_obj