trax/tf_numpy_and_keras.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7yuytuIllsv1"
      },
      "source": [
        "# Using Trax with TensorFlow NumPy and Keras\n",
        "\n",
        "This notebook ([run it in colab](https://colab.research.google.com/github/google/trax/blob/master/trax/tf_numpy_and_keras.ipynb)) shows how you can run [Trax](https://trax-ml.readthedocs.io/en/latest/) directly with [TensorFlow NumPy](https://www.tensorflow.org/api_docs/python/tf/experimental/numpy). You will also see how to use Trax layers and models inside [Keras](https://keras.io/) so you can use Trax in production, e.g., with [TensorFlow.js](https://www.tensorflow.org/js/) or [TensorFlow Serving](https://www.tensorflow.org/tfx/guide/serving).\n",
        "\n",
        "  1. **Trax with TensorFlow NumPy**: use Trax with [TensorFlow NumPy](https://www.tensorflow.org/api_docs/python/tf/experimental/numpy) without any code changes\n",
        "  1. **Convert Trax to Keras**: how to get a [Keras](https://keras.io/) layer for your Trax model and use it\n",
        "  1. **Exporting Trax Models for Deployment**: how to export Trax models to [TensorFlow SavedModel](https://www.tensorflow.org/guide/saved_model)\n",
        "  \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-LQ89rFFsEdk"
      },
      "source": [
        "## 1. Trax with TensorFlow NumPy\n",
        "\n",
        "In Trax, all computations rely on accelerated math operations happening in the `fastmath` module. This module can use different backends for acceleration. One of them is [TensorFlow NumPy](https://www.tensorflow.org/api_docs/python/tf/experimental/numpy) which uses [TensorFlow 2](https://www.tensorflow.org/) to accelerate the computations.\n",
        "\n",
        "The backend can be set using a call to `trax.fastmath.set_backend` as you'll see below. Currently available backends are `jax` (default), `tensorflow-numpy` and `numpy` (for debugging). The `tensorflow-numpy` backend uses [TensorFlow Numpy](https://www.tensorflow.org/api_docs/python/tf/experimental/numpy) for executing `fastmath` functions on TensorFlow, while the `jax` backend calls [JAX](https://github.com/google/jax) which lowers to TensorFlow XLA.\n",
        "\n",
        "You may see that `tensorflow-numpy` and `jax` backends show different speed and memory characteristics. You may also see different error messages when debugging since it might expose you to the internals of the backends. However for the most part, users can choose a backend and not worry about the internal details of these backends.\n",
        "\n",
        "Let's train the sentiment analysis model from the [Trax intro](https://colab.research.google.com/github/google/trax/blob/master/trax/intro.ipynb) using TensorFlow NumPy to see how it works."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BIl27504La0G"
      },
      "source": [
        "**General Setup**\n",
        "\n",
        "Execute the following few cells (once) before running any of the code samples."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "cellView": "form",
        "executionInfo": {
          "elapsed": 38104,
          "status": "ok",
          "timestamp": 1607390269924,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "oILRLCWN_16u"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "# Copyright 2020 Google LLC.\n",
        "\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "cellView": "both",
        "executionInfo": {
          "elapsed": 309,
          "status": "ok",
          "timestamp": 1607390270242,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "vlGjGoGMTt-D",
        "outputId": "279a980e-1e71-4080-9587-d89aeb17ebc6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "# Install and import Trax\n",
        "!pip install -q -U git+https://github.com/google/trax@master\n",
        "\n",
        "import os\n",
        "import numpy as np\n",
        "import trax"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O_3JcfZaT5oP"
      },
      "source": [
        "Here is how you can set the fastmath backend to `tensorflow-numpy` and verify that it's been set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "executionInfo": {
          "elapsed": 286,
          "status": "ok",
          "timestamp": 1607390270535,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "djTiSLcaNFGa",
        "outputId": "bac38e28-d1e5-41bd-9054-d85913fc2900"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensorflow-numpy\n"
          ]
        }
      ],
      "source": [
        "# Use the tensorflow-numpy backend.\n",
        "trax.fastmath.set_backend('tensorflow-numpy')\n",
        "print(trax.fastmath.backend_name())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "executionInfo": {
          "elapsed": 15126,
          "status": "ok",
          "timestamp": 1607390285667,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "AV5wrgjZ10yU",
        "outputId": "6385fbe2-5a8e-415c-8851-b5bef099e02f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "batch shapes = [(8, 2048), (8,), (8,)]\n"
          ]
        }
      ],
      "source": [
        "# Create data streams.\n",
        "train_stream = trax.data.TFDS('imdb_reviews', keys=('text', 'label'), train=True)()\n",
        "eval_stream = trax.data.TFDS('imdb_reviews', keys=('text', 'label'), train=False)()\n",
        "\n",
        "data_pipeline = trax.data.Serial(\n",
        "    trax.data.Tokenize(vocab_file='en_8k.subword', keys=[0]),\n",
        "    trax.data.Shuffle(),\n",
        "    trax.data.FilterByLength(max_length=2048, length_keys=[0]),\n",
        "    trax.data.BucketByLength(boundaries=[  32, 128, 512, 2048],\n",
        "                             batch_sizes=[512, 128,  32,    8, 1],\n",
        "                             length_keys=[0]),\n",
        "    trax.data.AddLossWeights()\n",
        "  )\n",
        "train_batches_stream = data_pipeline(train_stream)\n",
        "eval_batches_stream = data_pipeline(eval_stream)\n",
        "\n",
        "# Print example shapes.\n",
        "example_batch = next(train_batches_stream)\n",
        "print(f'batch shapes = {[x.shape for x in example_batch]}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "executionInfo": {
          "elapsed": 409,
          "status": "ok",
          "timestamp": 1607390286085,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "WoSz5plIyXOU",
        "outputId": "aa1db911-96fb-430b-8360-1a6e3f764cee"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Serial[\n",
            "  Embedding_8192_256\n",
            "  Mean\n",
            "  Dense_2\n",
            "]\n"
          ]
        }
      ],
      "source": [
        "# Create the model.\n",
        "from trax import layers as tl\n",
        "\n",
        "model = tl.Serial(\n",
        "    tl.Embedding(vocab_size=8192, d_feature=256),\n",
        "    tl.Mean(axis=1),  # Average on axis 1 (length of sentence).\n",
        "    tl.Dense(2),      # Classify 2 classes.\n",
        ")\n",
        "\n",
        "# You can print model structure.\n",
        "print(model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "executionInfo": {
          "elapsed": 79139,
          "status": "ok",
          "timestamp": 1607390365232,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "d6bIKUO-3Cw8",
        "outputId": "ba4199f4-cc31-459e-b46c-d14ec2f4ef68"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Step      1: Total number of trainable weights: 2097666\n",
            "Step      1: Ran 1 train steps in 1.01 secs\n",
            "Step      1: train WeightedCategoryCrossEntropy |  0.69292086\n",
            "Step      1: eval  WeightedCategoryCrossEntropy |  0.68457415\n",
            "Step      1: eval      WeightedCategoryAccuracy |  0.56406250\n",
            "\n",
            "Step    500: Ran 499 train steps in 19.92 secs\n",
            "Step    500: train WeightedCategoryCrossEntropy |  0.50587755\n",
            "Step    500: eval  WeightedCategoryCrossEntropy |  0.46716719\n",
            "Step    500: eval      WeightedCategoryAccuracy |  0.80625000\n",
            "\n",
            "Step   1000: Ran 500 train steps in 17.50 secs\n",
            "Step   1000: train WeightedCategoryCrossEntropy |  0.36375266\n",
            "Step   1000: eval  WeightedCategoryCrossEntropy |  0.44373559\n",
            "Step   1000: eval      WeightedCategoryAccuracy |  0.80000000\n",
            "\n",
            "Step   1500: Ran 500 train steps in 18.40 secs\n",
            "Step   1500: train WeightedCategoryCrossEntropy |  0.34449804\n",
            "Step   1500: eval  WeightedCategoryCrossEntropy |  0.34941847\n",
            "Step   1500: eval      WeightedCategoryAccuracy |  0.84687500\n",
            "\n",
            "Step   2000: Ran 500 train steps in 17.18 secs\n",
            "Step   2000: train WeightedCategoryCrossEntropy |  0.28685242\n",
            "Step   2000: eval  WeightedCategoryCrossEntropy |  0.50030373\n",
            "Step   2000: eval      WeightedCategoryAccuracy |  0.77539062\n"
          ]
        }
      ],
      "source": [
        "# Train the model.\n",
        "from trax.supervised import training\n",
        "\n",
        "# Training task.\n",
        "train_task = training.TrainTask(\n",
        "    labeled_data=train_batches_stream,\n",
        "    loss_layer=tl.WeightedCategoryCrossEntropy(),\n",
        "    optimizer=trax.optimizers.Adam(0.01),\n",
        "    n_steps_per_checkpoint=500,\n",
        ")\n",
        "\n",
        "# Evaluaton task.\n",
        "eval_task = training.EvalTask(\n",
        "    labeled_data=eval_batches_stream,\n",
        "    metrics=[tl.WeightedCategoryCrossEntropy(), tl.WeightedCategoryAccuracy()],\n",
        "    n_eval_batches=20  # For less variance in eval numbers.\n",
        ")\n",
        "\n",
        "# Training loop saves checkpoints to output_dir.\n",
        "output_dir = os.path.expanduser('~/output_dir/')\n",
        "training_loop = training.Loop(model,\n",
        "                              train_task,\n",
        "                              eval_tasks=[eval_task],\n",
        "                              output_dir=output_dir)\n",
        "\n",
        "# Run 2000 steps (batches).\n",
        "training_loop.run(2000)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "executionInfo": {
          "elapsed": 832,
          "status": "ok",
          "timestamp": 1607390366089,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "yuPu37Lp7GST",
        "outputId": "b95f944d-b5e8-44c6-829c-25c0b0b08f38"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "example input_str: The movie features another exceptional collaboration between director William Wyler and cinematographer Gregg Toland, the first after Toland worked on Citizen Kane. But the talent of both these men was focused on achieving a perfectly crafted movie, understood in the good old American sense as a great story. The technical aspects of the movie are covered so as the viewer gets absorbed into the action that takes place on the screen without submitting to the power of the image. Technique is seen as a vehicle of representation unlike in Citizen Kane where Welles' baroque style almost drew the attention from the story to the way the story was told. One of my favorite moves with deep focus in this film is the drama conveyed by the returning home welcoming of Homer and Al. If Homer's girl, Wilma comes towards him perfectly in focus, Al goes over to his wife also perfectly in focus. This is a brilliant move because it shows only through the use of the image the nature of these relationships as we will see them throughout the movie: Wilma loves Homer and she accepts him as he is, Al's wife loves him also but she feels unprepared to fully welcome him home. Also later in the film we find out that their marriage has not always been a bed of roses.\u003cbr /\u003e\u003cbr /\u003eWyler is a director whose force lies in being true to his work without feeling the need to boast. He wanted to show his audience how hard it was for the American soldiers returning from the war to fit into a society that either didn't understand them or treated them with contempt. With a perfect cast and great dialogue Goldwin and Wyler produced a movie that will forever be the template for any other returning home movie. The three hours which coincide with the \"rough cut\" because the test audience back then never felt for a moment that the action was slow and indeed every scene from the film seems perfectly justified. The whole thing is constructed beautifully, every character gets a fair amount of exposure, nothing is left to chance and it is quite pitiful that Hollywood nowadays never manages to bring so much character conflict to the screen. TBYOOL explores the depth of the American way of life, of the American family and society to an extent that makes other movies look like \"the children's hour\".\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\u003cpad\u003e\n",
            "Model returned sentiment activations: [[-1.6396211  1.6328843]]\n"
          ]
        }
      ],
      "source": [
        "# Run on an example.\n",
        "example_input = next(eval_batches_stream)[0][0]\n",
        "example_input_str = trax.data.detokenize(example_input, vocab_file='en_8k.subword')\n",
        "print(f'example input_str: {example_input_str}')\n",
        "sentiment_activations = model(example_input[None, :])  # Add batch dimension.\n",
        "print(f'Model returned sentiment activations: {np.asarray(sentiment_activations)}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8wgfJyhdihfR"
      },
      "source": [
        "## 2. Convert Trax to Keras\n",
        "\n",
        "Thanks to [TensorFlow NumPy](https://www.tensorflow.org/api_docs/python/tf/experimental/numpy) you can convert the model you just trained into a [Keras](https://keras.io/) layer using `trax.AsKeras`. This allows you to:\n",
        "\n",
        "* use Trax layers inside Keras models\n",
        "* run Trax models with existing Keras input pipelines\n",
        "* export Trax models to [TensorFlow SavedModel](https://www.tensorflow.org/guide/saved_model)\n",
        "\n",
        "When creating  a Keras layer from a Trax one, the Keras layer weights will get initialized to the ones the Trax layer had at the moment of creation. In this way, you can create Keras layers from pre-trained Trax models and save them as SavedModel as shown below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "executionInfo": {
          "elapsed": 322,
          "status": "ok",
          "timestamp": 1607390366418,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "bxSLRyjftuxH",
        "outputId": "6ec7180b-ff85-47e4-bba2-3634df913ad4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u003ctrax.trax2keras.AsKeras object at 0x7efff5a47a90\u003e\n",
            "Keras returned sentiment activations: [[-1.6396211  1.6328843]]\n"
          ]
        }
      ],
      "source": [
        "# Convert the model into a Keras layer, use the weights from model.\n",
        "keras_layer = trax.AsKeras(model)\n",
        "print(keras_layer)\n",
        "\n",
        "# Run the Keras layer to verify it returns the same result.\n",
        "sentiment_activations = keras_layer(example_input[None, :])\n",
        "print(f'Keras returned sentiment activations: {np.asarray(sentiment_activations)}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "executionInfo": {
          "elapsed": 3983,
          "status": "ok",
          "timestamp": 1607390370412,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "r8C-FoFGxGE1",
        "outputId": "0edfd1fa-2677-494a-f03f-2cc87324e88c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Keras returned sentiment activations: [[-1.6396211  1.6328843]]\n"
          ]
        }
      ],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "# Create a full Keras  model using the layer from Trax.\n",
        "inputs = tf.keras.Input(shape=(None,), dtype='int32')\n",
        "hidden = keras_layer(inputs) \n",
        "# You can add other Keras layers here operating on hidden.\n",
        "outputs = hidden\n",
        "keras_model = tf.keras.Model(inputs=inputs, outputs=outputs)\n",
        "print(keras_model)\n",
        "\n",
        "# Run the Keras model to verify it returns the same result.\n",
        "sentiment_activations = keras_model(example_input[None, :])\n",
        "print(f'Keras returned sentiment activations: {np.asarray(sentiment_activations)}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EQH1bvXwy5fE"
      },
      "source": [
        "## 3. Exporting Trax Models for Deployment\n",
        "\n",
        "You can export the Keras model to disk as [TensorFlow SavedModel](https://www.tensorflow.org/guide/saved_model). It's as simple as calling `keras_model.save` and allows you to use models with TF tools [TensorFlow.js](https://www.tensorflow.org/js/), [TensorFlow Serving](https://www.tensorflow.org/tfx/guide/serving) and [TensorFlow Lite](https://www.tensorflow.org/lite)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "executionInfo": {
          "elapsed": 1355,
          "status": "ok",
          "timestamp": 1607390371776,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 480
        },
        "id": "nQIJrOUgxRfK",
        "outputId": "62c028a5-da9e-40b1-d223-aa5f45b6a2aa"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Keras returned sentiment activations: [[-1.6396211  1.6328843]]\n"
          ]
        }
      ],
      "source": [
        "# Save the Keras model to output_dir.\n",
        "model_file = os.path.join(output_dir, \"model_checkpoint\")\n",
        "keras_model.save(model_file)\n",
        "\n",
        "# Load the model from SavedModel.\n",
        "loaded_model = tf.keras.models.load_model(model_file)\n",
        "\n",
        "# Run the loaded model to verify it returns the same result.\n",
        "sentiment_activations = loaded_model(example_input[None, :])\n",
        "print(f'Keras returned sentiment activations: {np.asarray(sentiment_activations)}')"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/public/tools/ml_python:ml_notebook",
        "kind": "private"
      },
      "name": "Using Trax with Keras",
      "provenance": [
        {
          "file_id": "1RNbQoOuzKsp_FoDqOFQX4mA--Wzt5ofq",
          "timestamp": 1596181556972
        },
        {
          "file_id": "https://github.com/google/trax/blob/master/trax/intro.ipynb",
          "timestamp": 1596178511100
        },
        {
          "file_id": "trax/intro.ipynb",
          "timestamp": 1595931762204
        },
        {
          "file_id": "1v1GvTkEFjMH_1c-bdS7JzNS70u9RUEHV",
          "timestamp": 1578964243645
        },
        {
          "file_id": "1SplqILjJr_ZqXcIUkNIk0tSbthfhYm07",
          "timestamp": 1572044421118
        },
        {
          "file_id": "intro.ipynb",
          "timestamp": 1571858674399
        },
        {
          "file_id": "1sF8QbqJ19ZU6oy5z4GUTt4lgUCjqO6kt",
          "timestamp": 1569980697572
        },
        {
          "file_id": "1EH76AWQ_pvT4i8ZXfkv-SCV4MrmllEl5",
          "timestamp": 1563927451951
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}