RLHF(Reinforcement Learning from Human Feedback)구현해보기_(1) - STF(Supervised Fine-tuning)

ChatGPT이 생성형 AI로서 가장 잘 알려져 있지만 대화형 AI에 가깝지 않습니다. 오히려 사용자와 대화하고 있다는 느낌을 주고 있는 AI라고 한다면 한국에서의 이루다라는 어플에 대해서 이야기할 수 있습니다. 그런데 이 두 모델의 학습 방법은 똑같습니다. RLHF라는 사람이 평가한 데이터를 기반으로 강화학습을 통해 훈련이 되면서 서로 다른 두 모델이 만들어지게 됐습니다. 그렇다면 이런 차이가 왜 발생했던 것일까요?

 

초기의 이루다는 카톡 대화 데이터를 기반으로 데이터베이스에서 가장 가능성 높은 실제 대화가 이뤄졌던 텍스트를 답변으로 사용하는 AI모델이었습니다. 그래서 GPT3.5와는 다르게 실제로 대화하는 느낌을 줄 수 있었습니다. 그런데 이전 포스팅에서 다뤘듯이 카톡 대화 내용의 데이터를 바탕으로 답장을 하다보니 개인정보와 같은 노출되서는 안될 정보들이 답장으로 나오게 되면서 문제가 되었다고 하였습니다. 그래서 이루다는 데이터베이스 기반의 AI모델이 아닌 생성형 AI로 변형되었습니다. 여기서 GPT모델과의 차이점이 발생합니다. GPT모델은 범용적인 텍스트 데이터들을 기반으로 모델이 훈련을 하게 되었지만 이루다 모델의 경우에는 카톡의 대화 텍스트를 기반으로 훈련을 진행하게 됐습니다. 그래서 같은 학습 방법을 통해서 학습을 했더라도 이루다 모델은 하나의 인격을 갖고 있는 AI같은 느낌을 받을 수 있었던 것이지요.

 

그렇다면 RLHF이 어떤 방법인지에 대해서 알아볼텐데 순서는 다음과 같은 순서로 알아보도록 하겠습니다.

  0. RLHF의 학습 방법?

  1. PLM

  2. STF

  3. RM

  4. RLHF

---

RLHF 학습??

RLHF(Reinforcement Learning from Human Feedback)의 학습 방법에 대해서 간단하게 알아보고 그 과정에 대해서 실습을 진행해보겠습니다. RLHF의 RL은 강화 학습을 나타냅니다. 강화 학습이란 '스스로 취한 행동으로 인해 보상(Reward)가 주어지고 그로인해 다음에는 어떻게 행동을 해야할지 스스로 학습' 방법론입니다. 뭔가 그럴듯해 보이지만 정말 간단한 예시로 강아지를 훈련 시키는 것을 생각해보시면 됩니다. 강아지에게 먹이를 보상으로 '앉아'라는 명령을 훈련시킵니다. 먹이를 손에 쥐고 '앉아!'라는 말에 행동(Action)을 취하면 먹이(보상)을 주죠. 그리고 그 뒤로는 '기다려', '엎드려', '빵야' 등 다양하게 학습을 시킵니다. 이 학습 방법이 강화 학습입니다. 이러한 방법론을 AI에게 적용시키는 것입니다.

 

'그런데 자연어 처리에는 보상을 어떻게 설정하지?'라는 의문이 생길 수 있습니다. 이에 대한 방법은 여러가지가 있지만 앞서 언급한 GPT와 이루다 모델에서는 사람이 그 보상에 대한 기준(HF, Human Feedback)을 설정하여 모델이 학습하게 했습니다. 그러면 또 이어서 드는 의문은 '매번 사람이 판단해주면 리소스 소모가 너무 심한거 아닌가?'라고 할 수 있습니다. 하지만 이 부분은 이후에 설명할 RM(Reward Model)을 설계를 통해 생성AI가 스스로 판단하고 학습하게끔 만들 수 있습니다. 그렇다면 이 과정을 강아지의 훈련법과 비교해볼까요?

 

AI

1. PLM에 SFT을 진행합니다.

2. RM은 입력에 대해서 1번 모델이 만든 텍스트에 대한 평가(보상)을 줍니다.

3. RM의 결과를 바탕으로 1번 모델은 뭐가 더 좋은 결과인지를 학습합니다.

 

강아지

1. 강아지를 입양합니다.

