[Pytorch Lightning] BART를 훈련해 Text를 요약해보자 - 2

• 1편 : 프로젝트 요약 및 LightningModule 설계
• 2편 : Dataset 구축 및 DataLoader 구축
• 3편 : logger 작성, wandb 연동 및 확인
• 4편 : 모델 로딩 및 실행결과(wandb) 확인

이번에는 Dataset 구축 및 DataLoader 구축에 대해 설명하려고 한다. 근데 생각을 해보았을 땐, 이 글을 보는 사람들은 데이터셋 구축은 관심이 없을 것이고 아무래도 어떻게 코드를 짜는지 궁금할 것이라 생각하기 때문에 데이터셋 구축에 대해서는 설명하지 않고 데이터셋이 어떻게 구성되어 있는지만 말하려고 한다.

 

 

 

Dataset 구축

# Text : Encoder 
# function : Decoder

# 예시
# 음료 개수, 메뉴 개수, 메뉴 개수 -> drink(quantity=1); menu(quantity=3); menu(quantity=2);
{"Text": "아이스 카페 라떼 4잔 주세요 그 외에 따듯한 그린티 라떼 줄래?", "function": "ICE_cafe_latte(quantity=4); HOT_greentea_latte(quantity=1); "}

데이터셋의 구성은 다음과 같다.

• Encoder에는 한국어로 되어있는 Text
• Decoder는 메뉴에 대해 요약된 Text

결론은 카페 메뉴에 대한 text가 들어갔을 때, 영어로 요약된 text가 출력되도록 하는 것이다.

 

DataLoader 구축

라이브러리

from datasets import Dataset
from torch.utils.data import DataLoader
from tokenizers.processors import TemplateProcessing

라이브러리는 다음과 같이 사용합니다.

 

데이터셋 로딩

# json을 dataset으로 변환
dataset = Dataset.from_json(fname)

 

토크나이저 설정

# cls token과 sep token이 설정되어있지 않으면 설정
if not tokenizer.cls_token:
    tokenizer.cls_token = tokenizer.eos_token
if not tokenizer.sep_token:
    tokenizer.sep_token = tokenizer.eos_token

토크나이저에 클래스 시작(cls_token)과 분리(sep_token) 토큰이 설정되어 있지 않다면, 각각을 문장 종료(eos_token) 토큰으로 설정

 

템플릿 프로세싱(포스트 프로세서 설정)

# TemplateProcessing
tokenizer._tokenizer.post_processor = TemplateProcessing(
    single=f"{tokenizer.cls_token} $0 {tokenizer.sep_token}",
    pair=f"{tokenizer.cls_token} $A {tokenizer.sep_token} $B:1 {tokenizer.sep_token}:1",
    special_tokens=[(tokenizer.cls_token, tokenizer.cls_token_id), (tokenizer.sep_token, tokenizer.sep_token_id)],
)

다음을 하는 이유는 다음과 같다. BART와 같이 모델이 요구하는 입력 형식에 맞추기 위함이다.

예를 들면 다음과 같다.

• 입력 : 아이스 카페 라테 4잔 주세요
• 출력 : [CLS] 아이스 카페 라뗴 4잔 주세요 [SEP]
• 입력 : 아이스 카페 라떼라테 4잔 주세요 그 외에 따듯한 그린티 라테 줘
• 출력 : [CLS] 아이스 카페 라뗴 4잔 주세요 [SEP] 그 외에 따듯한 그린티 라테 줘 [SEP]

다음과 같은 형식으로 처리한다.

 

전처리 함수

# 전처리 함수
def preprocess_function(examples):
    processed = {}
    inp = f'{examples["Text"]}'

    # 입력에 대한 토큰화
    tokenizer_input = tokenizer(
        inp,
        padding="max_length",
        max_length=max_length,
        truncation=True
    )
    processed["input_ids"] = tokenizer_input["input_ids"]
    processed["attention_mask"] = tokenizer_input["attention_mask"]

    # 훈련모드인 경우
    if mode == "train":
        #print(examples["code"])

        # 출력에 대한 토큰화
        tokenizer_output = tokenizer(
            examples["function"], 
            padding="max_length", 
            max_length=max_length, 
            truncation=True
        )
        processed["decoder_input_ids"] = tokenizer_output["input_ids"]
        processed["decoder_attention_mask"] = tokenizer_output["attention_mask"]

    # 토큰화된 배열 반환
    return processed

요약을 하면, 입력데이터를 BART의 출력에 맞도록 tokenizing 하는 것이라고 생각하면 된다.

