자연어 생성에 강화학습 적용

!pip install transformers

import tensorflow as tf
from transformers import TFAutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained('gpt2')
model = TFAutoModelForCausalLM.from_pretrained('gpt2')

input_ids = tokenizer.encode('I like this movie', return_tensors='tf')

result = model.generate(input_ids, max_length=50)

tokenizer.decode(result[0])

res = result[0]

from collections import Counter

def score(res):
    counter = Counter()
    tokens = res.numpy()
    for i in range(0, len(tokens) - 3):
        trigram = tokens[i:i+3]
        trigram = tuple(trigram)
        counter[trigram] += 1
    return len(counter) / sum(counter.values())

score(res)

with tf.GradientTape() as tape:
    output = model(result, labels=result)
    s = score(result[0])
    loss = (s - 0.5) * output[0]
    mean_loss = tf.reduce_mean(loss)

grads = tape.gradient(mean_loss, model.trainable_variables)

opt = tf.keras.optimizers.Adam(learning_rate=5e-4, epsilon=1e-08)

opt.apply_gradients(zip(grads, model.trainable_variables))

result = model.generate(input_ids, max_length=50)

tokenizer.decode(result[0])