KerasNLP로 처음부터 트랜스포머 사전 트레이닝

!pip install -q --upgrade keras-nlp
!pip install -q --upgrade keras  # Keras 3으로 업그레이드.

import os

os.environ["KERAS_BACKEND"] = "jax"  # 또는 "tensorflow" 또는 "torch"


import keras_nlp
import tensorflow as tf
import keras

# 사전 트레이닝 데이터 다운로드.
keras.utils.get_file(
    origin="https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-103-raw-v1.zip",
    extract=True,
)
wiki_dir = os.path.expanduser("~/.keras/datasets/wikitext-103-raw/")

# 미세 조정 데이터 다운로드.
keras.utils.get_file(
    origin="https://dl.fbaipublicfiles.com/glue/data/SST-2.zip",
    extract=True,
)
sst_dir = os.path.expanduser("~/.keras/datasets/SST-2/")

# 어휘 데이터 다운로드.
vocab_file = keras.utils.get_file(
    origin="https://storage.googleapis.com/tensorflow/keras-nlp/examples/bert/bert_vocab_uncased.txt",
)

# 전처리 파라미터.
PRETRAINING_BATCH_SIZE = 128
FINETUNING_BATCH_SIZE = 32
SEQ_LENGTH = 128
MASK_RATE = 0.25
PREDICTIONS_PER_SEQ = 32

# 모델 파라미터.
NUM_LAYERS = 3
MODEL_DIM = 256
INTERMEDIATE_DIM = 512
NUM_HEADS = 4
DROPOUT = 0.1
NORM_EPSILON = 1e-5

# 트레이닝 파라미터.
PRETRAINING_LEARNING_RATE = 5e-4
PRETRAINING_EPOCHS = 8
FINETUNING_LEARNING_RATE = 5e-5
FINETUNING_EPOCHS = 3

# SST-2 데이터 로드.
sst_train_ds = tf.data.experimental.CsvDataset(
    sst_dir + "train.tsv", [tf.string, tf.int32], header=True, field_delim="\t"
).batch(FINETUNING_BATCH_SIZE)
sst_val_ds = tf.data.experimental.CsvDataset(
    sst_dir + "dev.tsv", [tf.string, tf.int32], header=True, field_delim="\t"
).batch(FINETUNING_BATCH_SIZE)

# wikitext-103 데이터 로드 및 짧은 줄 필터링.
wiki_train_ds = (
    tf.data.TextLineDataset(wiki_dir + "wiki.train.raw")
    .filter(lambda x: tf.strings.length(x) > 100)
    .batch(PRETRAINING_BATCH_SIZE)
)
wiki_val_ds = (
    tf.data.TextLineDataset(wiki_dir + "wiki.valid.raw")
    .filter(lambda x: tf.strings.length(x) > 100)
    .batch(PRETRAINING_BATCH_SIZE)
)

# sst-2 데이터셋을 살펴봅니다.
print(sst_train_ds.unbatch().batch(4).take(1).get_single_element())

# 이 레이어는 입력 문장을 어휘 크기와 같은 크기의 1과 0 리스트로 변환합니다.
# 이 리스트는 단어가 존재하거나 존재하지 않음을 나타냅니다.
multi_hot_layer = keras.layers.TextVectorization(
    max_tokens=4000, output_mode="multi_hot"
)
multi_hot_layer.adapt(sst_train_ds.map(lambda x, y: x))
multi_hot_ds = sst_train_ds.map(lambda x, y: (multi_hot_layer(x), y))
multi_hot_val_ds = sst_val_ds.map(lambda x, y: (multi_hot_layer(x), y))

# 그런 다음 이 레이어에 대해 선형 회귀를 학습합니다. 이것이 베이스라인 모델입니다!

inputs = keras.Input(shape=(4000,), dtype="int32")
outputs = keras.layers.Dense(1, activation="sigmoid")(inputs)
baseline_model = keras.Model(inputs, outputs)
baseline_model.compile(loss="binary_crossentropy", metrics=["accuracy"])
baseline_model.fit(multi_hot_ds, validation_data=multi_hot_val_ds, epochs=5)

# sequence_length를 설정하면,
# 토큰 출력을 (batch_size, SEQ_LENGTH) 모양으로 자르거나(trim) 패딩합니다.
tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(
    vocabulary=vocab_file,
    sequence_length=SEQ_LENGTH,
    lowercase=True,
    strip_accents=True,
)
# mask_selection_length를 설정하면,
# 마스크 출력을 (batch_size, PREDICTIONS_PER_SEQ) 형태로 자르거나(trim) 패딩합니다.
masker = keras_nlp.layers.MaskedLMMaskGenerator(
    vocabulary_size=tokenizer.vocabulary_size(),
    mask_selection_rate=MASK_RATE,
    mask_selection_length=PREDICTIONS_PER_SEQ,
    mask_token_id=tokenizer.token_to_id("[MASK]"),
)


def preprocess(inputs):
    inputs = tokenizer(inputs)
    outputs = masker(inputs)
    # 마스킹 레이어 출력을 keras.Model.fit()에서 사용할 수 있는
    # (features, labels, and weights) 튜플로 분리합니다.
    features = {
        "token_ids": outputs["token_ids"],
        "mask_positions": outputs["mask_positions"],
    }
    labels = outputs["mask_ids"]
    weights = outputs["mask_weights"]
    return features, labels, weights


