미니어처 GPT로 텍스트 생성

# 백엔드를 TensorFlow로 설정합니다.
# 이 코드는 `tensorflow`와 `torch`에서 작동하며, JAX에서는 작동하지 않습니다.
# 이는 `causal_attention_mask()`에서 사용된
# JAX의 `jax.numpy.tile` 함수가 jit 스코프 내에서 동적 `reps` 인자를 지원하지 않기 때문입니다.
# JAX에서 이 코드를 작동시키려면,
# `causal_attention_mask` 함수 내부를 jit 컴파일을 방지하는 데코레이터로 감싸서 처리할 수 있습니다:
# `with jax.ensure_compile_time_eval():`.
import os

os.environ["KERAS_BACKEND"] = "tensorflow"

import keras
from keras import layers
from keras import ops
from keras.layers import TextVectorization
import numpy as np
import os
import string
import random
import tensorflow
import tensorflow.data as tf_data
import tensorflow.strings as tf_strings

def causal_attention_mask(batch_size, n_dest, n_src, dtype):
    """
    셀프 어텐션에서 점곱 행렬(dot product matrix)의 상단 절반을 마스킹합니다.
    이는 미래의 토큰에서 현재 토큰으로 정보가 흐르는 것을 방지합니다.
    오른쪽 아래 모서리에서 시작하여, 하단 삼각형에 1을 채웁니다.
    """
    i = ops.arange(n_dest)[:, None]
    j = ops.arange(n_src)
    m = i >= j - n_src + n_dest
    mask = ops.cast(m, dtype)
    mask = ops.reshape(mask, [1, n_dest, n_src])
    mult = ops.concatenate(
        [ops.expand_dims(batch_size, -1), ops.convert_to_tensor([1, 1])], 0
    )
    return ops.tile(mask, mult)


class TransformerBlock(layers.Layer):
    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1):
        super().__init__()
        self.att = layers.MultiHeadAttention(num_heads, embed_dim)
        self.ffn = keras.Sequential(
            [
                layers.Dense(ff_dim, activation="relu"),
                layers.Dense(embed_dim),
            ]
        )
        self.layernorm1 = layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = layers.LayerNormalization(epsilon=1e-6)
        self.dropout1 = layers.Dropout(rate)
        self.dropout2 = layers.Dropout(rate)

    def call(self, inputs):
        input_shape = ops.shape(inputs)
        batch_size = input_shape[0]
        seq_len = input_shape[1]
        causal_mask = causal_attention_mask(batch_size, seq_len, seq_len, "bool")
        attention_output = self.att(inputs, inputs, attention_mask=causal_mask)
        attention_output = self.dropout1(attention_output)
        out1 = self.layernorm1(inputs + attention_output)
        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output)
        return self.layernorm2(out1 + ffn_output)

class TokenAndPositionEmbedding(layers.Layer):
    def __init__(self, maxlen, vocab_size, embed_dim):
        super().__init__()
        self.token_emb = layers.Embedding(input_dim=vocab_size, output_dim=embed_dim)
        self.pos_emb = layers.Embedding(input_dim=maxlen, output_dim=embed_dim)

    def call(self, x):
        maxlen = ops.shape(x)[-1]
        positions = ops.arange(0, maxlen, 1)
        positions = self.pos_emb(positions)
        x = self.token_emb(x)
        return x + positions

vocab_size = 20000  # 상위 20,000개의 단어만 고려
maxlen = 80  # 최대 시퀀스 크기
embed_dim = 256  # 각 토큰에 대한 임베딩 크기
num_heads = 2  # 어텐션 헤드 수
feed_forward_dim = 256  # Transformer 내부 피드 포워드 네트워크의 히든 레이어 크기


def create_model():
    inputs = layers.Input(shape=(maxlen,), dtype="int32")
    embedding_layer = TokenAndPositionEmbedding(maxlen, vocab_size, embed_dim)
    x = embedding_layer(inputs)
    transformer_block = TransformerBlock(embed_dim, num_heads, feed_forward_dim)
    x = transformer_block(x)
    outputs = layers.Dense(vocab_size)(x)
    model = keras.Model(inputs=inputs, outputs=[outputs, x])
    loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
    model.compile(
        "adam",
        loss=[loss_fn, None],
    )  # Transformer 블록에서의 워드 임베딩 기반으로 한 최적화 및 손실 없음
    return model

!curl -O https://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz
!tar -xf aclImdb_v1.tar.gz

batch_size = 128

# 데이터셋은 각 리뷰가 개별 텍스트 파일에 포함되어 있습니다.
# 텍스트 파일은 네 개의 다른 폴더에 존재합니다.
# 모든 파일의 목록을 생성합니다.
filenames = []
directories = [
    "aclImdb/train/pos",
    "aclImdb/train/neg",
    "aclImdb/test/pos",
    "aclImdb/test/neg",
]
for dir in directories:
    for f in os.listdir(dir):
        filenames.append(os.path.join(dir, f))

