[Offline 강화학습 챗봇] Policy Gradient를 이용한 구현 도전기 - KoGPT2 Fine-tuning (1)

강화학습 챗봇에 대해 공부를 시작했을 때 처음으로 접한 논문이 "CHAI." 블로그에 정리글도 올렸다.

CHAI: A CHatbot AI for Task-Oriented Dialogue with OfflineReinforcement Learning

이미 글에 정리해뒀다 싶이 CHAI 논문에서는 언어 모델로 GPT2, 강화학습 기법으로는 Q-Learning 기법을 적용했다.

그럼 이 논문의 코드를 참고하면 가능하지 않을까? 해서 시작한 "Q-Learning을 이용한 강화학습 챗봇 구현" 도전.

결과부터 말하면 실패했다.

코드도 워낙 엉터리로 짰기에 삭제했다.

그래서...내가....이미 Q-Learning과 Policy Gradient Method를 비교하는 글을 올린 것이다.

그렇게 Q-Learning으로 시도하는 방법은 막을 내리고...

 

그러면 아래 논문과 같이 Policy Gradient를 사용하면 어떨까? 싶어 아래 논문의 코드를 참고해 구현을 시작했다.

Deep Reinforcement Learning for Dialogue Generation

 

※어디까지나 도전기일 뿐, 맞는 방법이 아니므로 따라하지 마시오.※

 

 

1. 원본 코드

1) 코드 출처

다행히 해당 논문 역시 코드가...있다!

GitHub - Ls-Dai/Deep-Reinforcement-Learning-for-Dialogue-Generation-in-PyTorch

그리고 코드 또한 CHAI에 비해서는 훨씬 간단한 편이다.

 

 

2) 코드 작동 과정

*코드 알고리즘 작동 과정을 간단하게 시각화

-원 논문에서 사용한 모델: Seq2Seq

-모델 학습 과정: 지도 학습을 이용해 모델 첫 학습 -> 강화학습을 이용해 추가적인 학습 진행

-강화학습의 보상 기준 3가지: 답변 난이도, 발화자의 답변 반복성, 전체 대화의 맥락 및 문법

 

 

 

2. 지도 학습을 이용한 모델 학습 코드 수정

*해당 논문은 2016년에 나온 논문으로, Seq2seq 모델을 이용함

→ 최근에 나온 자연어 생성에 성능이 뛰어난 모델인 KoGPT2로 변경해 이용하기로 결정

※KoGPT2 코드 참고

9. koGPT2 챗봇 만들기

 

1) 데이터 전처리 코드

(1) Voc 클래스 → Tokenizer 변경

*원본 코드

# handle dull_responses now
DULL_RESPONSES = ["I do not know what you are talking about.", "I do not know.", "You do not know.",
                  "You know what I mean.", "I know what you mean.", "You know what I am saying.", "You do not know anything."]

class Voc:
    def __init__(self, name):
        self.name = name
        self.trimmed = False
        self.word2index = {}
        self.word2count = {}
        self.index2word = {PAD_token: "PAD",
                           SOS_token: "SOS", EOS_token: "EOS"}
        self.num_words = 3  # Count SOS, EOS, PAD

    def add_sentence(self, sentence):
        for word in sentence.split(' '):
            self.add_word(word)

    def add_word(self, word):
        if word not in self.word2index:
            self.word2index[word] = self.num_words
            self.word2count[word] = 1
            self.index2word[self.num_words] = word
            self.num_words += 1
        else:
            self.word2count[word] += 1
            
    # Remove words below a certain count threshold
    def trim(self, min_count):
        if self.trimmed:
            return
        self.trimmed = True

        keep_words = []

        for k, v in self.word2count.items():
            if v >= min_count:
                keep_words.append(k)

        print('keep_words {} / {} = {:.4f}'.format(
            len(keep_words), len(self.word2index), len(
                keep_words) / len(self.word2index)
        ))

        # Reinitialize dictionaries
        self.word2index = {}
        self.word2count = {}
        self.index2word = {PAD_token: "PAD",
                           SOS_token: "SOS", EOS_token: "EOS"}
        self.num_words = 3  # Count default tokens

        for word in keep_words:
            self.add_word(word)