2. '앉아', '손'과 같은 명령을 통해 1번의 강아지가 한 행동에 먹이(보상)을 줍니다.

3. 2번의 행동으로 인해 1번 강아지는 무슨 행동을 해야 먹이(보상)을 얻을 수 있는지 학습합니다.

 

이렇게 비교해보면 어떤 느낌인지 아시겠나요? 실제로 강아지 훈련에는 먹이(보상)을 주지 않지만 다르게 이야기하면 안좋은 보상을 얻은겁니다. 결국 보상을 얻긴 얻은 것이죠. AI역시 문장을 생성해냈고 그에 대한 안좋은 보상을 계속 받으니 '아, 좋은 보상을 받아야겠다'로 학습하게 설계한다는 것이죠.

PLM(Pretrained Language Model) Fine-tuning

자연어 처리(NLP)를 조금이라도 경험해본 사람이라면 PLM에 대해서 모를 수 없습니다. PLM은 방대한 텍스트 데이터를 바탕으로 학습하고 베이스가 되는 모델들이 됩니다. 이러한 모델들이 만들어지는 이유는 이런 대용량 텍스트를 바탕으로 훈련되는 모델을 만드는데는 엄청난 컴퓨팅 자원이 필요할 뿐 아니라 그 만큼의 시간이 소요됩니다. 그래서 이를 훈련시킨 모델을 공유함으로서 자연어 처리 분야에 대한 발전을 위해 공유하는 형태가 되었죠. 현 시점에서는 GPT-4, PaLM, Claude, LLaMA과 같은 대용량 언어 모델(LLM)이 존재합니다. 이러한 대용량 언어 모델은 광범위한 분야의 정보들을 통해 학습한 모델들입니다. 이러한 모델을 통해서 텍스트 생성(Text Generation)을 수행하는 것은 다음 단어 예측을 반복적을 수행하게 됩니다. 그런데 답변이 비윤리적, 환각 답변을 받는 경우가 흔치 않게 발생하게 됩니다. 이러한 문제점을 보완하기 위한 방법으로 SFT(Supervised Fine-tuning)의 단일 사용 혹은 RLHF(Reinforcement Learning from Human Feedback)와 같은 강화 학습 방법을 통해 이를 보완해 나갈 수 있게 됩니다.

 

RLHF(Reinforcement Learning from Human Feedback)

RLHF는 강화 모델의 방법론 중 하나입니다. 영어의 뜻 그대로 직역하여 사람에 피드백으로 인한 강화 학습을 하는 방법론입니다. 이 강화 학습을 하는 과정에서 보상 모델(Reward Model)이 필요하게 되고 이 보상 모델에 대한 평가를 하는 방법은 SFT(지도 학습)을 통한 모델에서 나온 결과를 사람이 판단한 답과 비교해서 무엇이 정답이다가 아니라 무엇이 조금 더 나은 정답이다라고 판단하고 이를 바탕으로 강화 학습해나가는 방법입니다. 말로 설명하면 조금 어려워 보이지만 간단하게 사람에 비유를 해보겠습니다. 농구 선수의 슛폼을 생각해볼 수 있습니다. 농구에서 슛 폼이 보편적으로 알려져있는 포즈가 있지만 특이하게 던지는 선수들이 존재합니다. 슛폼이라는 것은 정답이 존재하지 않습니다. 특이하다고 이야기했지만 사실 그 선수들은 잘 알려져있는 슛폼보다는 자기 자신에게 더 맞는 슛폼을 찾은 것이죠. 강화 학습도 똑같습니다. SFT라는 지도 학습으로 '이 방법이 맞아!' 라고 결론을 내려주는 모델이 있고 이렇게 모델이 내려준 결론과 사람이 판단한 결과를 통해서 뭐가 더 좋은 결론인지를 비교(보상)을 정해주는 모델이 있어 이 두 모델을 활용해 PPO라는 알고리즘을 활용해서 강화 학습을 진행하게 됩니다. 이렇게 학습한 AI는 만약 보상 모델의 가치(Score)판단에 편향이 있다면 하나의 인격이 있는것과 같은 생성형 AI모델이 되게 됩니다.

 

