notebooks/tokun.demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "## Import deps"
      ],
      "metadata": {
        "id": "yd0W4Xjwm6vx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install mlable tokun\n",
        "\n",
        "%load_ext tensorboard"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "l5GL5MtSnAOC",
        "outputId": "1efdf330-0ad9-418a-94a1-a90f00fdb775"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting mlable\n",
            "  Downloading mlable-0.3.13-py3-none-any.whl (16 kB)\n",
            "Collecting tokun\n",
            "  Downloading tokun-0.6.1-py3-none-any.whl (14 kB)\n",
            "Requirement already satisfied: tensorflow>=2.14 in /usr/local/lib/python3.10/dist-packages (from mlable) (2.15.0)\n",
            "Requirement already satisfied: tensorflow-datasets>=4.9 in /usr/local/lib/python3.10/dist-packages (from tokun) (4.9.6)\n",
            "Requirement already satisfied: absl-py>=1.0.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (1.4.0)\n",
            "Requirement already satisfied: astunparse>=1.6.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (1.6.3)\n",
            "Requirement already satisfied: flatbuffers>=23.5.26 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (24.3.25)\n",
            "Requirement already satisfied: gast!=0.5.0,!=0.5.1,!=0.5.2,>=0.2.1 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (0.5.4)\n",
            "Requirement already satisfied: google-pasta>=0.1.1 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (0.2.0)\n",
            "Requirement already satisfied: h5py>=2.9.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (3.11.0)\n",
            "Requirement already satisfied: libclang>=13.0.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (18.1.1)\n",
            "Requirement already satisfied: ml-dtypes~=0.2.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (0.2.0)\n",
            "Requirement already satisfied: numpy<2.0.0,>=1.23.5 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (1.25.2)\n",
            "Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (3.3.0)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (24.1)\n",
            "Requirement already satisfied: protobuf!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.20.3 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (3.20.3)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (67.7.2)\n",
            "Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (1.16.0)\n",
            "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (2.4.0)\n",
            "Requirement already satisfied: typing-extensions>=3.6.6 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (4.12.2)\n",
            "Requirement already satisfied: wrapt<1.15,>=1.11.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (1.14.1)\n",
            "Requirement already satisfied: tensorflow-io-gcs-filesystem>=0.23.1 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (0.37.0)\n",
            "Requirement already satisfied: grpcio<2.0,>=1.24.3 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (1.64.1)\n",
            "Requirement already satisfied: tensorboard<2.16,>=2.15 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (2.15.2)\n",
            "Requirement already satisfied: tensorflow-estimator<2.16,>=2.15.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (2.15.0)\n",
            "Requirement already satisfied: keras<2.16,>=2.15.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow>=2.14->mlable) (2.15.0)\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.10/dist-packages (from tensorflow-datasets>=4.9->tokun) (8.1.7)\n",
            "Requirement already satisfied: dm-tree in /usr/local/lib/python3.10/dist-packages (from tensorflow-datasets>=4.9->tokun) (0.1.8)\n",
            "Requirement already satisfied: immutabledict in /usr/local/lib/python3.10/dist-packages (from tensorflow-datasets>=4.9->tokun) (4.2.0)\n",
            "Requirement already satisfied: promise in /usr/local/lib/python3.10/dist-packages (from tensorflow-datasets>=4.9->tokun) (2.3)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from tensorflow-datasets>=4.9->tokun) (5.9.5)\n",
            "Requirement already satisfied: pyarrow in /usr/local/lib/python3.10/dist-packages (from tensorflow-datasets>=4.9->tokun) (16.1.0)\n",
            "Requirement already satisfied: requests>=2.19.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow-datasets>=4.9->tokun) (2.31.0)\n",
            "Requirement already satisfied: simple-parsing in /usr/local/lib/python3.10/dist-packages (from tensorflow-datasets>=4.9->tokun) (0.1.5)\n",
            "Requirement already satisfied: tensorflow-metadata in /usr/local/lib/python3.10/dist-packages (from tensorflow-datasets>=4.9->tokun) (1.15.0)\n",
            "Requirement already satisfied: toml in /usr/local/lib/python3.10/dist-packages (from tensorflow-datasets>=4.9->tokun) (0.10.2)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from tensorflow-datasets>=4.9->tokun) (4.66.4)\n",
            "Requirement already satisfied: array-record>=0.5.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow-datasets>=4.9->tokun) (0.5.1)\n",
            "Requirement already satisfied: etils[enp,epath,epy,etree]>=1.6.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow-datasets>=4.9->tokun) (1.7.0)\n",
            "Requirement already satisfied: wheel<1.0,>=0.23.0 in /usr/local/lib/python3.10/dist-packages (from astunparse>=1.6.0->tensorflow>=2.14->mlable) (0.43.0)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from etils[enp,epath,epy,etree]>=1.6.0->tensorflow-datasets>=4.9->tokun) (2024.6.0)\n",
            "Requirement already satisfied: importlib_resources in /usr/local/lib/python3.10/dist-packages (from etils[enp,epath,epy,etree]>=1.6.0->tensorflow-datasets>=4.9->tokun) (6.4.0)\n",
            "Requirement already satisfied: zipp in /usr/local/lib/python3.10/dist-packages (from etils[enp,epath,epy,etree]>=1.6.0->tensorflow-datasets>=4.9->tokun) (3.19.2)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->tensorflow-datasets>=4.9->tokun) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->tensorflow-datasets>=4.9->tokun) (3.7)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->tensorflow-datasets>=4.9->tokun) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->tensorflow-datasets>=4.9->tokun) (2024.6.2)\n",
            "Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow>=2.14->mlable) (2.27.0)\n",
            "Requirement already satisfied: google-auth-oauthlib<2,>=0.5 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow>=2.14->mlable) (1.2.0)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow>=2.14->mlable) (3.6)\n",
            "Requirement already satisfied: tensorboard-data-server<0.8.0,>=0.7.0 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow>=2.14->mlable) (0.7.2)\n",
            "Requirement already satisfied: werkzeug>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow>=2.14->mlable) (3.0.3)\n",
            "Requirement already satisfied: docstring-parser~=0.15 in /usr/local/lib/python3.10/dist-packages (from simple-parsing->tensorflow-datasets>=4.9->tokun) (0.16)\n",
            "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow>=2.14->mlable) (5.3.3)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.10/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow>=2.14->mlable) (0.4.0)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.10/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow>=2.14->mlable) (4.9)\n",
            "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.10/dist-packages (from google-auth-oauthlib<2,>=0.5->tensorboard<2.16,>=2.15->tensorflow>=2.14->mlable) (2.0.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.1.1 in /usr/local/lib/python3.10/dist-packages (from werkzeug>=1.0.1->tensorboard<2.16,>=2.15->tensorflow>=2.14->mlable) (2.1.5)\n",
            "Requirement already satisfied: pyasn1<0.7.0,>=0.4.6 in /usr/local/lib/python3.10/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow>=2.14->mlable) (0.6.0)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.10/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<2,>=0.5->tensorboard<2.16,>=2.15->tensorflow>=2.14->mlable) (3.2.2)\n",
            "Installing collected packages: mlable, tokun\n",
            "Successfully installed mlable-0.3.13 tokun-0.6.1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import datetime\n",
        "import functools\n",
        "import itertools\n",
        "import math\n",
        "import os\n",
        "import urllib.request\n",
        "\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "import mlable.metrics\n",
        "import mlable.io\n",
        "\n",
        "import tokun.evaluation\n",
        "import tokun.meta\n",
        "import tokun.model\n",
        "import tokun.pipeline"
      ],
      "metadata": {
        "id": "yo18wu1bm6ee"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"Tensorflow version \" + tf.__version__)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sbjr53W8ots9",
        "outputId": "a05dd1b4-a520-43b8-a305-010b9a31188c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Tensorflow version 2.15.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_pQCOmISAQBu"
      },
      "source": [
        "## Setup the GPU / TPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vFIMfPmgQa0h",
        "outputId": "f85ce277-dbbb-402f-a44e-0e1c52979168"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "<tensorflow.python.distribute.tpu_strategy.TPUStrategyV2 object at 0x7c8aa5755090>\n"
          ]
        }
      ],
      "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": "0t1jfsJlM3SX"
      },
      "source": [
        "## Defining The Metadata"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# TOGGLE ######################################################################\n",
        "\n",
        "BINARY = True"
      ],
      "metadata": {
        "id": "u6fxCXsWF2w4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8Z74MlibMWnu"
      },
      "outputs": [],
      "source": [
        "# PARAMETERS ##################################################################\n",
        "\n",
        "N_SEQUENCE_AXIS = 1\n",
        "N_FEATURE_AXIS = -1\n",
        "\n",
        "N_TOKEN_DIM = [4, 16] # G, for each block\n",
        "N_INPUT_DIM = 256 # U_i (bytes)\n",
        "N_EMBED_DIM = 256 # E\n",
        "N_OUTPUT_DIM = 8 if BINARY else 256 # U_o (8 bits)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jm6y63XRBz07"
      },
      "outputs": [],
      "source": [
        "# DERIVED #####################################################################\n",
        "\n",
        "N_TOKEN_SIZES = list(itertools.accumulate(N_TOKEN_DIM, lambda x, y: x * y)) # in bytes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xfIZb86Fg0dQ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0244ac30-015a-4314-a9bd-727393b56b9b"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "('model.keras', <http.client.HTTPMessage at 0x7c897e52f9d0>)"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ],
      "source": [
        "# IMPORT ######################################################################\n",
        "\n",
        "LABEL = '8.6'\n",
        "VERSION = tokun.meta.version(token_dim=N_TOKEN_DIM, sequence_axis=N_SEQUENCE_AXIS, input_dim=N_INPUT_DIM, embed_dim=N_EMBED_DIM, output_dim=N_OUTPUT_DIM)\n",
        "\n",
        "URL_IMPORT = 'https://github.com/apehex/tokun/raw/main/models/{}/{}/{}/{}.keras'.format(*VERSION, LABEL)\n",
        "PATH_IMPORT = 'model.keras'\n",
        "\n",
        "urllib.request.urlretrieve(URL_IMPORT, PATH_IMPORT)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R-iGpXb15m83"
      },
      "outputs": [],
      "source": [
        "# SAMPLES #####################################################################\n",
        "\n",
        "SAMPLES = [\n",
        "    \"\"\"A variational autoencoder is a generative model with a prior and noise distribution respectively. Usually such models are trained using the expectation-maximization meta-algorithm (e.g. probabilistic PCA, (spike & slab) sparse coding). Such a scheme optimizes a lower bound of the data likelihood, which is usually intractable, and in doing so requires the discovery of q-distributions, or variational posteriors. These q-distributions are normally parameterized for each individual data point in a separate optimization process. However, variational autoencoders use a neural network as an amortized approach to jointly optimize across data points. This neural network takes as input the data points themselves, and outputs parameters for the variational distribution. As it maps from a known input space to the low-dimensional latent space, it is called the encoder.\"\"\",\n",
        "    \"\"\"Une unité lexicale ou token lexical ou plus simplement token est un couple composé d'un nom et d'une valeur optionnelle (e.g. 135677).\"\"\",\n",
        "    \"\"\"class AutoEncoder(tf.keras.models.Model):\\n    def __init__(self, token_dim: int, encoding_dim: int, embedding_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, batch_dim=batch_dim)\\n        self._decoder = Decoder(token_dim=token_dim, encoding_dim=encoding_dim, embedding_dim=embedding_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, 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, batch_dim=batch_dim)\\n    self._decoder = Decoder(token_dim=token_dim, encoding_dim=encoding_dim, embedding_dim=embedding_dim, batch_dim=batch_dim)\\n\\n  def call(self, x: tf.Tensor) -> tf.Tensor:\\n    return self._decoder(self._encoder(x))\"\"\",\n",
        "    \"\"\"위키백과, 우리 모두의 백과사전.\\nt-분포 확률적 임베딩(t-SNE)은 데이터의 차원 축소에 사용되는 기계 학습 알고리즘 중 하나로, 2002년 샘 로이스Sam Rowise와 제프리 힌튼에 의해 개발되었다.[1] t-SNE는 비선형 차원 축소 기법으로, 고차원 데이터를 특히 2, 3차원 등으로 줄여 가시화하는데에 유용하게 사용된다. 구체적으로 t-SNE는 비슷한 데이터는 근접한 2, 3차원의 지점으로, 다른 데이터는 멀리 떨어진 지점으로 맵핑한다.\"\"\",]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dEyFtkcFNGe4"
      },
      "source": [
        "## Init"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# METRICS #####################################################################\n",
        "\n",
        "_Accuracy = mlable.metrics.BinaryGroupAccuracy if BINARY else mlable.metrics.CategoricalGroupAccuracy\n",
        "_Loss = tf.keras.losses.BinaryCrossentropy if BINARY else tf.keras.losses.CategoricalCrossentropy"
      ],
      "metadata": {
        "id": "HOhNK7hiYaaY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iEpY1-vFIFX7"
      },
      "outputs": [],
      "source": [
        "# INIT ########################################################################\n",
        "\n",
        "with DISTRIBUTION_STRATEGY.scope():\n",
        "    # metrics\n",
        "    byte_accuracy = _Accuracy(group=1, name='byte_accuracy')\n",
        "    character_accuracy = _Accuracy(group=4, name='character_accuracy')\n",
        "    token_accuracy = _Accuracy(group=N_TOKEN_SIZES[-1], name='token_accuracy')\n",
        "    # weights\n",
        "    MODEL = tf.keras.models.load_model(PATH_IMPORT, compile=False)\n",
        "    # compilation\n",
        "    MODEL.compile(\n",
        "        optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001),\n",
        "        loss=_Loss(from_logits=False, label_smoothing=0., axis=-1, reduction='sum_over_batch_size', name='ce_loss'),\n",
        "        metrics=[byte_accuracy, character_accuracy, token_accuracy])\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EuRwWdjpPQBM",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2a477d05-b2d5-4da9-bb46-22af04b463d0"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model: \"auto_encoder\"\n",
            "_________________________________________________________________\n",
            " Layer (type)                Output Shape              Param #   \n",
            "=================================================================\n",
            " encoder (Encoder)           multiple                  1377792   \n",
            "                                                                 \n",
            " decoder (Decoder)           multiple                  1382656   \n",
            "                                                                 \n",
            "=================================================================\n",
            "Total params: 2760448 (10.53 MB)\n",
            "Trainable params: 2760448 (10.53 MB)\n",
            "Non-trainable params: 0 (0.00 Byte)\n",
            "_________________________________________________________________\n"
          ]
        }
      ],
      "source": [
        "MODEL.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FHtROW1K1R7c"
      },
      "source": [
        "## Tokenize"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# ENCODE ######################################################################\n",
        "\n",
        "__s = SAMPLES[1]\n",
        "__x = tokun.pipeline.preprocess(text=__s, token_size=math.prod(N_TOKEN_DIM), expand=N_SEQUENCE_AXIS * [1])\n",
        "__e = MODEL._encoder(__x) # final embedding = input for another model"
      ],
      "metadata": {
        "id": "gw_bhpP6JX8-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# VIEW ########################################################################\n",
        "\n",
        "print(len(__s)) # original text length, roughly S\n",
        "print(__x.shape) # 4 * S, with padding\n",
        "print(__e.shape) # 4 * S // 64"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "M5DINgnlMCwC",
        "outputId": "248d819a-d2d0-4cce-cce8-a1103e476db6"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "134\n",
            "(1, 576, 256)\n",
            "(1, 9, 256)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Detokenize"
      ],
      "metadata": {
        "id": "maFagKtSLgPv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# DECODE ######################################################################\n",
        "\n",
        "__p = MODEL._decoder(__e)\n",
        "__y = tokun.pipeline.postprocess(__p, binary=BINARY, random=False)\n",
        "__o = tokun.pipeline.unpack(__y)"
      ],
      "metadata": {
        "id": "cLuBORzbMNcq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print('# INPUT ################################################################\\n\\n' + __s)\n",
        "print('\\n# OUTPUT ###############################################################\\n\\n' + __o[0])\n",
        "print('\\n# SCORE ################################################################\\n\\n' + str(tokun.evaluation.compare(__s, __o[0])))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mZkdNGWoMalz",
        "outputId": "3a35bfe9-a7f9-45fb-9165-932e4412e19b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "# INPUT ################################################################\n",
            "\n",
            "Une unité lexicale ou token lexical ou plus simplement token est un couple composé d'un nom et d'une valeur optionnelle (e.g. 135677).\n",
            "\n",
            "# OUTPUT ###############################################################\n",
            "\n",
            "Une unité lexicale ou token lexical ou plus simplement token est un couple composé d'un nom et d'une valeur optionnelle (e.g. 135677).\n",
            "\n",
            "# SCORE ################################################################\n",
            "\n",
            "1.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Robustness"
      ],
      "metadata": {
        "id": "dLjTuH3IMyza"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# SAMPLE ######################################################################\n",
        "\n",
        "__x = tokun.pipeline.preprocess(\"\"\"Une unité lexicale ou token lexical ou plus simplement token est un couple composé d'un nom et d'une valeur optionnelle (e.g. 135677).\"\"\", token_size=math.prod(N_TOKEN_DIM), expand=N_SEQUENCE_AXIS * [1])\n",
        "__e = MODEL._encoder(__x)"
      ],
      "metadata": {
        "id": "-NRq3hZHM_6D"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7ysd903RYTt8"
      },
      "outputs": [],
      "source": [
        "# NOISE #######################################################################\n",
        "\n",
        "__std = tf.math.reduce_std(__e, axis=1)\n",
        "__noise = tf.random.normal(shape=(256,), mean=0., stddev=tf.math.reduce_mean(__std).numpy())"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# DETOKENIZE ##################################################################\n",
        "\n",
        "print(tokun.pipeline.postprocess(MODEL._decoder(__e), binary=BINARY, random=False)) # original\n",
        "print(tokun.pipeline.postprocess(MODEL._decoder(__e + 0.8 * __std), binary=BINARY, random=False)) # with structured noise\n",
        "print(tokun.pipeline.postprocess(MODEL._decoder(__e + 1. * __noise), binary=BINARY, random=False)) # with random noise"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nKyqASAJs2xl",
        "outputId": "d2a0805d-700f-41b3-819c-b13c6334d36b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Une unité lexicale ou token lexical ou plus simplement token est un couple composé d'un nom et d'une valeur optionnelle (e.g. 135677).\n",
            "Une unité lexicale ou token lexical ou plus simplement token est un couple composé d'un nom et d'une valeur optionnelle (e.g. 135677).\n",
            "Une unité lexicale ou token lexical ou plus simplement token est un couple composé d'un nom et d'une valeur optionnelle (e.g. 135677).\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Text Similarity"
      ],
      "metadata": {
        "id": "y7yyDMCmmbYD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# SAMPLES #####################################################################\n",
        "\n",
        "REFERENCE = \"\"\"A couple is enjoying fresh sushis at home.\"\"\"\n",
        "CANDIDATES = [\n",
        "    \"\"\"A couple is cooking fresh sushis at home.\"\"\",\n",
        "    \"\"\"Ana and Bob are enjoying sushis while watching TV.\"\"\",\n",
        "    \"\"\"A couple is enjoying california rolls.\"\"\",\n",
        "    \"\"\"You can order fresh sushis from home.\"\"\",\n",
        "    \"\"\"Enjoying fresh sushis at home a couple is.\"\"\",\n",
        "    \"\"\"Sushis are the best!\"\"\",]"
      ],
      "metadata": {
        "id": "khdOpTtlnLqe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# METRIC ######################################################################\n",
        "\n",
        "__f = tf.keras.losses.CosineSimilarity(axis=-1, reduction=tf.keras.losses.Reduction.SUM_OVER_BATCH_SIZE)\n",
        "\n",
        "# REFERENCE ###################################################################\n",
        "\n",
        "__xr, __er, __pr, __yr = tokun.pipeline.sample(model=MODEL, text=REFERENCE, token_size=math.prod(N_TOKEN_DIM), expand=N_SEQUENCE_AXIS * [1])\n",
        "\n",
        "# SHAPE #######################################################################\n",
        "\n",
        "__sr = list(__er.shape) # shape\n",
        "__rr = len(__sr) # rank"
      ],
      "metadata": {
        "id": "bpKV8Qvhwtvy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# DISTANCES ###################################################################\n",
        "\n",
        "DISTANCES = {}\n",
        "\n",
        "for __c in CANDIDATES:\n",
        "  __xl, __el, __pl, __yl, __ol = tokun.pipeline.sample(model=MODEL, text=__c, token_size=math.prod(__sr[1:2] + N_TOKEN_DIM), expand=N_SEQUENCE_AXIS * [1], binary=BINARY, random=False)\n",
        "  __el = tf.slice(__el, begin=__rr * [0], size=__sr) # match reference shape\n",
        "  DISTANCES[__c] = __f(y_true=__er, y_pred=__el).numpy()"
      ],
      "metadata": {
        "id": "SNHPvydHlsGW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# SORT ########################################################################\n",
        "\n",
        "DISTANCES = dict(sorted(DISTANCES.items(), key=lambda __x: __x[1]))"
      ],
      "metadata": {
        "id": "Vj90AWFUy2no"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "DISTANCES"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sKTTjT5KxDUK",
        "outputId": "c501eaa3-f581-4c1c-8953-e7fe66dad173"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'A couple is enjoying california rolls.': -0.9971239,\n",
              " 'A couple is cooking fresh sushis at home.': -0.9952307,\n",
              " 'Sushis are the best!': -0.9944132,\n",
              " 'Enjoying fresh sushis at home a couple is.': -0.9937381,\n",
              " 'You can order fresh sushis from home.': -0.993692,\n",
              " 'Ana and Bob are enjoying sushis while watching TV.': -0.9927745}"
            ]
          },
          "metadata": {},
          "execution_count": 79
        }
      ]
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "gpuType": "V28",
      "machine_shape": "hm",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}