<jupyter_start><jupyter_text>Training DistilKoBERT on KorNLI dataset* kornli dataset: https://jeongukjae.github.io/tfds-korean/datasets/kornli.html* distilkobert    * encoder: https://tfhub.dev/jeongukjae/distilkobert_cased_L-3_H-768_A-12/1    * preprocessor: https://tfhub.dev/jeongukjae/distilkobert_cased_preprocess/1  Install pacakges<jupyter_code>!pip install -q \
    tensorflow-text \
    tfds-korean<jupyter_output><empty_output><jupyter_text>Prepare environments<jupyter_code>import math

import tensorflow as tf
import tensorflow_text as text
import tensorflow_hub as hub
import tensorflow_datasets as tfds

import tfds_korean.kornli
import os

os.environ["TFHUB_MODEL_LOAD_FORMAT"] = "UNCOMPRESSED"
cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
strategy = tf.distribute.TPUStrategy(cluster_resolver)<jupyter_output>INFO:tensorflow:Deallocate tpu buffers before initializing tpu system.<jupyter_text>Set up hyperparameters and build models<jupyter_code>BATCH_SIZE = 64
LEARNING_RATE = 1e-4
EPOCHS = 5
WARMUP_RATE = 0.05
def create_preprocessing_model():
    preprocessor = hub.load("https://tfhub.dev/jeongukjae/distilkobert_cased_preprocess/1")
    tokenize = hub.KerasLayer(preprocessor.tokenize)
    bert_pack_inputs = hub.KerasLayer(preprocessor.bert_pack_inputs)

    text_inputs = [
        tf.keras.Input([], dtype=tf.string),
        tf.keras.Input([], dtype=tf.string),
    ]
    tokens = [tokenize(item) for item in text_inputs]
    model_inputs = bert_pack_inputs(tokens)
    return tf.keras.Model(text_inputs, model_inputs)


def create_model():
    encoder = hub.KerasLayer("https://tfhub.dev/jeongukjae/distilkobert_cased_L-3_H-768_A-12/1", trainable=True)
    inputs = {
        "input_word_ids": tf.keras.Input([None], dtype=tf.int32, name="input_word_ids"),
        "input_mask": tf.keras.Input([None], dtype=tf.int32, name="input_mask"),
    }
    logit = encoder(inputs)['pooled_output']
    logit = tf.keras.layers.Dropout(0.1)(logit)
    pred = tf.keras.layers.Dense(3)(logit)
    model = tf.keras.Model(inputs, pred)
    model.summary()
    return model
class BertScheduler(tf.keras.optimizers.schedules.LearningRateSchedule):
    def __init__(self, rate, warmup_ratio, total_steps, name=None):
        super().__init__()

        self.rate = rate
        self.warmup_ratio = warmup_ratio
        self.total_steps = float(total_steps)
        self.warmup_steps = warmup_ratio * total_steps
        self.name = name

    def __call__(self, step):
        with tf.name_scope("BertScheduler"):
            total_steps = tf.convert_to_tensor(self.total_steps, name="total_steps")
            warmup_steps = tf.convert_to_tensor(self.warmup_steps, name="warmup_steps")

            current_step = step + 1.0

            return self.rate * tf.cond(
                current_step < warmup_steps,
                lambda: self.warmup(current_step, warmup_steps),
                lambda: self.decay(current_step, total_steps, warmup_steps),
            )

    @tf.function
    def warmup(self, step, warmup_steps):
        return step / tf.math.maximum(tf.constant(1.0), warmup_steps)

    @tf.function
    def decay(self, step, total_steps, warmup_steps):
        return tf.math.maximum(
            tf.constant(0.0), (total_steps - step) / tf.math.maximum(tf.constant(1.0), total_steps - warmup_steps)
        )<jupyter_output><empty_output><jupyter_text>Prepare datasets<jupyter_code>def get_dataset(preprocessor, batch_size):
    with tf.device('/job:localhost'):
        # batch_size=-1 is a way to load the dataset into memory
        in_memory_ds = tfds.load("kornli", batch_size=-1, shuffle_files=True)

    tfds_info = tfds.builder("kornli").info
    train_ds = tf.data.Dataset.from_tensor_slices(in_memory_ds['mnli_train'])
    dev_ds = tf.data.Dataset.from_tensor_slices(in_memory_ds['xnli_dev'])
    test_ds = tf.data.Dataset.from_tensor_slices(in_memory_ds['xnli_test'])
    num_examples = tfds_info.splits['mnli_train'].num_examples

    train_ds = (
        train_ds
        .shuffle(num_examples, reshuffle_each_iteration=True)
        .batch(batch_size, drop_remainder=True)
        .map(lambda x: (preprocessor([x['sentence1'], x['sentence2']]), x['gold_label']), num_parallel_calls=tf.data.AUTOTUNE)
    )
    dev_ds = (
        dev_ds
        .batch(batch_size)
        .map(lambda x: (preprocessor([x['sentence1'], x['sentence2']]), x['gold_label']), num_parallel_calls=tf.data.AUTOTUNE)
    )
    test_ds = (
        test_ds
        .batch(batch_size)
        .map(lambda x: (preprocessor([x['sentence1'], x['sentence2']]), x['gold_label']), num_parallel_calls=tf.data.AUTOTUNE)
    )
    return (train_ds, dev_ds, test_ds), num_examples<jupyter_output><empty_output><jupyter_text>Run train<jupyter_code>preprocessor = create_preprocessing_model()


with strategy.scope():
    (train_ds, dev_ds, test_ds), num_examples = get_dataset(preprocessor, BATCH_SIZE)
    print("Element spec:", train_ds.element_spec)
    print("Num examples:", num_examples)
    steps_per_epoch = math.ceil(num_examples / BATCH_SIZE)
    print("steps per epoch:", steps_per_epoch)
    num_train_steps = steps_per_epoch * EPOCHS
    print("total num steps:", num_train_steps)

    model = create_model()
    model.compile(
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        optimizer=tf.keras.optimizers.Adam(learning_rate=BertScheduler(LEARNING_RATE, WARMUP_RATE, num_train_steps)),
        metrics=['acc']
    )

    model.fit(
        train_ds,
        epochs=EPOCHS,
        validation_data=dev_ds,
    )
    model.evaluate(test_ds)<jupyter_output>WARNING:tensorflow:From /usr/local/lib/python3.7/dist-packages/tensorflow_datasets/core/dataset_builder.py:622: get_single_element (from tensorflow.python.data.experimental.ops.get_single_element) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.data.Dataset.get_single_element()`.<jupyter_text>Save Model<jupyter_code>save_options = tf.saved_model.SaveOptions(experimental_io_device='/job:localhost')
model.save("distilkobert_kornli", include_optimizer=False, options=save_options)<jupyter_output>WARNING:absl:Found untraced functions such as restored_function_body, restored_function_body, restored_function_body, restored_function_body, restored_function_body while saving (showing 5 of 225). These functions will not be directly callable after loading.
{'fraction_non_alphanumeric': 0.11368838699832139, 'fraction_numerical': 0.008545704257591942, 'mean_word_length': 4.394238683127572, 'pattern_counts': {'":': 2, '<': 23, '<?xml version=': 0, '>': 22, 'https://': 5, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '72656', 'license': 'Apache-2.0', 'n_tokens_mistral': 2326, 'n_tokens_neox': 2251, 'n_words': 342, 'path': 'distilkobert_kornli_tpu.ipynb', 'repo_name': 'jeongukjae/distilkobert-tfhub-examples', 'size': '6553'}