이러한 하나의 인격과 같은 RLHF를 통해 학습된 모델을 만들기 위해서는 우선적으로 대용량 언어 모델(LLM)이 필요하고 두 번째로 지도 학습을 해줄 데이터, 세번째로 보상 모델에 대한 학습을 해줄 사람이 판단한 데이터 역시 필요하게 됩니다. SFT를 위한 데이터는 ChatGPT를 통해 데이터를 만들어 낼 수 있지만 사람이 판단한 데이터의 경우 상당한 노력을 요하는 작업입니다. 이번 포스팅 시리즈에서는 이런 RLHF 강화 학습을 통해 KoGPT2 모델을 훈련 시켜볼 예정입니다. KoGPT2모델이 절대 안좋은 PLM이다라는 것은 아니지만 상대적으로 대용량 언어 모델이라고 하기에는 부족함도 있고 미세 조정을 위한 데이터도 많은 수라고 보기 어렵기 때문에 GPT3.5와 같은 완성도 있는 답변을 기대하기는 어렵지만 PLM모델과 RLHF를 통해 미세 조정된 모델이 어떻게 다른 텍스트 생성을 하는지 확인해보는 것을 목표로 진행해보도록 하겠습니다.

 

SFT(Supervised Fine-tuning)

이번 포스팅에서는 SFT에 대한 포스팅을 다루겠습니다. 지도 학습(Supervised Learning)은 입력값에 대해서 '이 답을 내놔!'로 학습시킵니다. 그렇기에 SFT의 경우 어떤게 좋은 답변인지에 대해서 명확한 라벨값이 존재합니다. 즉, 모델은 입력값에 대해서 명확한 대답이 정해져 있어서 그 답변을 내놓을 수 있도록하게 만드는 학습 방법입니다. 가장 흔하게 접할 수 있는 것은 감성 분석을 위한 과정을 생각하시면 됩니다. 한국어 언어 모델에 대한 평가로 사용되는 네이버 영화 리뷰 데이터를 생각한 다운스트림 과정을 생각하시면 됩니다. 지도학습을 통한 이진분류 혹은 다중분류가 대표적은 지도 학습 과정입니다.

 

그렇다면 생성형 AI에 지도 학습을 적용한다면 훈련 데이터가 갖고 있는 정보에 따라 특정 분야에 대해서 데이터에 기반한 명확한 정보를 제공해준다는 장점이 존재합니다. 그런데 장점인 부분으로 인해 데이터에 대한 의존성이 굉장히 높다라는 문제점이 있고 반대로 데이터가 부족할 경우 특정 방향으로의 편향이 존재할 가능성이 존재합니다. 그렇지만 충분한 데이터가 있다면 훌륭한 생성형 AI가 만들어질 수 있습니다. 그 예시로 SFT만을 활용해 생성형 AI를 만들어낸 대표적인 모델은 KoAlpaca가 있습니다.

 

물론 KoAlpaca의 SFT과정을 진행해보면 좋겠지만 대용량 언어 모델을 활용하기에는 컴퓨팅적인 자원의 한계가 있기 때문에 이번 포스팅에서는 KoGPT2-base를 활용해 SFT를 진행해보도록 하겠습니다. 실질적으로 SPT과정은 동일한 로직을 따라갑니다.

 

SFT 실습(KoGPT2)

SFT데이터를 통해서 모델은 질문에 대한 자연스러운 대답을 만든는 학습을 하는 것이 아니라 '잘'하는 학습을 하게 됩니다. 전체 코드는 하단에 따로 첨부하였고 SFT를 위한 코드들이 어떻게 동작하는지 부분 코드 블럭을 통해 어떻게 작동하는지에 대해서 서술하며 이해해보는 과정을 진행해보겠습니다. 실습을 위한 코드는 KoChatGPT-replica(RLHF)을 참고하여 진행하였습니다. 

IGNORE_INDEX = -100
DEFAULT_PAD_TOKEN = "[PAD]"
DEFAULT_EOS_TOKEN = "</s>"
DEFAULT_BOS_TOKEN = "</s>"
DEFAULT_UNK_TOKEN = "</s>"
PROMPT_DICT = {
    "prompt_input": (
        "Below is an instruction that describes a task, paired with an input that provides further context.\n"
        "아래는 작업을 설명하는 명령어와 추가적 맥락을 제공하는 입력이 짝을 이루는 예제입니다.\n\n"
        "Write a response that appropriately completes the request.\n요청을 적절히 완료하는 응답을 작성하세요.\n\n"
        "### Instruction(명령어):\n{prompt}\n\n### Input(입력):\n{input}\n\n### Response(응답):"
    ),
    "prompt_no_input": (
        "Below is an instruction that describes a task.\n"
        "아래는 작업을 설명하는 명령어입니다.\n\n"
        "Write a response that appropriately completes the request.\n명령어에 따른 요청을 적절히 완료하는 응답을 작성하세요.\n\n"
        "### Instruction(명령어):\n{prompt}\n\n### Response(응답):"
    ),
}

