Pytorch CustomDataset 선언

강태원·2023년 12월 3일
1

학습

사투리를 표준어로 바꾸기 위해서 사용해본 모델은 KoGPT와 KoBART 모델이다.

1. KoGPT

GPT 모델은 문장 생성 영역에서 주로 쓰인다.
그럼에도 기계 번역에 쓰인다는 글을 본 적이 있었다.

따라서 이전에 KoGPT 모델을 학습 시켜본 코드가 있었기 때문에 먼저 학습을 돌려봤지만,,
성능이 좀.. 구렸다.

이 때 학습했던 코드와 결과물이 유실되어 따로 첨부하지 못하는 점이 많이 아쉽다.

코드를 체계적으로 관리하자..

2. KoBART

실제로 이번 프로젝트에 쓰였던 모델이다!

BART 모델에 대한 논문 리뷰는 이 글에서 자세히 설명되어있다.
https://chloelab.tistory.com/34

쉽게 얘기하자면 "니 밥 뭇나?"를 입력값이라 생각했을 때, BART 모델은 여기에 노이즈를 추가해서 문장을 오염시킨 뒤에, "너 밥 먹었어?"로 다시 복원하는 작업을 통해 학습을 하게 된다.

https://github.com/SKT-AI/KoBART 이 레포지토리에 KoBart가 pretrained되어 있긴하지만
https://huggingface.co/heegyu/kobart-text-style-transfer 에서 내가 원하는 형식과 비슷하게 학습되어 있는 모델을 발견했다.

따라서 이 모델의 가중치를 받아온 뒤에 학습을 하는 게 좀 더 높은 성능이 나올 것이라고 판단했다.

학습 코드

Pytorch Custom Dataset

transformers 라이브러리를 설치해주자

pip install transformers
# dataset.py
import pandas as pd
from torch.utils.data import Dataset
class TextStyleTransferDataset(Dataset):
    def __init__(self, df,tokenizer):
      self.df = df
      self.tokenizer = tokenizer
    
    def __len__(self):
        return len(self.df)

    def __getitem__(self, index):
        row=self.df.iloc[index]
        text1=row[0]
        text2=row[1]
        target_style_name = '표준어'
		
        # Tokenizer를 허깅페이스의 레포지토리에서 가져왔기 때문에
        # 테스트 코드를 참고해서 모델 입력 형식을 만들었다.
        encoder_text = f"{target_style_name} 말투로 변환:{text1}"
        decoder_text = f"{text2}{self.tokenizer.eos_token}"
        model_inputs = self.tokenizer(encoder_text, max_length=64, truncation=True)

        with self.tokenizer.as_target_tokenizer():
          labels = tokenizer(decoder_text, max_length=64, truncation=True)
        model_inputs['labels'] = labels['input_ids']
        del model_inputs['token_type_ids']

        return model_inputs
# 지난번에 만들어뒀던 데이터셋을 불러와서 데이터프레임을 만들어준다
def make_df(data_root):
    df = pd.read_csv(f'{data_root}/data.tsv',sep='\t')
    # 주로 쓰이는 방법 같지는 않지만,, Train Data와 Test Data를 8:2 비율로 나눠준다
    rate=int(len(df)*0.2)
    df_train,df_test = df[rate:],df[:rate]

    print(f'Train DataFrame length : {len(df_train)},Test DataFrame length : {len(df_test)}')
    return df_train,df_test
# 위 함수와 클래스를 종합해서 Dataset들을 반환해준다
def make_dataset(data_root,tokenizer):
    df_train,df_test = make_df(data_root)

    train_dataset = TextStyleTransferDataset(df_train,tokenizer)
    test_dataset = TextStyleTransferDataset(df_test,tokenizer)

    return train_dataset,test_dataset

다음 포스트에는 이 스크립트에서 정의한 TextStyleTransferDataset으로 데이터셋을 생성하고 Train Arguments를 정의한 뒤 학습하는 코드를 알아보도록 하자

profile
가치를 창출하는 개발자! 가 목표입니다

0개의 댓글