(모든 코드를 일일이 분석하지는 않고 간단하게만 분석)

*Voc 클래스: 불러오는 모든 대화 데이터 내 문장을 단어로 쪼갠 후 해당 단어들을 가져와 저장 (일종의 단어 사전)

-add_sentence: 하나의 문장을 가져온 후, 해당 문장 내 단어를 클래스 사전에 추가

-add_word: 단어 추가 및 각 단어 별 인덱스 지정

-trim: 최소한의 빈도(사용자 지정)를 넘지 못하는 단어는 사전에서 제거

 

*Voc 클래스의 용도를 보면 KoGPT2 구성 중에서 토크나이저 용도와 비슷하다.

⇒ Voc 클래스 대신 Tokenizer 사용

*수정 코드

U_TKN = "<usr>"
S_TKN = "<sys>"
BOS = "<bos>"
EOS = "<eos>"
MASK = "<unused0>"
SENT = "<unused1>"
PAD = "<pad>"

TOKENIZER = PreTrainedTokenizerFast.from_pretrained("skt/kogpt2-base-v2",
            bos_token=BOS, eos_token=EOS, unk_token="<unk>",
            pad_token=PAD, mask_token=MASK)

# TOKENIZER.get_vocab(): TOKENIZER 단어 사전 불러오기
# TOKENIZER.add_tokens(): TOKENIZER에 새로운 단어 추가

-토크나이저를 이용해 새 단어 추가 및 이미 학습된 단어 사전 불러오기 가능

 

(2) 텍스트 전처리 및 데이터 쪼개기 코드 수정

*원본 코드

def unicode_to_ascii(s):
    return ''.join(
        c for c in unicodedata.normalize('NFD', s)
        if unicodedata.category(c) != 'Mn')

# Lowercase, trim, and remove non-letter characters
def normalize_string(s):
    s = unicode_to_ascii(s.lower().strip())
    s = re.sub(r"([.!?])", r" \1", s)
    s = re.sub(r"[^a-zA-Z.!?]+", r" ", s)
    s = re.sub(r"\s+", r" ", s).strip()
    return s

# Read query/response pairs and return a voc object
def read_vocs(datafile, corpus_name):
    print("Reading lines...")
    # Read the file and split into lines
    lines = open(datafile, encoding='utf-8').\
        read().strip().split('\n')
    # Split every line into pairs and normalize
    pairs = [[normalize_string(s) for s in l.split('\t')] for l in lines]
    voc = Voc(corpus_name)
    return voc, pairs

# Returns True iff both sentences in a pair 'p' are under the MAX_LENGTH threshold
def filter_pair(p, max_length=15):
    # Input sequences need to preserve the last word for EOS token
    return len(p[0].split(' ')) < max_length and len(p[1].split(' ')) < max_length

# Filter pairs using filterPair condition
def filter_pairs(pairs):
    return [pair for pair in pairs if filter_pair(pair)]

# Using the functions defined above, return a populated voc object and pairs list
def load_prepare_data(corpus, corpus_name, datafile, save_dir):
    print("Start preparing training data ...")
    voc, pairs = read_vocs(datafile, corpus_name)
    print("Read {!s} sentence pairs".format(len(pairs)))
    pairs = filter_pairs(pairs)
    print("Trimmed to {!s} sentence pairs".format(len(pairs)))
    print("Counting words...")
    for d in DULL_RESPONSES:
        voc.add_sentence(d)
    for pair in pairs:
        voc.add_sentence(pair[0])
        voc.add_sentence(pair[1])
    return voc, pairs