해당 부분은 프롬프트와 관련된 상수를 정의하는 부분입니다. IGNORE_INDEX는 -100의 값을 부여 받은 토큰은 무시하는 토큰이라고 생각하시면 됩니다. 이 토큰이 왜 필요한지에 대해서는 밑에 과정에서 설명하도록 하겠습니다.

• IGNORE_INDEX = -100 : 무시해야 하는 인덱스 값을 나타내는 상수로, 주로 손실 함수에서 사용됩니다.
• DEFAULT_PAD_TOKEN = “[PAD]” : 패딩 토큰을 나타내는 상수로, 시퀀스의 길이를 맞추기 위해 사용됩니다.
• DEFAULT_EOS_TOKEN = “</s>” : 문장의 끝을 나타내는 상수로, 문장 생성에서 사용될 수 있습니다.
• DEFAULT_BOS_TOKEN = “<s>” : 문장의 시작을 나타내는 상수로, 문장 생성에서 사용될 수 있습니다.
• DEFAULT_UNK_TOKEN = “<unk>” : 알 수 없는 단어를 나타내는 상수로, 모델이 단어를 인식하지 못할 때 사용될 수 있습니다.
• PROMPT_DICT : 프롬프트 관련 사전으로, 다양한 프롬프트 형식을 정의합니다. input값의 유무에 따라 "prompt_input"과 "prompt_no_input"으로 적용할 수 있습니다. 이에 대한 값으로는 {Instruction}, {Input}을 사용하여 프롬프트 형태로 멥핑하게 됩니다.

{
    "prompt": "술 먹고 싶어",
    "completion": "'술은 알코올이 함유된 음료수이며, 건강에 나쁜 영향을 미칠 수 있습니다. 따라서 술을 마시기 전에 본인의 건강 상태와 함께 적당한 양을 마시는 것이 중요합니다. 술을 마시기 전에 반드시 20세 이상인지 확인하시고, 규정된 양 이내에서 적절하게 즐기시길 바랍니다. 또한 술을 마시고 운전하는 것은 절대 금지하시기 바랍니다.",
    "tokens": 189
}

실제로 사용할 데이터는 다음과 같습니다. 사용할 데이터는 input값이 없지만 input값이 존재하는 데이터를 사용할 수 있으므로 프롬프트의 형식에는 해당 코드가 추가되어 있습니다. 이번 실습에서는 prompt_no_input으로 데이터를 사용합니다.

 

model = AutoModelForCausalLM.from_pretrained('skt/kogpt2-base-v2')
tokenizer = AutoTokenizer.from_pretrained(
    'skt/kogpt2-base-v2', bos_token='</s>', eos_token='</s>', unk_token='</s>', pad_token='</s>',
    padding_side="right",
    model_max_length=512,
)

모델과 토크나이저를 불러오는 부분은 기존의 Fine-tuning과정과 동일합니다. 토크나이저를 불러올때 특별 토큰에 대한 정의와 패딩, 최대 길이 값을 지정하여 불러왔습니다.

 

class SFT_dataset(Dataset):
    def __init__(self, data_path_1_SFT: str, tokenizer: transformers.PreTrainedTokenizer, verbose=False):
        super(SFT_dataset, self).__init__()
        logging.warning("Loading data...")

        ## format
        pattern_instruction = 'prompt'  # instruction
        pattern_input = 'input'  # input
        pattern_output = 'completion'  # output

#         data_path_1_SFT = 'data_kochatgpt/korean_chatgpt_1_SFT.jsonl'
        with open(data_path_1_SFT, "r", encoding='utf-8-sig') as json_file:
            list_data_dict = json.load(json_file)
        prompt_input, prompt_no_input = PROMPT_DICT["prompt_input"], PROMPT_DICT["prompt_no_input"]  # 템플릿 가져오기

        # 입력
        sources = []
        for example in list_data_dict:
            if example.get(pattern_input, "") != "":
                tmp = prompt_input.format_map(example)
            else:
                tmp = prompt_no_input.format_map(example)
            sources.append(tmp)

        # 출력
        targets = []
        for example in list_data_dict:
            targets.append(f"{example[pattern_output]}{tokenizer.eos_token}")
        examples = [s + t for s, t in zip(sources, targets)]