print(f"{len(filenames)} files")

# 텍스트 파일에서 데이터셋을 생성합니다.
random.shuffle(filenames)
text_ds = tf_data.TextLineDataset(filenames)
text_ds = text_ds.shuffle(buffer_size=256)
text_ds = text_ds.batch(batch_size)


def custom_standardization(input_string):
    """HTML 줄 바꿈 태그를 제거하고 구두점을 처리합니다."""
    lowercased = tf_strings.lower(input_string)
    stripped_html = tf_strings.regex_replace(lowercased, "<br />", " ")
    return tf_strings.regex_replace(stripped_html, f"([{string.punctuation}])", r" \1")


# 벡터화 레이어를 생성하고 텍스트에 맞춥니다.
vectorize_layer = TextVectorization(
    standardize=custom_standardization,
    max_tokens=vocab_size - 1,
    output_mode="int",
    output_sequence_length=maxlen + 1,
)
vectorize_layer.adapt(text_ds)
vocab = vectorize_layer.get_vocabulary()  # 토큰 인덱스로부터 단어를 다시 얻기 위함


def prepare_lm_inputs_labels(text):
    """
    단어 시퀀스를 1 위치씩 이동시켜 위치 (i)의 대상이 위치 (i+1)에 있는 단어가 되도록 합니다.
    모델은 위치 (i)까지의 모든 단어를 사용하여 다음 단어를 예측할 것입니다.
    """
    text = tensorflow.expand_dims(text, -1)
    tokenized_sentences = vectorize_layer(text)
    x = tokenized_sentences[:, :-1]
    y = tokenized_sentences[:, 1:]
    return x, y


text_ds = text_ds.map(prepare_lm_inputs_labels, num_parallel_calls=tf_data.AUTOTUNE)
text_ds = text_ds.prefetch(tf_data.AUTOTUNE)

class TextGenerator(keras.callbacks.Callback):
    """트레이닝된 모델에서 텍스트를 생성하는 콜백.
    1. 모델에 시작 프롬프트를 제공
    2. 다음 토큰에 대한 확률 예측
    3. 다음 토큰을 샘플링하여 다음 입력에 추가

    Arguments:
        max_tokens: 정수, 프롬프트 이후에 생성할 토큰의 수.
        start_tokens: 정수 리스트, 시작 프롬프트의 토큰 인덱스.
        index_to_word: 문자열 리스트, TextVectorization 레이어에서 얻음.
        top_k: 정수, `top_k` 토큰 예측에서 샘플링.
        print_every: 정수, 이 만큼의 에포크 이후에 출력.
    """

    def __init__(
        self, max_tokens, start_tokens, index_to_word, top_k=10, print_every=1
    ):
        self.max_tokens = max_tokens
        self.start_tokens = start_tokens
        self.index_to_word = index_to_word
        self.print_every = print_every
        self.k = top_k

    def sample_from(self, logits):
        logits, indices = ops.top_k(logits, k=self.k, sorted=True)
        indices = np.asarray(indices).astype("int32")
        preds = keras.activations.softmax(ops.expand_dims(logits, 0))[0]
        preds = np.asarray(preds).astype("float32")
        return np.random.choice(indices, p=preds)

    def detokenize(self, number):
        return self.index_to_word[number]

    def on_epoch_end(self, epoch, logs=None):
        start_tokens = [_ for _ in self.start_tokens]
        if (epoch + 1) % self.print_every != 0:
            return
        num_tokens_generated = 0
        tokens_generated = []
        while num_tokens_generated <= self.max_tokens:
            pad_len = maxlen - len(start_tokens)
            sample_index = len(start_tokens) - 1
            if pad_len < 0:
                x = start_tokens[:maxlen]
                sample_index = maxlen - 1
            elif pad_len > 0:
                x = start_tokens + [0] * pad_len
            else:
                x = start_tokens
            x = np.array([x])
            y, _ = self.model.predict(x, verbose=0)
            sample_token = self.sample_from(y[0][sample_index])
            tokens_generated.append(sample_token)
            start_tokens.append(sample_token)
            num_tokens_generated = len(tokens_generated)
        txt = " ".join(
            [self.detokenize(_) for _ in self.start_tokens + tokens_generated]
        )
        print(f"generated text:\n{txt}\n")


# 시작 프롬프트 토큰화
word_to_index = {}
for index, word in enumerate(vocab):
    word_to_index[word] = index

start_prompt = "this movie is"
start_tokens = [word_to_index.get(_, 1) for _ in start_prompt.split()]
num_tokens_generated = 40
text_gen_callback = TextGenerator(num_tokens_generated, start_tokens, vocab)

model = create_model()

model.fit(text_ds, verbose=2, epochs=25, callbacks=[text_gen_callback])