def trim_rare_words(voc, pairs, min_count=3):
    # Trim words used under the MIN_COUNT from the voc
    voc.trim(min_count)
    # Filter out pairs with trimmed words
    keep_pairs = []
    for pair in pairs:
        input_sentence = pair[0]
        output_sentence = pair[1]
        keep_input = True
        keep_output = True
        # Check input sentence
        for word in input_sentence.split(' '):
            if word not in voc.word2index:
                keep_input = False
                break
        # Check output sentence
        for word in output_sentence.split(' '):
            if word not in voc.word2index:
                keep_output = False
                break

        # Only keep pairs that do not contain trimmed word(s) in their input or output sentence
        if keep_input and keep_output:
            keep_pairs.append(pair)

    print("Trimmed from {} pairs to {}, {:.4f} of total".format(
        len(pairs), len(keep_pairs), len(keep_pairs) / len(pairs)))
    return keep_pairs

def token_to_id(voc, sentence):
    return [voc.word2index[word] for word in sentence.split(' ')] + [EOS_token]

def zero_padding(l, fillvalue=PAD_token):
    return list(itertools.zip_longest(*l, fillvalue=fillvalue))

def binary_matrix(l, value=PAD_token):
    m = []
    for i, seq in enumerate(l):
        m.append([])
        for token in seq:
            if token == PAD_token:
                m[i].append(0)
            else:
                m[i].append(1)
    return m

# Returns padded input sequence tensor and lengths
def input_var(l, voc):
    indexes_batch = [token_to_id(voc, sentence) for sentence in l]
    print("Voc-input_var: indexes_batch", indexes_batch)
    lengths = torch.tensor([len(indexes) for indexes in indexes_batch])
    pad_list = zero_padding(indexes_batch)
    pad_var = torch.LongTensor(pad_list)
    print("Voc-input_var: lengths: ", lengths, "pad_list", pad_list[:3], "pad_var", pad_var[:10])
    print("shape pad_var: ", pad_var.shape)
    return pad_var, lengths

# Returns padded target sequence tensor, padding mask, and max target length
def output_var(l, voc):
    indexes_batch = [token_to_id(voc, sentence) for sentence in l]
    print("Voc-output_var: indexes_batch", indexes_batch)
    max_target_len = max([len(indexes) for indexes in indexes_batch])
    pad_list = zero_padding(indexes_batch)
    mask = binary_matrix(pad_list)
    mask = torch.BoolTensor(mask)
    pad_var = torch.LongTensor(pad_list)
    print("Voc-output_var: max_target_len: ", max_target_len, "pad_list", pad_list[:3], "mask: ", mask[:10], "pad_var", pad_var[:10])
    print("shape", "mask: ", mask.shape, "pad_var", pad_var.shape)
    return pad_var, mask, max_target_len

# Returns all items for a given batch of pairs
def batch_2_train_data(voc, pair_batch):
    pair_batch.sort(key=lambda x: len(x[0].split(" ")), reverse=True)
    input_batch, output_batch = [], []
    for pair in pair_batch:
        input_batch.append(pair[0])
        output_batch.append(pair[1])
    inp, lengths = input_var(input_batch, voc)
    output, mask, max_target_len = output_var(output_batch, voc)
    print("Voc-batch_2_train_data-inp: ", inp[:10], "lengths", lengths)
    print("Voc-batch_2_train_data-output: ", output[:10], "mask: ", mask[:10], "max_target_len: ", max_target_len)
    return inp, lengths, output, mask, max_target_len

# Load/Assemble voc and pairs
voc, pairs = load_prepare_data(corpus, corpus_name, datafile, save_dir)
# Print some pairs to validate
print("\npairs:")
for pair in pairs[:10]:
    print(pair)
pairs = trim_rare_words(voc, pairs, min_count=3)

*normalize_string: 텍스트 내 다양한 문장부호 제거

*read_vocs: 주어진 데이터셋 내 문장에서 각 단어들의 다양한 문장부호를 제거한 후, 클래스 Voc을 불러와 후에 해당 단어들을 Voc 클래스, 즉 단어 사전에 추가

*filter_pair: 원래 대화 데이터셋에서 쪼갠 두 문장마다 모든 문장이 최대 길이를 넘는지 아닌지 확인