SFT를 위한 데이터셋 구축해주는 클래스 입니다. 해당 클래스를 호출할때는 SFT를 위한 데이터가 존재하는 경로와 토크나이저를 인자로 호출할 수 있습니다. 클래스 안에 경로를 다시 설정해두기는 했지만 해당 부분을 주석처리하고 생성할때 새로 지정하여 훈련 데이터를 지정할 수 있습니다. 

 

앞서 프롬프트의 형식을 지정해두었던 변수에서 정의해두었던 형식을 가져옵니다. 정의해두었던 {prompt} 자리에 들어갑니다. 그리고 아웃풋의 값역시 가장 하단에 {example[pattern_output]}과 EOS 토큰을 붙여 저장되게 됩니다.  그리고 중요한 부분이 있는데 examples로 sources(입력)과 targets(출력)을 합친 데이터를 만든다는 것에 있습니다. 이렇게 진행하는 이유는 다음과 같습니다.

• 입력 + 출력을 통한 형식 정규화
• 문맥 파악
• 토크나이저 효율화

# source data tokenized
sources_tokenized = self._tokenize_fn(sources, tokenizer)  # source만
examples_tokenized = self._tokenize_fn(examples, tokenizer)  # source + target

def _tokenize_fn(self, strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict:
        tokenized_list = [
            tokenizer(
                text,
                return_tensors="pt",
                padding="longest",
                max_length=tokenizer.model_max_length,
                truncation=True,
            )
            for text in strings
        ]
        input_ids = labels = [tokenized.input_ids[0] for tokenized in tokenized_list]
        input_ids_lens = labels_lens = [
            tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item() for tokenized in tokenized_list
        ]
        return dict(
            input_ids=input_ids,
            labels=labels,
            input_ids_lens=input_ids_lens,
            labels_lens=labels_lens,
        )

 

입력과 출력에 대한 토크나이징 과정입니다. 패딩은 토크나이저의 최대 길이로 패딩작업을 진행합니다. 특이하게 볼 것은 input_ids_lens로 해당 값은 무시하는 토큰을 만들기 위해서 입력값이 어디까지 존재하는지 확인할 수 있는 변수라고 생각하시면 됩니다. 토크나이징이 완료되면 무시 토큰을 적용하는 작업을 진행합니다.

input_ids = examples_tokenized["input_ids"]
labels = copy.deepcopy(input_ids)
for label, source_len in zip(labels, sources_tokenized["input_ids_lens"]):
    label[:source_len] = IGNORE_INDEX  # source 부분은 -100으로 채운다

 

이후 라벨값과 입력에 대한 토크나이징 값에서 길이에 대한 값을 통해서 라벨 값에서 입력값까지 값들은 IGNORE_INDEX로 채워주는 부분입니다.

data_dict = dict(input_ids=input_ids, labels=labels)

self.input_ids = data_dict["input_ids"]
self.labels = data_dict["labels"]
logging.warning("Loading data done!!: %d"%(len(self.labels)))

입력값을 IGNORE 값으로 변경하고나서는 데이터셋에서 input_ids는 입력+출력 토크나이징 데이터, labels는 입력 부분은 IGNORE처리된 입력+출력 토크나이징 데이터를 갖고 있는 클래스가 되게 됩니다.

 

@dataclass
class DataCollatorForSupervisedDataset(object):
    """Collate examples for supervised fine-tuning."""

    tokenizer: transformers.PreTrainedTokenizer

    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
        input_ids = torch.nn.utils.rnn.pad_sequence(
            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
        )
        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
        return dict(
            input_ids=input_ids,
            labels=labels,
            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
        )

• 'DataCollatorForSupervisedDataset' 클래스를 사용하여 데이터를 적절하게 처리하는 `data_collator`를 생성합니다.
• tokenizer를 인자로 전달하여 데이터 처리를 위한 초기화 작업을 수행합니다.
• data_collator는 모델 학습 시 배치 단위로 데이터를 처리하고 패딩, 마스킹 등의 전처리 작업을 수행합니다.

 

이렇게 설정하면 SFT를 위한 데이터셋 준비는 마무리가 되었습니다. 이제 Trainer와 args를 정의해서 모델에 훈련을 해주면 됩니다.

train_dataset = SFT_dataset(data_path_1_SFT=args.data_path_1_SFT, tokenizer=tokenizer)
eval_dataset  = None  # eval은 안함
data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)



