notebooks/llaminate.student.pretrain.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3xXM7DoPpds1"
      },
      "source": [
        "## Import deps"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W24EKFXaO5yC"
      },
      "outputs": [],
      "source": [
        "!pip install -U datasets mlable tokun llaminate"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FdR0Yr-S3RqK"
      },
      "outputs": [],
      "source": [
        "!huggingface-cli login\n",
        "\n",
        "%load_ext tensorboard"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VXU-Ebl2pddk"
      },
      "outputs": [],
      "source": [
        "import datetime\n",
        "import functools\n",
        "import itertools\n",
        "import math\n",
        "import os\n",
        "import random\n",
        "import urllib.request\n",
        "\n",
        "import datasets as hd\n",
        "import tensorflow as tf\n",
        "\n",
        "import mlable.data\n",
        "import mlable.metrics\n",
        "\n",
        "import tokun.data\n",
        "import tokun.evaluation\n",
        "import tokun.meta\n",
        "import tokun.model\n",
        "import tokun.pipeline\n",
        "\n",
        "import llaminate.model\n",
        "import llaminate.pipeline\n",
        "import llaminate.utils"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Pn1ywhSrpin9"
      },
      "outputs": [],
      "source": [
        "print(\"Tensorflow version \" + tf.__version__)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_pQCOmISAQBu"
      },
      "source": [
        "## Setup the GPU / TPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n_LfBoYAQa4d"
      },
      "outputs": [],
      "source": [
        "# MIXED PRECISION #############################################################\n",
        "\n",
        "tf.keras.mixed_precision.set_global_policy('mixed_bfloat16')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vFIMfPmgQa0h"
      },
      "outputs": [],
      "source": [
        "# DEVICES #####################################################################\n",
        "\n",
        "tf.debugging.set_log_device_placement(False)\n",
        "\n",
        "CPU = tf.config.list_logical_devices('CPU')\n",
        "GPU = tf.config.list_logical_devices('GPU')\n",
        "TPU = tf.config.list_logical_devices('TPU')\n",
        "\n",
        "if TPU:\n",
        "    RESOLVER = tf.distribute.cluster_resolver.TPUClusterResolver()\n",
        "    tf.config.experimental_connect_to_cluster(RESOLVER)\n",
        "    tf.tpu.experimental.initialize_tpu_system(RESOLVER)\n",
        "    DISTRIBUTION_STRATEGY = tf.distribute.TPUStrategy(RESOLVER)\n",
        "elif GPU:\n",
        "    DISTRIBUTION_STRATEGY = tf.distribute.MirroredStrategy(GPU)\n",
        "else:\n",
        "    DISTRIBUTION_STRATEGY = tf.distribute.MirroredStrategy(CPU)\n",
        "\n",
        "print(DISTRIBUTION_STRATEGY)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f9066X5EOyAX"
      },
      "source": [
        "## Mode"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lFSPMtQaO1fu"
      },
      "outputs": [],
      "source": [
        "# TOGGLE ######################################################################\n",
        "\n",
        "IMPORT = False\n",
        "DOWNLOAD = False\n",
        "TRAINING = True\n",
        "BINARY = True"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0t1jfsJlM3SX"
      },
      "source": [
        "## Defining The Metadata"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_jbqmcLOg1-5"
      },
      "outputs": [],
      "source": [
        "# PIPELINE ####################################################################\n",
        "\n",
        "BATCH_CONFIG = {\n",
        "    'batch_size': 32,\n",
        "    'drop_remainder': True,\n",
        "    'num_parallel_calls': tf.data.AUTOTUNE,}\n",
        "\n",
        "PIPELINE_CONFIG = {\n",
        "    'batch_dim': BATCH_CONFIG['batch_size'],\n",
        "    'sample_dim': 2048 * 8,\n",
        "    'input_dim': 8,\n",
        "    'data_weight': 1.0,\n",
        "    'padding_weight': 0.0001,\n",
        "    'separator': '\\x1d',}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8Z74MlibMWnu"
      },
      "outputs": [],
      "source": [
        "# MODEL PARAMETERS ############################################################\n",
        "\n",
        "LLAMINATE_CONFIG = {\n",
        "  'num_layers': 12,\n",
        "  'num_heads': 16,\n",
        "  'input_dim': PIPELINE_CONFIG['input_dim'],\n",
        "  'embed_dim': 1024,\n",
        "  'head_dim': 1024 // 16,\n",
        "  'hidden_dim': 1024 * 4,\n",
        "  'epsilon': 1e-6,}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "anEMVSIJiuLY"
      },
      "outputs": [],
      "source": [
        "# DERIVED PARAMETERS ##########################################################\n",
        "\n",
        "LLAMINATE_META = {\n",
        "    'version': '{}x{}x{}'.format(LLAMINATE_CONFIG['num_layers'], LLAMINATE_CONFIG['input_dim'], LLAMINATE_CONFIG['embed_dim']),\n",
        "    'path': 'llaminate.keras',\n",
        "    'url': '',}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P2bgdk4P02n8"
      },
      "outputs": [],
      "source": [
        "# TRAINING PARAMETERS #########################################################\n",
        "\n",
        "OPTIMIZER_CONFIG = {\n",
        "    'learning_rate': 0.001 * (0.1 if IMPORT else 1.0),\n",
        "    'weight_decay': 0.01,\n",
        "    'beta_1': 0.9,\n",
        "    'beta_2': 0.95,\n",
        "    'clipnorm': 1.0,}\n",
        "\n",
        "SCHEDULER_CONFIG = {\n",
        "    'initial_learning_rate': OPTIMIZER_CONFIG['learning_rate'],\n",
        "    'decay_steps': 16384 * 8,\n",
        "    'alpha': 0.01,\n",
        "    'name': 'cosine_lr',\n",
        "    'warmup_target': None,\n",
        "    'warmup_steps': 0,}\n",
        "\n",
        "METRICS_CONFIG = {\n",
        "    'depth': 8,\n",
        "    'threshold': 0.6,}\n",
        "\n",
        "LOSS_CONFIG = {\n",
        "    'from_logits': False,\n",
        "    'label_smoothing': 0.,\n",
        "    'axis': -1,\n",
        "    'reduction': 'sum_over_batch_size',\n",
        "    'name': 'loss',}\n",
        "\n",
        "CHECKPOINT_CONFIG = {\n",
        "    'filepath': LLAMINATE_META['path'],\n",
        "    'monitor': 'val_loss',\n",
        "    'mode': 'auto',\n",
        "    'save_freq': 'epoch',\n",
        "    'save_best_only': False,\n",
        "    'save_weights_only': False,\n",
        "    'verbose': 1,}\n",
        "\n",
        "TENSORBOARD_CONFIG = {\n",
        "    'log_dir': os.path.join('.logs/', *LLAMINATE_META['version'], datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\")),\n",
        "    'histogram_freq': 1,\n",
        "    'embeddings_freq': 0,\n",
        "    'profile_batch': (128, 256),\n",
        "    'write_graph': False,\n",
        "    'write_images': True,}\n",
        "\n",
        "TRAINING_CONFIG = {\n",
        "    'epochs': 8,\n",
        "    'batch_size': None,\n",
        "    'validation_split': None,\n",
        "    'validation_freq': list(range(1, 9)),\n",
        "    # 'class_weight': {__c: 1. if __c == 0 else 1. for __c in range(256)}, # there are 3 times more 0s than other bytes\n",
        "    'verbose': 1,}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tWxgap423z1a"
      },
      "outputs": [],
      "source": [
        "# DATASETS ####################################################################\n",
        "\n",
        "# TODO bigcode/the-stack\n",
        "# TODO ArmelR/stack-exchange-instruction\n",
        "\n",
        "DATASETS_CONFIG = {\n",
        "    'pt-fineweb-edu': {\n",
        "        'path': 'HuggingFaceFW/fineweb-edu',\n",
        "        'name': 'sample-10BT',\n",
        "        'splits': [f'train[{__p}%:{__p + 10}%]' for __p in range(0, 100, 10)],\n",
        "        'features': ['text'],},\n",
        "    # 'pt-wikipedia': {\n",
        "    #     'path': 'wikimedia/wikipedia',\n",
        "    #     'name': '20231101.en',\n",
        "    #     'splits': [f'train[{__p}%:{__p + 9}%]' for __p in range(0, 80, 8)],\n",
        "    #     'features': ['text'],},\n",
        "    # 'tp-wikipedia-1': {\n",
        "    #     'path': 'wikimedia/wikipedia',\n",
        "    #     'name': '20231101.en',\n",
        "    #     'splits': [f'train[{__p}%:{__p + 1}%]' for __p in range(80, 90, 1)],\n",
        "    #     'features': ['text'],},\n",
        "    # 'tp-wikipedia-2': {\n",
        "    #     'path': 'wikimedia/wikipedia',\n",
        "    #     'name': '20231101.en',\n",
        "    #     'splits': [f'train[{__p}%:{__p + 1}%]' for __p in range(90, 100, 1)],\n",
        "    #     'features': ['text'],},\n",
        "    # 'ft-retro-ascii-art': {\n",
        "    #     'path': 'jdpressman/retro-ascii-art-v1',\n",
        "    #     'name': None,\n",
        "    #     'train': 'train',\n",
        "    #     'splits': [f'train[{__p}%:{__p + 10}%]+validation[{__p}%:{__p + 10}%]' for __p in range(0, 100, 10)],\n",
        "    #     'features': ['prompt', 'art_aic'],},\n",
        "    # 'ft-stack-exchange': {\n",
        "    #     'path': 'Alignment-Lab-AI/Stack-Exchange-April',\n",
        "    #     'name': None,\n",
        "    #     'splits': [f'train[{__p}%:{__p + 10}%]' for __p in range(0, 100, 10)],\n",
        "    #     'features': ['question', 'answer'],},\n",
        "    'ft-math': {\n",
        "        'path': 'hendrycks/competition_math',\n",
        "        'name': None,\n",
        "        'splits': [f'train[{__p}%:{__p + 10}%]+test[{__p}%:{__p + 10}%]' for __p in range(0, 100, 10)],\n",
        "        'features': ['problem', 'solution'],},}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AzxAYGC9ZmOH"
      },
      "source": [
        "## Loading The Weights"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jm6y63XRBz07"
      },
      "outputs": [],
      "source": [
        "# DERIVED PARAMETERS ##########################################################\n",
        "\n",
        "if IMPORT and DOWNLOAD:\n",
        "    urllib.request.urlretrieve(LLAMINATE_META['url'], LLAMINATE_META['path'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dEyFtkcFNGe4"
      },
      "source": [
        "## Loading The Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "39aImJwK68wr"
      },
      "outputs": [],
      "source": [
        "# DOWNLOAD ####################################################################\n",
        "\n",
        "DATASETS = {\n",
        "    __name: [\n",
        "        hd.load_dataset(path=__args['path'], name=__args['name'], split=__s).to_tf_dataset(shuffle=True, batch_size=None)\n",
        "        for __s in __args['splits']]\n",
        "    for __name, __args in DATASETS_CONFIG.items()}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ovWpdtkifbgg"
      },
      "source": [
        "## Checking The Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "10Ret-lA4Emo"
      },
      "outputs": [],
      "source": [
        "# STATS #######################################################################\n",
        "\n",
        "STATS = {__n: mlable.data.stats(dataset=DATASETS[__n][0], features=DATASETS_CONFIG[__n]['features'], count=2048) for __n in DATASETS}\n",
        "\n",
        "print(STATS)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5__iAHh41CZH"
      },
      "outputs": [],
      "source": [
        "__b = iter(DATASETS['pt-fineweb-edu'][0])\n",
        "next(__b)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8cheN52OEchs"
      },
      "source": [
        "## Preprocess"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hCMihca3F2yQ"
      },
      "outputs": [],
      "source": [
        "# ITERATE #####################################################################\n",
        "\n",
        "with DISTRIBUTION_STRATEGY.scope():\n",
        "    for __name in DATASETS:\n",
        "        # specialized preprocessing fn\n",
        "        __preprocess = llaminate.pipeline.preprocess_factory(\n",
        "            features=DATASETS_CONFIG[__name]['features'],\n",
        "            **PIPELINE_CONFIG)\n",
        "        # apply\n",
        "        for __idx in range(len(DATASETS[__name])):\n",
        "            DATASETS[__name][__idx] = DATASETS[__name][__idx].batch(**BATCH_CONFIG).map(__preprocess, num_parallel_calls=tf.data.AUTOTUNE)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qxz2uQx-g5SR"
      },
      "outputs": [],
      "source": [
        "# CONCATENATE #################################################################\n",
        "\n",
        "FINE_TRAIN = functools.reduce(lambda __l, __r: __l.concatenate(__r), DATASETS['pt-fineweb-edu'][:-1])\n",
        "FINE_TEST = DATASETS['pt-fineweb-edu'][-1]\n",
        "\n",
        "DATASET_TRAIN = functools.reduce(lambda __l, __r: __l.concatenate(__r), [DATASETS[__n][__i] for __n in (set(DATASETS.keys()) - {'ft-retro-ascii-art'}) for __i in range(len(DATASETS[__n]) - 1)]) # - {'pt-wikipedia'}\n",
        "DATASET_TEST = functools.reduce(lambda __l, __r: __l.concatenate(__r), [DATASETS[__n][-1] for __n in (set(DATASETS.keys()) - {'ft-retro-ascii-art'})]) # - {'pt-wikipedia'}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nJFgVaASv0bY"
      },
      "outputs": [],
      "source": [
        "# CHECK DATASET ###############################################################\n",
        "\n",
        "__X, __T, __W = next(iter(FINE_TRAIN.take(1)))\n",
        "\n",
        "print(FINE_TRAIN.element_spec)\n",
        "print(FINE_TEST.element_spec)\n",
        "\n",
        "print(DATASET_TRAIN.element_spec)\n",
        "print(DATASET_TEST.element_spec)\n",
        "\n",
        "print('fine: {:,} / {:,} samples'.format(FINE_TRAIN.cardinality().numpy(), FINE_TEST.cardinality().numpy()))\n",
        "print('total: {:,} / {:,} samples'.format(DATASET_TRAIN.cardinality().numpy(), DATASET_TEST.cardinality().numpy()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S39n2JmXG6yv"
      },
      "source": [
        "## Initializing The Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EuDcDORB2Gzl"
      },
      "outputs": [],
      "source": [
        "# METRICS #####################################################################\n",
        "\n",
        "_Accuracy = mlable.metrics.BinaryGroupAccuracy if BINARY else mlable.metrics.RawGroupAccuracy\n",
        "_Loss = tf.keras.losses.BinaryCrossentropy if BINARY else tf.keras.losses.MeanSquaredError"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AzD3vE7ZlB-Z"
      },
      "outputs": [],
      "source": [
        "# OVERALL SCOPE ###############################################################\n",
        "\n",
        "with DISTRIBUTION_STRATEGY.scope():\n",
        "    # COSINE LR ###############################################################\n",
        "    cosine_lr = tf.keras.optimizers.schedules.CosineDecay(**SCHEDULER_CONFIG)\n",
        "    OPTIMIZER_CONFIG['learning_rate'] = cosine_lr\n",
        "\n",
        "    # METRICS #################################################################\n",
        "    byte_accuracy = _Accuracy(group=1, name='byte_accuracy', **METRICS_CONFIG)\n",
        "    character_accuracy = _Accuracy(group=4, name='character_accuracy', **METRICS_CONFIG)\n",
        "    token_accuracy = _Accuracy(group=PIPELINE_CONFIG['input_dim'], name='token_accuracy', **METRICS_CONFIG)\n",
        "\n",
        "    # WEIGHTS #################################################################\n",
        "    LLAMINATE = llaminate.model.Transformer(**LLAMINATE_CONFIG)\n",
        "    if IMPORT and os.path.isfile(LLAMINATE_META['path']): LLAMINATE = tf.keras.models.load_model(LLAMINATE_META['path'], compile=False)\n",
        "\n",
        "    # BUILD ###################################################################\n",
        "    LLAMINATE(__X)\n",
        "\n",
        "    # COMPILE #################################################################\n",
        "    LLAMINATE.compile(\n",
        "        optimizer=tf.keras.optimizers.AdamW(**OPTIMIZER_CONFIG),\n",
        "        loss=_Loss(**LOSS_CONFIG),\n",
        "        weighted_metrics=[byte_accuracy, character_accuracy, token_accuracy])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Uz1lRnWtMSUm"
      },
      "outputs": [],
      "source": [
        "LLAMINATE.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jRkNkXthBwar"
      },
      "source": [
        "## Train"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "beTpALmzFdu1"
      },
      "outputs": [],
      "source": [
        "# TRAIN #######################################################################\n",
        "\n",
        "if TRAINING:\n",
        "    with DISTRIBUTION_STRATEGY.scope():\n",
        "        # callbacks\n",
        "        cp_callback = tf.keras.callbacks.ModelCheckpoint(**CHECKPOINT_CONFIG)\n",
        "        tb_callback = tf.keras.callbacks.TensorBoard(**TENSORBOARD_CONFIG)\n",
        "        # model fitting\n",
        "        TRAINING_HISTORY = LLAMINATE.fit(\n",
        "            x=DATASETS['pt-fineweb-edu'][0].prefetch(tf.data.AUTOTUNE),\n",
        "            validation_data=DATASETS['pt-fineweb-edu'][-1].take(128).prefetch(tf.data.AUTOTUNE),\n",
        "            callbacks=[cp_callback, tb_callback],\n",
        "            **TRAINING_CONFIG)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FHtROW1K1R7c"
      },
      "source": [
        "## Dataviz"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EBwqrcl7naB_"
      },
      "outputs": [],
      "source": [
        "__i = iter(DATASETS['pt-fineweb-edu'][-1])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UHPxwBWYnjpz"
      },
      "outputs": [],
      "source": [
        "__x, __t, __w = next(__i)\n",
        "__y = LLAMINATE(__x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z8QX4X4Dqdl4"
      },
      "outputs": [],
      "source": [
        "__s = llaminate.pipeline.postprocess(__y)\n",
        "__s[:4]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dvWo5zt8oK2V"
      },
      "outputs": [],
      "source": [
        "__s = llaminate.pipeline.postprocess(__t)\n",
        "__s[:4]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EiRVgG-oSfb7"
      },
      "outputs": [],
      "source": [
        "# DATA ########################################################################\n",
        "\n",
        "SAMPLES = [\n",
        "    \"\"\"위키백과, 우리 모두의 백과사전.\\nt-분포 확률적 임베딩(t-SNE)은 데이터의 차원 축소에 사용되는 기계 학습 알고리즘 중 하나로, 2002년 샘 로이스Sam Rowise와 제프리 힌튼에 의해 개발되었다.[1] t-SNE는 비선형 차원 축소 기법으로, 고차원 데이터를 특히 2, 3차원 등으로 줄여 가시화하는데에 유용하게 사용된다. 구체적으로 t-SNE는 비슷한 데이터는 근접한 2, 3차원의 지점으로, 다른 데이터는 멀리 떨어진 지점으로 맵핑한다.\"\"\",\n",
        "    \"\"\"class Encoder(tf.keras.models.Model):\\n    def __init__(self, depth: int, token_dim: int, encoding_dim: int, embedding_dim: int, latent_dim: int, batch_dim: int=None, attention: bool=False, **kwargs) -> None:\\n        super(Encoder, self).__init__(**kwargs)\\n        self._encoder = tf.keras.Sequential([\\n            tf.keras.Input(shape=(encoding_dim,), batch_size=batch_dim, name='input'), # (B * G ^ D, U)\\n            tf.keras.layers.Dense(units=embedding_dim, activation=None, use_bias=False, kernel_initializer='glorot_uniform', bias_initializer=None, name='embed-1'),] # (B * G ^ D, U) => (B * G ^ D, E)\\n            + [tokun.layers.TokenizeBlock(left_axis=-2, right_axis=-1, token_dim=token_dim, latent_dim=latent_dim, attention=attention, name='tokenize' + (__i + 1) * '-4') for __i in range(depth)]) # (B * G ^ i, E) => (B * G ^ (i-1), E)\\n\\n    def call(self, x: tf.Tensor) -> tf.Tensor:\\n        return self._encoder(x)\\n\"\"\",\n",
        "    \"\"\"class AutoEncoder(tf.keras.models.Model):\\n    def __init__(self, token_dim: int, encoding_dim: int, embedding_dim: int, latent_dim: int, batch_dim: int=None, **kwargs) -> None:\\n        super(AutoEncoder, self).__init__(**kwargs)\\n        self._encoder = Encoder(token_dim=token_dim, encoding_dim=encoding_dim, embedding_dim=embedding_dim, latent_dim=latent_dim, batch_dim=batch_dim)\\n        self._decoder = Decoder(token_dim=token_dim, encoding_dim=encoding_dim, embedding_dim=embedding_dim, latent_dim=latent_dim, batch_dim=batch_dim)\\n\\n    def call(self, x: tf.Tensor) -> tf.Tensor:\\n        return self._decoder(self._encoder(x))\"\"\",\n",
        "    \"\"\"class AutoEncoder(tf.keras.models.Model):\\n  def __init__(self, token_dim: int, encoding_dim: int, embedding_dim: int, latent_dim: int, batch_dim: int=None, **kwargs) -> None:\\n    super(AutoEncoder, self).__init__(**kwargs)\\n    self._encoder = Encoder(token_dim=token_dim, encoding_dim=encoding_dim, embedding_dim=embedding_dim, latent_dim=latent_dim, batch_dim=batch_dim)\\n    self._decoder = Decoder(token_dim=token_dim, encoding_dim=encoding_dim, embedding_dim=embedding_dim, latent_dim=latent_dim, batch_dim=batch_dim)\\n\\n  def call(self, x: tf.Tensor) -> tf.Tensor:\\n    return self._decoder(self._encoder(x))\"\"\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8naQsR7i8iea"
      },
      "outputs": [],
      "source": [
        "# CACHE #######################################################################\n",
        "\n",
        "__cache = llaminate.utils.create_cache(batch_dim=N_BATCH_DIM, cache_dim=N_CACHE_DIM, head_dim=N_HEAD_DIM, num_layers=N_LAYERS_NUM, num_heads=N_HEADS_NUM)\n",
        "__step = 4"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HubWVf_L_Frh"
      },
      "outputs": [],
      "source": [
        "# PREPROCESS ##################################################################\n",
        "\n",
        "__prompt = \"\"\"Skynet is an artificial neural network-based conscious group mind and artificial general superintelligence system that serves as the antagonistic force of the Terminator franchise.\"\"\"\n",
        "__inputs = tokun.pipeline.preprocess(text=__prompt, token_size=PIPELINE_CONFIG['sample_dim'], expand_dims=[1], output_dtype=tf.uint8)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Gcf9vBg_KKza"
      },
      "outputs": [],
      "source": [
        "__inputs = mlable.shaping.divide(__inputs, input_axis=-2, output_axis=-1, factor=PIPELINE_CONFIG['input_dim'], insert=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Pt7OdZw22vGM"
      },
      "outputs": [],
      "source": [
        "# PREDICT #####################################################################\n",
        "\n",
        "__predictions = LLAMINATE(__inputs)\n",
        "__outputs = llaminate.pipeline.postprocess(__predictions)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i6HogAyPLGjM"
      },
      "outputs": [],
      "source": [
        "tokun.pipeline.chunk(__prompt, size=4)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1izEKX5VKuaR"
      },
      "outputs": [],
      "source": [
        "__outputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tm0qRLA9r7zL"
      },
      "outputs": [],
      "source": [
        "__batch = iter(DATASETS['pt-wikipedia'][1])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GzM21G9YsIT7"
      },
      "outputs": [],
      "source": [
        "__x, __y, __m = next(__batch)\n",
        "__p = LLAMINATE(inputs=__x, training=True, mask=None)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "00AG38z5d_VW"
      },
      "outputs": [],
      "source": [
        "__yt = tokun.pipeline.interpret(__y, binary=True)\n",
        "__yp = tokun.pipeline.interpret(__p, binary=True)\n",
        "__it = tokun.pipeline.decode(__x)\n",
        "__ot = tokun.pipeline.decode(__yt)\n",
        "__op = tokun.pipeline.decode(__yp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zkpSSNd8mhRz"
      },
      "outputs": [],
      "source": [
        "print(__it[:2])\n",
        "print(__ot[:2])\n",
        "print(__op[:2])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kgt0SQDbaopC"
      },
      "source": [
        "## Logs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "eJmv4xjnTH4t"
      },
      "outputs": [],
      "source": [
        "%tensorboard --logdir .logs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5dviQuRo_Ebb"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "gpuType": "V28",
      "machine_shape": "hm",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}