[PyTorch, KoGPT2] Fine-tuning하고 문장 생성하기(w/ full code)

GPT2는 주어진 텍스트의 다음 단어를 잘 예측할 수 있도록 학습된 언어모델이며 문장 생성에 최적화 되어 있습니다. KoGPT2는 부족한 한국어 성능을 극복하기 위해 40GB 이상의 텍스트로 학습된 한국어 디코더 언어모델입니다.
출처 : https://github.com/SKT-AI/KoGPT2

이렇게 기학습된 kogpt2 모델을 가지고 원하는 텍스트를 생성해봅시다.
지금부터 소개해드릴 예시는 kogpt2가 신문 기사의 제목을 생성해낼 수 있는 모델을 파인튜닝으로 만들어 보는 것입니다.

사용한 데이터의 예시는 아래와 같습니다.

예시 데이터는 제가 직접 파인튜닝한 모델로 생성해낸 문장들입니다.
정확하고 정교한 모델을 만들기 위해서는 예시로 올려둔 데이터가 아닌 정제된 데이터로 사용하시길 권장합니다.

디렉토리 구조는 아래와 같습니다.

📦pytorch-kogpt2-example
┣ 📂data
 ┃ ┗ 📜train.csv
 ┣ 📜dataloader.py
 ┣ 📜generate.py
 ┣ 📜train.py
 ┗ 📜utils.py

DataLoader

import pandas as pd
from torch.utils.data import Dataset, DataLoader


class GPTDataset(Dataset):
    def __init__(self, tokenizer, file_path):
        data = pd.read_csv(file_path)
        concats = [
            label + "|" + text for label, text in zip(data["target"], data["text"])
        ]
        self.item = tokenizer(
            concats,
            return_tensors="pt",
            padding="max_length",
            truncation=True,
            max_length=32,
        )["input_ids"]
        self.length = len(concats)

    def __getitem__(self, i):
        return self.item[i]

    def __len__(self):
        return self.length


def GPTDataLoader(tokenizer, file_path, batch_size):
    data = GPTDataset(tokenizer, file_path)
    return DataLoader(data, batch_size=batch_size)

기사의 분류를 입력으로 넣으면 기사 제목을 출력으로 내어주는 모델을 만들 계획이기 때문에, 학습 시에는 기사의 분류와 기사의 제목을 특수 기호로 합쳐 입력으로 넣어줍니다.

입력 데이터 예시

"스포츠 | 프로야구 삼성의 첫 대결...SK 두산 꺾고 SO리그 4연패 탈출",
"사회 | 박영수 특검 삼성그룹·글루코어 결함설 공식 부인",
"생활문화 | 신간 괜찮은 일터·착한 작가"
...

Fine-tuning

import os
import argparse
import torch
from tqdm import tqdm
from transformers import (
    GPT2LMHeadModel,
    AutoTokenizer,
    AdamW,
    get_linear_schedule_with_warmup,
)
from dataloader import GPTDataLoader
from utils import generate


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_name", default="skt/kogpt2-base-v2", type=str)
    parser.add_argument("--data_dir", default="./data", type=str)
    parser.add_argument("--batch_size", default=32, type=int)
    parser.add_argument("--epochs", default=10, type=int)
    parser.add_argument("--lr", default=2e-5, type=float)
    parser.add_argument("--warmup_steps", default=200, type=int)
    args = parser.parse_args()

    BASE_DIR = os.getcwd()
    DATA_DIR = os.path.join(BASE_DIR, args.data_dir)

    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
    tokenizer.add_special_tokens({"pad_token": "<pad>"})

    # Load dataset
    train_dataloader = GPTDataLoader(
        tokenizer, os.path.join(DATA_DIR, "train.csv"), args.batch_size
    )

    # Load model
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = GPT2LMHeadModel.from_pretrained(args.model_name).to(device)
    model.train()

    # Set optimizer, scheduler
    optimizer = AdamW(model.parameters(), lr=args.lr)
    scheduler = get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=-1
    )

    min_loss = int(1e-9)
    for epoch in range(args.epochs):
        print(f"Training epoch {epoch}")
        for input_text in tqdm(train_dataloader):
            input_tensor = input_text.to(device)
            outputs = model(input_tensor, labels=input_tensor)
            loss = outputs[0]

            optimizer.zero_grad()
            model.zero_grad()
            loss.backward()
            optimizer.step()
            scheduler.step()

        print(f"epoch {epoch} loss {outputs[0].item():0.2f}")

        # Generate examples on epoch end
        labels = ["IT과학", "경제", "사회", "생활문화", "세계", "스포츠", "정치"]
        for label in labels:
            gen = generate(label, tokenizer, model, 1)
            print(f"{label} : {gen[0]}")

        # Save best model
        if outputs[0].item() < min_loss:
            min_loss = outputs[0].item()
            model.save_pretrained("./best_model")

    print("Training Done!")

GPT는 self-supervised learning을 기반으로 학습을 진행하기 때문에, 따로 라벨을 만들 필요 없이 텍스트를 입력해주는 것만으로도 fine tuning이 가능합니다.

예시 코드를 보면 매 에폭이 끝날 때마다 생성 예시들을 보여주어서 매 에폭마다 생성 성능을 정성적으로 평가할 수 있게끔 장치해두었습니다.

Generate

import torch
from tqdm import tqdm


def generate(input_text, tokenizer, model, num):
    sentence_list = []
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    token_ids = tokenizer(input_text + "|", return_tensors="pt")["input_ids"].to(device)
    for cnt in tqdm(range(num)):
        gen_ids = model.generate(
            token_ids,
            max_length=32,
            repetition_penalty=2.0,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
            bos_token_id=tokenizer.bos_token_id,
            use_cache=True,
            do_sample=True,
        )
        sentence = tokenizer.decode(gen_ids[0])
        sentence = sentence[sentence.index("|") + 1 :]
        if "<pad>" in sentence:
            sentence = sentence[: sentence.index("<pad>")].rstrip()
        sentence = sentence.replace("<unk>", " ").split("\n")[0]

        if cnt % 100 == 0 and cnt != 0:
            print(sentence)
        sentence_list.append(sentence)
    return sentence_list

문장을 생성하는 generate 원본 함수에는 다양한 옵션들이 있습니다.
공식 링크를 참조해주세요.

소스코드에는 문장을 생성한 후 저장하는 코드 또한 추가해두었습니다.

소스코드 링크 : https://github.com/JLake310/pytorch-kogpt2-example