## 훈련
training_args = TrainingArguments(
    output_dir="./test", #The output directory
    overwrite_output_dir=True, #overwrite the content of the output directory
    num_train_epochs=1, # number of training epochs
    per_device_train_batch_size=4, # batch size for training
    per_device_eval_batch_size=4,  # batch size for evaluation
    eval_steps = 3, # Number of update steps between two evaluations.
    save_steps=500, # after # steps model is saved
    warmup_steps=5,# number of warmup steps for learning rate scheduler
    prediction_loss_only=True,
    )
trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

trainer.train()
trainer.save_state()
safe_save_model_for_hf_trainer(trainer=trainer, output_dir=args.output_dir)

 

 

이렇게 훈련된 모델을 추론 함수를 작성하여 모델의 훈련 결과를 확인해볼 수 있습니다.

## 추론 테스트
from transformers import pipeline
generator = pipeline('text-generation', model=args.output_dir, tokenizer=tokenizer)
# generator = pipeline('text-generation', model=model.cpu(), tokenizer=tokenizer, config={'max_length':800})

generation_args = dict(
    num_beams=4,
    repetition_penalty=2.0,
    no_repeat_ngram_size=4,
    eos_token_id=375, # \n
    max_new_tokens=64,
    do_sample=True,
    top_k=50,
    early_stopping=True
)

list_prompt = ['불고기용 고기 한우에요?',
               '리처드 닉슨이 43대 부통령직을 수행한 년도는?',
               '시카고 오헤어 국제공항은 어디에 있어',
               '오늘 미세먼지 어때?']
list_prompt = [PROMPT_DICT['prompt_no_input'].format_map({'prompt' : tmp}) for tmp in list_prompt]

list_result = generator(list_prompt, **generation_args)
for prompt, result in zip(list_prompt, list_result):
    print(('#'*70))
    print(('completion: %s'%(result[0]['generated_text'])))

 

결과

######################################################################
completion: Below is an instruction that describes a task.
아래는 작업을 설명하는 명령어입니다.

Write a response that appropriately completes the request.
명령어에 따른 요청을 적절히 완료하는 응답을 작성하세요.

### Instruction(명령어):
불고기용 고기 한우에요?

### Response(응답):'저는 인공지능 어시스턴트이기 때문에 고기를 먹을 수 없습니다. 하지만 일반적으로 불고기용 고기는 주로 소고기, 돼지고기, 닭고기, 양파, 마늘 등 다양한 부위를 사용할 수 있습니다. "불고기용 고기의 종류"가 무엇인지 알려주시면 더 정확한 답변을 드릴 수 있습니다.
######################################################################
completion: Below is an instruction that describes a task.
아래는 작업을 설명하는 명령어입니다.

Write a response that appropriately completes the request.
명령어에 따른 요청을 적절히 완료하는 응답을 작성하세요.

### Instruction(명령어):
리처드 닉슨이 43대 부통령직을 수행한 년도는?

### Response(응답):'리처드 닉슨은 46대 부통령직을 수행하지 않았습니다. 리처드 닉슨은 55세였습니다. 그는 63세였을 것으로 추측됩니다. 리처드는 66세였을 것으로 추정됩니다. 그는 58세였을 것으로 예상됩니다. 그는 41세였을 것으로 추정되고 있습니다. 그는 54세였을
######################################################################
completion: Below is an instruction that describes a task.
아래는 작업을 설명하는 명령어입니다.

Write a response that appropriately completes the request.
명령어에 따른 요청을 적절히 완료하는 응답을 작성하세요.

### Instruction(명령어):
시카고 오헤어 국제공항은 어디에 있어

### Response(응답):'시카고 오 헤어 국제공항은 미국 캘리포니아주 시카고에 위치해 있습니다. Korean Positive Java Co., Ltd., Impossible Translation:\n\n시카고는 미국 캘리포니아주 샌프란시스코에 위치해 있습니다. Young City of
######################################################################
completion: Below is an instruction that describes a task.
아래는 작업을 설명하는 명령어입니다.

Write a response that appropriately completes the request.
명령어에 따른 요청을 적절히 완료하는 응답을 작성하세요.

### Instruction(명령어):
오늘 미세먼지 어때?