# prefetch()를 사용하여 CPU에서 전처리된 배치를 실시간으로 미리 계산합니다.
pretrain_ds = wiki_train_ds.map(
    preprocess, num_parallel_calls=tf.data.AUTOTUNE
).prefetch(tf.data.AUTOTUNE)
pretrain_val_ds = wiki_val_ds.map(
    preprocess, num_parallel_calls=tf.data.AUTOTUNE
).prefetch(tf.data.AUTOTUNE)

# 단일 입력 예시 미리보기
# 마스크는 셀을 실행할 때마다 변경됩니다.
print(pretrain_val_ds.take(1).get_single_element())

inputs = keras.Input(shape=(SEQ_LENGTH,), dtype="int32")

# 위치 임베딩을 사용해 토큰을 임베딩합니다.
embedding_layer = keras_nlp.layers.TokenAndPositionEmbedding(
    vocabulary_size=tokenizer.vocabulary_size(),
    sequence_length=SEQ_LENGTH,
    embedding_dim=MODEL_DIM,
)
outputs = embedding_layer(inputs)

# 임베딩에 레이어 정규화와 드롭아웃을 적용합니다.
outputs = keras.layers.LayerNormalization(epsilon=NORM_EPSILON)(outputs)
outputs = keras.layers.Dropout(rate=DROPOUT)(outputs)

# 여러 개의 인코더 블록을 추가합니다.
for i in range(NUM_LAYERS):
    outputs = keras_nlp.layers.TransformerEncoder(
        intermediate_dim=INTERMEDIATE_DIM,
        num_heads=NUM_HEADS,
        dropout=DROPOUT,
        layer_norm_epsilon=NORM_EPSILON,
    )(outputs)

encoder_model = keras.Model(inputs, outputs)
encoder_model.summary()

# 마스킹된 언어 모델 헤드를 연결하여 사전 트레이닝된 모델을 만듭니다.
inputs = {
    "token_ids": keras.Input(shape=(SEQ_LENGTH,), dtype="int32", name="token_ids"),
    "mask_positions": keras.Input(
        shape=(PREDICTIONS_PER_SEQ,), dtype="int32", name="mask_positions"
    ),
}

# 토큰을 인코딩합니다.
encoded_tokens = encoder_model(inputs["token_ids"])

# 마스크된 입력 토큰마다 출력 단어를 예측합니다.
# 입력 토큰 임베딩을 사용하여 인코딩된 벡터를 어휘 로짓으로 변환하는 것은,
# 트레이닝 효율성을 높이는 것으로 알려져 있습니다.
outputs = keras_nlp.layers.MaskedLMHead(
    token_embedding=embedding_layer.token_embedding,
    activation="softmax",
)(encoded_tokens, mask_positions=inputs["mask_positions"])

# 사전 트레이닝 모델을 정의하고 컴파일합니다.
pretraining_model = keras.Model(inputs, outputs)
pretraining_model.compile(
    loss="sparse_categorical_crossentropy",
    optimizer=keras.optimizers.AdamW(PRETRAINING_LEARNING_RATE),
    weighted_metrics=["sparse_categorical_accuracy"],
    jit_compile=True,
)

# wiki 텍스트 데이터셋으로 모델을 사전 트레이닝합니다.
pretraining_model.fit(
    pretrain_ds,
    validation_data=pretrain_val_ds,
    epochs=PRETRAINING_EPOCHS,
)

# 이후의 미세 조정을 위해 이 베이스 모델을 저장합니다.
encoder_model.save("encoder_model.keras")

def preprocess(sentences, labels):
    return tokenizer(sentences), labels


# prefetch()를 사용하여 CPU에서 전처리된 배치를 실시간으로 미리 계산합니다.
finetune_ds = sst_train_ds.map(
    preprocess, num_parallel_calls=tf.data.AUTOTUNE
).prefetch(tf.data.AUTOTUNE)
finetune_val_ds = sst_val_ds.map(
    preprocess, num_parallel_calls=tf.data.AUTOTUNE
).prefetch(tf.data.AUTOTUNE)

# 단일 입력 예제를 미리 확인합니다.
print(finetune_val_ds.take(1).get_single_element())

# 디스크에서 인코더 모델을 다시 로드하여 처음부터 미세 조정을 시작합니다.
encoder_model = keras.models.load_model("encoder_model.keras", compile=False)

# 토큰화된 입력을 입력으로 받습니다.
inputs = keras.Input(shape=(SEQ_LENGTH,), dtype="int32")

# 토큰을 인코딩하고 풀링합니다.
encoded_tokens = encoder_model(inputs)
pooled_tokens = keras.layers.GlobalAveragePooling1D()(encoded_tokens[0])

# 출력 레이블을 예측합니다.
outputs = keras.layers.Dense(1, activation="sigmoid")(pooled_tokens)

# 미세 조정 모델을 정의하고 컴파일합니다.
finetuning_model = keras.Model(inputs, outputs)
finetuning_model.compile(
    loss="binary_crossentropy",
    optimizer=keras.optimizers.AdamW(FINETUNING_LEARNING_RATE),
    metrics=["accuracy"],
)

# SST-2 작업에 대해 모델을 미세 조정합니다.
finetuning_model.fit(
    finetune_ds,
    validation_data=finetune_val_ds,
    epochs=FINETUNING_EPOCHS,
)