• train mode : encoder, decoder 모두 tokenizing
• not train mode : encoder만 tokenizing

그 이유는 테스트 모드의 경우에는 decoder는 예측을 하는 것이기 때문에 정답이 없다.

 

데이터셋 전처리 적용

# dataset에 대해 전처리 함수로 매핑, columns제거 및 torch tensor로 변환
dataset = dataset.map(preprocess_function,remove_columns=dataset.column_names).with_format("torch")

dataset에 대해 전처리를 적용한 뒤, 이전에 필요 없는 column을 삭제한  뒤, pytorch 형태로 변환한다.

 

데이터셋 분할

dataset = dataset.train_test_split(0.2)

train_dataloader = DataLoader(dataset['train'], shuffle=True, batch_size=train_batch_size, num_workers=8, pin_memory=True)
valid_dataloader = DataLoader(dataset['test'], shuffle=True, batch_size=valid_batch_size, num_workers=8, pin_memory=True)

return train_dataloader, valid_dataloader

train과 valid를 8 : 2로 분할하는 코드이다. 이렇게 되면 dataset은 'train'과 'test'로 데이터를 8 : 2로 나눈다. 

• shuffle : data를 무작위로 섞는다.
• batch_size : batch_size를 조절한다. ram이 적다면, 낮게설정하고 ram이 많다면 높게 설정하는 것이 좋다.
• num_workers : 데이터 로딩을 위해 사용할 별도의 작업 프로세스 수를 지정, 8이면, 8개의 프로세스가 병렬로 데이터를 로딩한다.
• pin_memory : 데이터를 더 빠르게 GPU로 전송할 수 있게 하여, 특히 GPU를 사용할 때 데이터 전송 병목을 줄여준다.

다음과 같이 분할하여 train_dataloader과 valid_dataloader 형식으로 반환한다.

 

전체 코드

import json
import pandas as pd

from datasets import Dataset
from torch.utils.data import DataLoader
from tokenizers.processors import TemplateProcessing

def OrderDataLoader(fname, tokenizer, train_batch_size,valid_batch_size, max_length,split=False, mode="train"):
    """
    Build Data Loader

    """

    # json읽어와서 dataset으로 변환
    dataset = Dataset.from_json(fname)
    
    # cls token과 sep token이 설정되어있지 않으면 설정
    if not tokenizer.cls_token:
        tokenizer.cls_token = tokenizer.eos_token
    if not tokenizer.sep_token:
        tokenizer.sep_token = tokenizer.eos_token

    # TemplateProcessing
    tokenizer._tokenizer.post_processor = TemplateProcessing(
        single=f"{tokenizer.cls_token} $0 {tokenizer.sep_token}",
        pair=f"{tokenizer.cls_token} $A {tokenizer.sep_token} $B:1 {tokenizer.sep_token}:1",
        special_tokens=[(tokenizer.cls_token, tokenizer.cls_token_id), (tokenizer.sep_token, tokenizer.sep_token_id)],
    )

    # 전처리 함수
    def preprocess_function(examples):
        processed = {}
        inp = f'{examples["Text"]}'
       
        # 입력에 대한 토큰화
        tokenizer_input = tokenizer(
            inp,
            padding="max_length",
            max_length=max_length,
            truncation=True
        )
        processed["input_ids"] = tokenizer_input["input_ids"]
        processed["attention_mask"] = tokenizer_input["attention_mask"]
        
        # 훈련모드인 경우
        if mode == "train":
            #print(examples["code"])
             
            # 출력에 대한 토큰화
            tokenizer_output = tokenizer(
                examples["function"], 
                padding="max_length", 
                max_length=max_length, 
                truncation=True
            )
            processed["decoder_input_ids"] = tokenizer_output["input_ids"]
            processed["decoder_attention_mask"] = tokenizer_output["attention_mask"]
        
        # 토큰화된 배열 반환
        return processed
    

    # dataset에 대해 전처리 함수로 매핑, columns제거 및 torch tensor로 변환
    dataset = dataset.map(preprocess_function,remove_columns=dataset.column_names).with_format("torch")

    # dataset을 dataloader로 변환

    dataset = dataset.train_test_split(0.2)

    train_dataloader = DataLoader(dataset['train'], shuffle=True, batch_size=train_batch_size, num_workers=8, pin_memory=True)
    valid_dataloader = DataLoader(dataset['test'], shuffle=True, batch_size=valid_batch_size, num_workers=8, pin_memory=True)

    return train_dataloader, valid_dataloader