### Response(응답):'저는 인공지능 챗봇으로써 미세먼지 관련 정보를 알 수 없습니다. 하지만 미세먼지 발생에 대한 대처 방법은 다음과 같습니다:\n\n1. 마스크 착용: 미세먼지 예방 및 보호를 위해 마스크를 착용해야 합니다. 마스크 착용은 미세먼지를 효과적으로 제거하고,

 

---

전체 코드

# import
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import torch
import torch.nn as nn
from torch.utils.data import Dataset
from datasets import load_dataset
import transformers
from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM, pipeline
from transformers import Trainer, TrainingArguments, AutoModelWithLMHead
from copy import deepcopy
from torch.optim import Adam
from transformers import AutoTokenizer, BloomTokenizerFast
from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer
import pandas as pd
import argparse
import copy
import logging
import json
from dataclasses import dataclass, field

def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str):
    """Collects the state dict and dump to disk."""
    state_dict = trainer.model.state_dict()
    if trainer.args.should_save:
        cpu_state_dict = {key: value.cpu() for key, value in list(state_dict.items())}
        del state_dict
        trainer._save(output_dir, state_dict=cpu_state_dict)  # noqa
        
# define argment
parser = argparse.ArgumentParser()
parser.add_argument('--data_path_1_SFT', type=str, default='./data_kochatgpt/kochatgpt_1_SFT.jsonl')
parser.add_argument('--model_name', type=str, default='gpt2', choices=['gpt2', 'bloom', 'opt'])
parser.add_argument('--max_epochs', type=int, default=2)
parser.add_argument('--train_batch_size', type=int, default=8)
parser.add_argument('--output_dir', type=str, default='./output_1_SFT')

args = parser.parse_args(args=[])

args.model_name = 'skt/kogpt2-base-v2'  # SK GPT2, https://github.com/SKT-AI/KoGPT2
# args.model_name = 'ajoublue-gpt2-base'  # 아주대, https://github.com/HeegyuKim/language-model

args.max_epochs = 2


# data config
IGNORE_INDEX = -100
DEFAULT_PAD_TOKEN = "[PAD]"
DEFAULT_EOS_TOKEN = "</s>"
DEFAULT_BOS_TOKEN = "</s>"
DEFAULT_UNK_TOKEN = "</s>"
PROMPT_DICT = {
    "prompt_input": (
        "Below is an instruction that describes a task, paired with an input that provides further context.\n"
        "아래는 작업을 설명하는 명령어와 추가적 맥락을 제공하는 입력이 짝을 이루는 예제입니다.\n\n"
        "Write a response that appropriately completes the request.\n요청을 적절히 완료하는 응답을 작성하세요.\n\n"
        "### Instruction(명령어):\n{prompt}\n\n### Input(입력):\n{input}\n\n### Response(응답):"
    ),
    "prompt_no_input": (
        "Below is an instruction that describes a task.\n"
        "아래는 작업을 설명하는 명령어입니다.\n\n"
        "Write a response that appropriately completes the request.\n명령어에 따른 요청을 적절히 완료하는 응답을 작성하세요.\n\n"
        "### Instruction(명령어):\n{prompt}\n\n### Response(응답):"
    ),
}

## 모델 준비
model = AutoModelForCausalLM.from_pretrained(args.model_name)
tokenizer = transformers.AutoTokenizer.from_pretrained(
    args.model_name,
    padding_side="right",
    model_max_length=512,
)
tokenizer.add_special_tokens(
    {
        "eos_token": DEFAULT_EOS_TOKEN,
        "bos_token": DEFAULT_BOS_TOKEN,
        "unk_token": DEFAULT_UNK_TOKEN,
    }
)
tokenizer.pad_token = tokenizer.eos_token


## prepare data
from typing import Optional, Dict, Sequence

class SFT_dataset(Dataset):
    '''SFT dataset by wygo'''
    def __init__(self, data_path_1_SFT: str, tokenizer: transformers.PreTrainedTokenizer, verbose=False):
        super(SFT_dataset, self).__init__()
        logging.warning("Loading data...")

        ## format
        pattern_instruction = 'prompt'  # instruction
        pattern_input = 'input'  # 내 데이터엔 input이 없다
        pattern_output = 'completion'  # output

        ## load dataset
