{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3825a3f9",
   "metadata": {},
   "source": [
    "# Text Generation using T5\n",
    "\n",
    "* This tutorial is intended to provide, a familiarity in how to use ```T5``` for text-generation tasks.\n",
    "* No training is involved in this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7be005d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install tf-transformers\n",
    "\n",
    "!pip install transformers\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "393be3cf",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ec2a48ab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensorflow version 2.7.0\n",
      "Devices [PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'), PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # Supper TF warnings\n",
    "\n",
    "import tensorflow as tf\n",
    "print(\"Tensorflow version\", tf.__version__)\n",
    "print(\"Devices\", tf.config.list_physical_devices())\n",
    "\n",
    "from tf_transformers.models import T5Model, T5TokenizerTFText\n",
    "from tf_transformers.core import TextGenerationChainer\n",
    "from tf_transformers.text import TextDecoder, TextDecoderSerializable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76fe06fe",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "8414afe2",
   "metadata": {},
   "source": [
    "### Load T5 Model \n",
    "\n",
    "* 1. Note `use_auto_regressive=True`, argument. This is required for any models to enable text-generation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "55a67314",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2bf18e0b91af4b5984a0bc1412960cf2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/773k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f34463285e0545ea881fae01b53ed1f3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/1.32M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aa40618dc2194e4f8ea78d368527ef01",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/1.17k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:absl:Saving t5-small tokenizer to /var/folders/vq/4fxns8l55gq8_msgygbyb51h0000gn/T/tftransformers_tokenizer_cache/t5-small\n",
      "INFO:absl:Loading t5-small tokenizer to /var/folders/vq/4fxns8l55gq8_msgygbyb51h0000gn/T/tftransformers_tokenizer_cache/t5-small/spiece.model\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Metal device set to: Apple M1\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "409e6c62a0d84661a64547ba48f5a844",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/242M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1c71328a883d44289bc3a541f751f644",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/11.5k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d03a8ca4ffbe4d49b698337ea9f6942e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/712 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Inconsistent references when loading the checkpoint into this object graph. Either the Trackable object references in the Python program have changed in an incompatible way, or the checkpoint was generated in an incompatible program.\n",
      "\n",
      "Two checkpoint references resolved to different objects (<tf_transformers.models.encoder_decoder.encoder_decoder.EncoderDecoder object at 0x17b752e20> and <keras.engine.input_layer.InputLayer object at 0x17b752340>).\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Inconsistent references when loading the checkpoint into this object graph. Either the Trackable object references in the Python program have changed in an incompatible way, or the checkpoint was generated in an incompatible program.\n",
      "\n",
      "Two checkpoint references resolved to different objects (<tf_transformers.models.encoder_decoder.encoder_decoder.EncoderDecoder object at 0x17b752e20> and <keras.engine.input_layer.InputLayer object at 0x17b752340>).\n",
      "INFO:absl:Successful ✅✅: Model checkpoints matched and loaded from /Users/sarathrnair/.cache/huggingface/hub/tftransformers__t5-small.main.699b12fe9601feda4892ca82c07e800f3c1da440/ckpt-1\n",
      "INFO:absl:Successful ✅: Loaded model from tftransformers/t5-small\n"
     ]
    }
   ],
   "source": [
    "model_name = 't5-small'\n",
    "\n",
    "tokenizer = T5TokenizerTFText.from_pretrained(model_name, dynamic_padding=True, truncate=True, max_length=256)\n",
    "model = T5Model.from_pretrained(model_name, use_auto_regressive=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5ab11b9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "d1812648",
   "metadata": {},
   "source": [
    "### Serialize and load\n",
    "\n",
    "* The most recommended way of using a Tensorflow model is to load it after serializing.\n",
    "* The speedup, especially for text generation is up to 50x times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5157a4b4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Metal device set to: Apple M1\n"
     ]
    }
   ],
   "source": [
    "# Save as serialized\n",
    "model_dir = 'MODELS/t5'\n",
    "model.save_transformers_serialized(model_dir)\n",
    "\n",
    "# Load\n",
    "loaded = tf.saved_model.load(model_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b837f9e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "22f1229e",
   "metadata": {},
   "source": [
    "### Text-Generation\n",
    "\n",
    "* . We can pass ```tf.keras.Model``` also to ```TextDecoder```, but ```SavedModel``` this is recommended"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6cbe7af8",
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder = TextDecoder(model=loaded)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "557c0764",
   "metadata": {},
   "source": [
    "### Greedy Decoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a6c4517f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[b'Das Haus ist wunderbar und wir m\\xc3\\xb6chten hier sein :)'\n",
      " b'Elle est belle::::'], shape=(2,), dtype=string)\n"
     ]
    }
   ],
   "source": [
    "texts = ['translate English to German: The house is wonderful and we wish to be here :)', \n",
    "         'translate English to French: She is beautiful']\n",
    "\n",
    "inputs = tokenizer({'text': tf.constant(texts)})\n",
    "\n",
    "predictions = decoder.decode(inputs, \n",
    "                             mode='greedy', \n",
    "                             max_iterations=64, \n",
    "                             eos_id=tokenizer.eos_token_id)\n",
    "print(tokenizer._tokenizer.detokenize(tf.squeeze(predictions['predicted_ids'], axis=1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "087e74ea",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "b47fedc0",
   "metadata": {},
   "source": [
    "### Beam Decoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d733ff03",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[b'Das Haus ist wunderbar und wir m\\xc3\\xb6chten hier sein :)'\n",
      " b'Elle est belle::::'], shape=(2,), dtype=string)\n"
     ]
    }
   ],
   "source": [
    "predictions = decoder.decode(inputs, \n",
    "                             mode='beam',\n",
    "                             num_beams=3,\n",
    "                             max_iterations=64,\n",
    "                             eos_id=tokenizer.eos_token_id)\n",
    "print(tokenizer._tokenizer.detokenize(predictions['predicted_ids'][:, 0, :]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5c7063e",
   "metadata": {},
   "source": [
    "### Top K Nucleus Sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "61c94a1c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[b'Das Haus ist wunderbar und wir m\\xc3\\xb6chten hier sein :)'\n",
      " b'Elle est belle::::'], shape=(2,), dtype=string)\n"
     ]
    }
   ],
   "source": [
    "predictions = decoder.decode(inputs, \n",
    "                             mode='top_k_top_p',\n",
    "                             top_k=50,\n",
    "                             top_p=0.7,\n",
    "                             num_return_sequences=3,\n",
    "                             max_iterations=64,\n",
    "                             eos_id=tokenizer.eos_token_id)\n",
    "print(tokenizer._tokenizer.detokenize(predictions['predicted_ids'][:, 0, :]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8bea8f8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "ed72aea7",
   "metadata": {},
   "source": [
    "### Advanced Serialization (include preprocessing + Decoding Together)\n",
    "\n",
    "* What if we can bundle all this into a single model and serialize it ?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "5c27e58c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Inconsistent references when loading the checkpoint into this object graph. Either the Trackable object references in the Python program have changed in an incompatible way, or the checkpoint was generated in an incompatible program.\n",
      "\n",
      "Two checkpoint references resolved to different objects (<tf_transformers.models.encoder_decoder.encoder_decoder.EncoderDecoder object at 0x38790b400> and <keras.engine.input_layer.InputLayer object at 0x3827cc880>).\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Inconsistent references when loading the checkpoint into this object graph. Either the Trackable object references in the Python program have changed in an incompatible way, or the checkpoint was generated in an incompatible program.\n",
      "\n",
      "Two checkpoint references resolved to different objects (<tf_transformers.models.encoder_decoder.encoder_decoder.EncoderDecoder object at 0x38790b400> and <keras.engine.input_layer.InputLayer object at 0x3827cc880>).\n",
      "INFO:absl:Successful ✅✅: Model checkpoints matched and loaded from /Users/sarathrnair/.cache/huggingface/hub/tftransformers__t5-small.main.699b12fe9601feda4892ca82c07e800f3c1da440/ckpt-1\n",
      "INFO:absl:Successful ✅: Loaded model from tftransformers/t5-small\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using default `decoder_start_token_id` 0 from the model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:Found untraced functions such as tf_transformers/t5_encoder_layer_call_fn, tf_transformers/t5_encoder_layer_call_and_return_conditional_losses, tf_transformers/t5_decoder_layer_call_fn, tf_transformers/t5_decoder_layer_call_and_return_conditional_losses, word_embeddings_layer_call_fn while saving (showing 5 of 1140). These functions will not be directly callable after loading.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: MODELS/t5_serialized/assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: MODELS/t5_serialized/assets\n"
     ]
    }
   ],
   "source": [
    "model_dir = 'MODELS/t5_serialized/'\n",
    "# Load Auto Regressive Version\n",
    "model = T5Model.from_pretrained(model_name=model_name, use_auto_regressive=True)\n",
    "# Assume we are doing beam decoding\n",
    "text_generation_kwargs = {'mode': 'beam', \n",
    "                         'num_beams': 3,\n",
    "                          'max_iterations': 32,\n",
    "                          'eos_id': tokenizer.eos_token_id\n",
    "                         }\n",
    "# TextDecoderSerializable - makes decoding serializable\n",
    "decoder = TextDecoderSerializable(model=model, **text_generation_kwargs)\n",
    "# TextGenerationChainer - joins tokenizer + TextDecoderSerializable\n",
    "model_fully_serialized = TextGenerationChainer(tokenizer.get_model(), decoder)\n",
    "model_fully_serialized = model_fully_serialized.get_model()\n",
    "# Save as saved_model\n",
    "model_fully_serialized.save_serialized(model_dir, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69f06466",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "14220f4a",
   "metadata": {},
   "source": [
    "### Load Advanced Model and Generate text\n",
    "\n",
    "* How nice is it? All done by model, no overhead of anything (tokenization, decoding, generating)\n",
    "\n",
    "* 1. TextDecoderSerializable - very advances serializable decoder written in pure tensorflow ops\n",
    "* 2. TextGenerationChainer -  very simple ```tf.keras.layers.Layer``` wrapper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "d1a0941a",
   "metadata": {},
   "outputs": [],
   "source": [
    "loaded = tf.saved_model.load(model_dir)\n",
    "model = loaded.signatures['serving_default']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "b89d4ed2",
   "metadata": {},
   "outputs": [],
   "source": [
    "texts = ['translate English to German: The house is wonderful and we wish to be here :)', \n",
    "         'translate English to French: She is beautiful']\n",
    "\n",
    "predictions = model(**{'text': tf.constant(texts)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "1a7cc7d6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[b'Das Haus ist wunderbar und wir m\\xc3\\xb6chten hier sein :)'\n",
      " b'Elle est belle::::'], shape=(2,), dtype=string)\n"
     ]
    }
   ],
   "source": [
    "print(predictions['decoded_text'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aef0c646",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,md:myst"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}