*filter_pairs: 원래 대화 데이터셋에서 쪼갠 두 문장끼리 모두 최대 길이 이하인 경우, 해당 두 문장끼리 하나의 데이터로 묶어서 반환 (아래 그림 설명 추가)
*load_prepare_data: 대화 데이터셋이 주어지면 해당 데이터 내 단어를 이용해 단어 사전 생성 + 두 문장씩 쪼갬

 

filter_pairs 함수를 시각화한 그림

 

⇒ KoGPT2 코드를 이용할 때는 csv 형식의 대화 데이터셋과 토크나이저를 사용하므로, 따로 단어 사전 클래스 정의할 필요 없이, 위와 같은 형태로 데이터를 쪼개는 코드 제작

*수정 코드

first = []
second = []
for i in range(len(data)):
  turn = data.iloc[i]
  for j in range(0, 8):  #사용하는 데이터셋이 최대 8개의 열밖에 존재하지 않아서 8로 지정한 것
    if (j==6) or (pd.isnull(turn[j+2])):
      first.append(str(turn[j])+SEP)
      second.append(str(turn[j+1])+EOS)
      break
    else:
      first.append(str(turn[j])+SEP)
      second.append(str(turn[j+1])+SEP)

new_data = pd.DataFrame({"first": first, "second": second})

 

*데이터 쪼개기 전

 

*데이터 쪼갠 후

-주어진 대화 데이터셋에서 각 발화자의 문장을 구분할 수 있도록 "<sep>" 토큰 추가

(사실 실제로 <sep> 토큰이 추가된 시점은 이보다 훨씬 뒤다)

 

 

2) 모델 학습 코드 수정: Seq2seq → KoGPT2

*원본 코드

class EncoderRNN(nn.Module):
    def __init__(self, hidden_size, embedding, n_layers=1, dropout=0):
        super(EncoderRNN, self).__init__()
        self.n_layers = n_layers
        self.hidden_size = hidden_size
        self.embedding = embedding
        ...
        
    def forward(self, input_seq, input_lengths, hidden=None):
        # Convert word indexes to embeddings
        embedded = self.embedding(input_seq)
        # Pack padded batch of sequences for RNN module
        ...

# Luong attention layer
class Attn(nn.Module):
    def __init__(self, method, hidden_size):
        super(Attn, self).__init__()
        self.method = method
        ...

    def dot_score(self, hidden, encoder_output):
        print("Attn-dot_score:", "hidden", hidden[:3], "encoder_output", encoder_output[:3])
        return torch.sum(hidden * encoder_output, dim=2)

    def general_score(self, hidden, encoder_output):
        energy = self.attn(encoder_output)
        return torch.sum(hidden * energy, dim=2)

    def concat_score(self, hidden, encoder_output):
        energy = self.attn(torch.cat(
            (hidden.expand(encoder_output.size(0), -1, -1), encoder_output), 2)).tanh()
        return torch.sum(self.v * energy, dim=2)

    def forward(self, hidden, encoder_outputs):
        # Calculate the attention weights (energies) based on the given method
        if self.method == 'general':
            attn_energies = self.general_score(hidden, encoder_outputs)
        elif self.method == 'concat':
            attn_energies = self.concat_score(hidden, encoder_outputs)
        elif self.method == 'dot':
            attn_energies = self.dot_score(hidden, encoder_outputs)
        ...

class LuongAttnDecoderRNN(nn.Module):
    def __init__(self, attn_model, embedding, hidden_size, output_size, n_layers=1, dropout=0.1):
        super(LuongAttnDecoderRNN, self).__init__()

        # Keep for reference
        self.attn_model = attn_model
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.n_layers = n_layers
        self.dropout = dropout
        ...

    def forward(self, input_step, last_hidden, encoder_outputs):
        # Note: we run this one step (word) at a time
        # Get embedding of current input word
        embedded = self.embedding(input_step)
        embedded = self.embedding_dropout(embedded)
        # Forward through unidirectional GRU
        ...