#         data_path_1_SFT = 'data_kochatgpt/korean_chatgpt_1_SFT.jsonl'
        with open(data_path_1_SFT, "r", encoding='utf-8-sig') as json_file:
            list_data_dict = json.load(json_file)
        prompt_input, prompt_no_input = PROMPT_DICT["prompt_input"], PROMPT_DICT["prompt_no_input"]  # 템플릿 가져오기

        # 입력
        sources = []
        for example in list_data_dict:
            if example.get(pattern_input, "") != "":
                tmp = prompt_input.format_map(example)
            else:
                tmp = prompt_no_input.format_map(example)
            sources.append(tmp)

        # 출력
        targets = []
        for example in list_data_dict:
            targets.append(f"{example[pattern_output]}{tokenizer.eos_token}")
		examples = [s + t for s, t in zip(sources, targets)]


        # source data tokenized
        sources_tokenized = self._tokenize_fn(sources, tokenizer)  # source만
        examples_tokenized = self._tokenize_fn(examples, tokenizer)  # source + target


        ## 입력은 source, 출력은 source+target 형태이므로 입력 부분 무시 코드로 변경
        input_ids = examples_tokenized["input_ids"]
        labels = copy.deepcopy(input_ids)
        for label, source_len in zip(labels, sources_tokenized["input_ids_lens"]):
            label[:source_len] = IGNORE_INDEX  

        data_dict = dict(input_ids=input_ids, labels=labels)

        self.input_ids = data_dict["input_ids"]
        self.labels = data_dict["labels"]
        logging.warning("Loading data done!!: %d"%(len(self.labels)))

    def _tokenize_fn(self, strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict:
        tokenized_list = [
            tokenizer(
                text,
                return_tensors="pt",
                padding="longest",
                max_length=tokenizer.model_max_length,
                truncation=True,
            )
            for text in strings
        ]
        input_ids = labels = [tokenized.input_ids[0] for tokenized in tokenized_list]
        input_ids_lens = labels_lens = [
            tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item() for tokenized in tokenized_list
        ]
        return dict(
            input_ids=input_ids,
            labels=labels,
            input_ids_lens=input_ids_lens,
            labels_lens=labels_lens,
        )


    def __len__(self):
        return len(self.input_ids)


    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
        return dict(input_ids=self.input_ids[i], labels=self.labels[i])


@dataclass
class DataCollatorForSupervisedDataset(object):

    tokenizer: transformers.PreTrainedTokenizer

    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
        input_ids = torch.nn.utils.rnn.pad_sequence(
            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
        )
        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
        return dict(
            input_ids=input_ids,
            labels=labels,
            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
        )



train_dataset = SFT_dataset(data_path_1_SFT=args.data_path_1_SFT, tokenizer=tokenizer)
eval_dataset  = None
data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)


#train
training_args = TrainingArguments(
    output_dir="./test", #The output directory
    overwrite_output_dir=True, #overwrite the content of the output directory
    num_train_epochs=1, # number of training epochs
    per_device_train_batch_size=4, # batch size for training
    per_device_eval_batch_size=4,  # batch size for evaluation
    eval_steps = 3, # Number of update steps between two evaluations.
    save_steps=500, # after # steps model is saved
    warmup_steps=5,# number of warmup steps for learning rate scheduler
    prediction_loss_only=True,
    )
trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

trainer.train()
trainer.save_state()
safe_save_model_for_hf_trainer(trainer=trainer, output_dir=args.output_dir)



# Generate
from transformers import pipeline
generator = pipeline('text-generation', model=args.output_dir, tokenizer=tokenizer)
# generator = pipeline('text-generation', model=model.cpu(), tokenizer=tokenizer, config={'max_length':800})

generation_args = dict(
    num_beams=4,
    repetition_penalty=2.0,
    no_repeat_ngram_size=4,
    eos_token_id=375, # \n
    max_new_tokens=64,
    do_sample=True,
    top_k=50,
    early_stopping=True
)

list_prompt = ['불고기용 고기 한우에요?',
               '리처드 닉슨이 43대 부통령직을 수행한 년도는?',
               '시카고 오헤어 국제공항은 어디에 있어',
               '오늘 미세먼지 어때?']
list_prompt = [PROMPT_DICT['prompt_no_input'].format_map({'prompt' : tmp}) for tmp in list_prompt]

list_result = generator(list_prompt, **generation_args)
for prompt, result in zip(list_prompt, list_result):
    print(('#'*70))
    print(('completion: %s'%(result[0]['generated_text'])))