class GreedySearchDecoder(nn.Module):
    def __init__(self, encoder, decoder):
        super(GreedySearchDecoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder

    def forward(self, input_seq, input_length, max_length):
        # device choice
        USE_CUDA = torch.cuda.is_available()
        device = torch.device("cuda" if USE_CUDA else "cpu")

        # Forward input through encoder model
        encoder_outputs, encoder_hidden = self.encoder(input_seq, input_length)
        # Prepare encoder's final hidden layer to be first hidden input to the decoder
        ...

*원 논문에서는 Encoder와 Decoder 구조에 Attention 메커니즘을 사용하는 Seq2seq 모델 사용

⇒ KoGPT2를 이용하기 위해 KoGPT2에 맞는 데이터셋 및 모델 정의

*수정 코드

class ChatbotDataset(Dataset):
  def __init__(self, chats, max_len):
    self._data = chats
    self.max_len = max_len
    self.tokenizer = TOKENIZER

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

  def __getitem__(self, idx):  # 로드한 챗봇 데이터를 차례차례 DataLoader로 넘겨주는 메서드
    turn = self._data.iloc[idx]
    q = turn[0]
    a = turn[1]

    q_toked = self.tokenizer.tokenize(q)
    q_len = len(q_toked)
    a_toked = self.tokenizer.tokenize(a)
    a_len = len(a_toked)

    #질문의 길이가 최대길이보다 크면
    if q_len > self.max_len:
      a_len = self.max_len - q_len        #답변의 길이를 최대길이 - 질문길이
      if a_len <= 0:
        q_toked = q_toked[-(int(self.max_len / 2)) :]
        q_len = len(q_toked)
        a_len = self.max_len - q_len
      a_toked = a_toked[:a_len]
      a_len = len(a_toked)

    #질문의 길이 + 답변의 길이가 최대길이보다 크면
    if q_len + a_len > self.max_len:
      a_len = self.max_len - q_len        #답변의 길이를 최대길이 - 질문길이
      if a_len <= 0:
        q_toked = q_toked[-(int(self.max_len / 2)):]
        q_len = len(q_toked)
        a_len = self.max_len - q_len              #답변의 길이를 최대길이 - 질문길이
      a_toked = a_toked[:a_len]
      a_len = len(a_toked)

    # 답변 labels = [mask, mask, ...., mask, ..., <bos>,..답변.. <eos>, <pad>....]
    labels = [self.tokenizer.mask_token] * q_len + a_toked[1:]

    # mask = 질문길이 0 + 답변길이 1 + 나머지 0
    mask = [0] * q_len + [1] * a_len + [0] * (self.max_len - q_len - a_len)
    # 답변 labels을 index 로 만든다.
    labels_ids = self.tokenizer.convert_tokens_to_ids(labels)
    # 최대길이만큼 PADDING
    while len(labels_ids) < self.max_len:
      labels_ids += [self.tokenizer.pad_token_id]

    token_ids = self.tokenizer.convert_tokens_to_ids(q_toked + a_toked)
    while len(token_ids) < self.max_len:
      token_ids += [self.tokenizer.pad_token_id]

    return (token_ids, np.array(mask), labels_ids)
    
def collate_batch(batch):
    data = np.array([item[0] for item in batch])
    mask = np.array([item[1] for item in batch])
    label = np.array([item[2] for item in batch])
    return torch.LongTensor(data).to(device), torch.LongTensor(mask).to(device), torch.LongTensor(label).to(device)

train_set = ChatbotDataset(ChatbotData, max_len=128)
train_dataloader = DataLoader(train_set, batch_size=4, num_workers=0, shuffle=True, collate_fn=collate_batch)

learning_rate = 3e-5
criterion = torch.nn.CrossEntropyLoss(reduction="none").to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
Sneg = -1e18

*학습 시 문장의 최대 길이: 128

*배치 사이즈: 4 (구글 코랩 프로 플러스의 메모리 한계)