FR-EN-subword-Tokenizer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "20580905-b71b-4f76-a064-cae690feb71b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "import os\n",
    "import pathlib\n",
    "import re\n",
    "import string\n",
    "import sys\n",
    "import tempfile\n",
    "import time\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import tensorflow_datasets as tfds\n",
    "import tensorflow_text as text\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "753e93e4-26a1-4786-8b22-e1eb8328cfa7",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.get_logger().setLevel('ERROR')\n",
    "pwd = pathlib.Path.cwd()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "4565db22-2e67-48f1-a48c-635eb8075e71",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_name = 'wmt14_translate/fr-en'\n",
    "data_dir = 'nlp_lab_dataset/'\n",
    "train_samples = 60000  \n",
    "\n",
    "\n",
    "\n",
    "# Load the dataset with specified splits\n",
    "ds_splits = tfds.load(dataset_name, split=['train', 'validation', 'test'], data_dir=data_dir)\n",
    "\n",
    "# Take a subset of the training set\n",
    "train_ds = ds_splits[0].take(train_samples)\n",
    "val_ds = ds_splits[1]\n",
    "test_ds = ds_splits[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "95a68d94-fea0-44c4-a694-6ce0e6dc6a2e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training set samples:\n",
      "> Examples in English:\n",
      "In his briefing on economic development, Al Horner will give you details of programs we fund to foster partnerships between the private sector and First Nations and Inuit communities, in areas like resource development projects, for example.\n",
      "(b) Positive aspects\n",
      "Crop insurance payments include only government crop insurance programs; private hail insurance payments are excluded.\n",
      "\n",
      "> Examples in French:\n",
      "Dans sa présentation sur le développement économique, M. Al Horner vous donnera des détails sur les programmes que nous finançons pour favoriser l'établissement de partenariats entre le secteur privé et les collectivités des Premières nations et inuites dans des domaines comme celui de l'exploitation des ressources naturelles.\n",
      "b) Aspects positifs\n",
      "Les indemnités d’assurance-récolte comprennent uniquement celles des programmes publics; les indemnités de l’assurance-grêle privée sont exclues.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-19 21:47:56.958686: W tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence\n"
     ]
    }
   ],
   "source": [
    "# Example of how to use the subsets\n",
    "print(\"Training set samples:\")\n",
    "for batch in train_ds.batch(3).take(1):\n",
    "    print('> Examples in English:')\n",
    "    en_examples = batch[\"en\"].numpy()\n",
    "    for en in en_examples:\n",
    "        print(en.decode(\"utf-8\"))\n",
    "\n",
    "    print()\n",
    "    \n",
    "    print('> Examples in French:')\n",
    "    fr_examples = batch[\"fr\"].numpy()\n",
    "    for fr in fr_examples:\n",
    "        print(fr.decode(\"utf-8\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "b95ce320-c89d-41d2-b340-0a9ee081e1ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_en = train_ds.map(lambda train: train[\"en\"])\n",
    "train_fr = train_ds.map(lambda train: train[\"fr\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c0354f3-0470-49a0-b170-e0fb6453f7a5",
   "metadata": {},
   "source": [
    "### Generate the vocabulary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "37937ff6-faae-49bf-ae29-87e456b3cfce",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow_text.tools.wordpiece_vocab import bert_vocab_from_dataset as bert_vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "29c56efe-04fb-43d6-8cf9-80859d2c03ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "bert_tokenizer_params=dict(lower_case=True)\n",
    "reserved_tokens=[\"[PAD]\", \"[UNK]\", \"[START]\", \"[END]\"]\n",
    "\n",
    "bert_vocab_args = dict(\n",
    "    # The target vocabulary size\n",
    "    vocab_size = 8000,\n",
    "    # Reserved tokens that must be included in the vocabulary\n",
    "    reserved_tokens=reserved_tokens,\n",
    "    # Arguments for `text.BertTokenizer`\n",
    "    bert_tokenizer_params=bert_tokenizer_params,\n",
    "    # Arguments for `wordpiece_vocab.wordpiece_tokenizer_learner_lib.learn`\n",
    "    learn_params={},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "b43372aa-c0bd-462d-9bb9-be100d7b1e58",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-19 21:47:58.443344: W tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 25s, sys: 1.81 s, total: 1min 27s\n",
      "Wall time: 1min 20s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "fr_vocab = bert_vocab.bert_vocab_from_dataset(\n",
    "    train_fr.batch(1000).prefetch(2),\n",
    "    **bert_vocab_args\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "fcf95734-804a-4568-b7a2-8c748edbc4f3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['[PAD]', '[UNK]', '[START]', '[END]', '!', '\"', '#', '$', '%', '&']\n",
      "['ʼ', 'ˆ', 'ˇ', 'α', 'β', 'γ', 'δ', 'ε', 'η', 'θ']\n",
      "['propriete', 'attention', 'vos', 'assurance', '##ure', 'debat', 'http', 'donner', 'eux', 'protocole']\n",
      "['##\\uf76d', '##\\uf76e', '##\\uf76f', '##\\uf770', '##\\uf772', '##\\uf773', '##\\uf774', '##\\uf775', '##\\uf7e9', '##\\uf8e7']\n"
     ]
    }
   ],
   "source": [
    "print(fr_vocab[:10])\n",
    "print(fr_vocab[100:110])\n",
    "print(fr_vocab[1000:1010])\n",
    "print(fr_vocab[-10:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "544d49e2-42e9-4ab6-93ba-2cd09a9e7b02",
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_vocab_file(filepath, vocab):\n",
    "  with open(filepath, 'w') as f:\n",
    "    for token in vocab:\n",
    "      print(token, file=f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "a976f986-68fc-4f84-8d4b-431f22db31a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "write_vocab_file('fr_vocab.txt', fr_vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "f12218f4-fea8-45c8-a111-82eec146193c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-19 21:49:18.259709: W tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 9s, sys: 1.14 s, total: 1min 11s\n",
      "Wall time: 1min 6s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "en_vocab = bert_vocab.bert_vocab_from_dataset(\n",
    "    train_en.batch(1000).prefetch(2),\n",
    "    **bert_vocab_args\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "05cab7d0-211e-4e2d-b720-1139c985a1e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['[PAD]', '[UNK]', '[START]', '[END]', '!', '\"', '#', '$', '%', '&']\n",
      "['ʼ', 'ˆ', 'α', 'β', 'γ', 'δ', 'ε', 'η', 'θ', 'ι']\n",
      "['forces', 'days', 'final', 'pay', 'un', '29', 'attention', 'capital', 'prevention', 'previous']\n",
      "['##\\uf766', '##\\uf767', '##\\uf769', '##\\uf76e', '##\\uf76f', '##\\uf772', '##\\uf774', '##\\uf775', '##\\uf8e7', '##�']\n"
     ]
    }
   ],
   "source": [
    "print(en_vocab[:10])\n",
    "print(en_vocab[100:110])\n",
    "print(en_vocab[1000:1010])\n",
    "print(en_vocab[-10:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "94d76f5a-06cb-4786-b1a3-1751dc87d146",
   "metadata": {},
   "outputs": [],
   "source": [
    "write_vocab_file('en_vocab.txt', en_vocab)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a49c39d-6d33-46ec-8e7b-9ed4a56ca8d0",
   "metadata": {},
   "source": [
    "### Build the tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "id": "303dd595-2275-4f1a-b98a-94bc39232351",
   "metadata": {},
   "outputs": [],
   "source": [
    "fr_tokenizer = text.BertTokenizer('fr_vocab.txt', **bert_tokenizer_params)\n",
    "en_tokenizer = text.BertTokenizer('en_vocab.txt', **bert_tokenizer_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "9b3cf15c-54ce-4e19-bc44-17de58adbcd1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "b'In his briefing on economic development, Al Horner will give you details of programs we fund to foster partnerships between the private sector and First Nations and Inuit communities, in areas like resource development projects, for example.'\n",
      "b'(b) Positive aspects'\n",
      "b'Crop insurance payments include only government crop insurance programs; private hail insurance payments are excluded.'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-19 21:50:24.245867: W tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence\n"
     ]
    }
   ],
   "source": [
    "for batch in train_ds.batch(3).take(1):\n",
    "    en_examples = batch[\"en\"]\n",
    "    fr_examples = batch[\"fr\"]\n",
    "    for ex in en_examples:\n",
    "        print(ex.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "id": "e521a05c-6028-44b0-874d-95f6b569c304",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[295, 431, 3203, 298, 400, 331, 15, 1438, 49, 5212, 390, 316, 1013, 351, 1735, 292, 700, 324, 802, 294, 2793, 1740, 379, 291, 777, 588, 293, 378, 363, 293, 2442, 721, 15, 295, 519, 593, 1155, 331, 617, 15, 296, 691, 17]\n",
      "[11, 43, 12, 1278, 1310]\n",
      "[4060, 1164, 1409, 574, 391, 354, 4060, 1164, 700, 30, 777, 5578, 1578, 1164, 1409, 306, 3922, 17]\n"
     ]
    }
   ],
   "source": [
    "# Tokenize the examples -> (batch, word, word-piece)\n",
    "token_batch = en_tokenizer.tokenize(en_examples)\n",
    "# Merge the word and word-piece axes -> (batch, tokens)\n",
    "token_batch = token_batch.merge_dims(-2,-1)\n",
    "\n",
    "for ex in token_batch.to_list():\n",
    "  print(ex)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "id": "1d25ea75-2517-4ad1-bb99-5001c57df366",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(3,), dtype=string, numpy=\n",
       "array([b'in his briefing on economic development , al h ##orn ##er will give you details of programs we fund to foster partnerships between the private sector and first nations and inuit communities , in areas like resource development projects , for example .',\n",
       "       b'( b ) positive aspects',\n",
       "       b'crop insurance payments include only government crop insurance programs ; private ha ##il insurance payments are excluded .'],\n",
       "      dtype=object)>"
      ]
     },
     "execution_count": 98,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Lookup each token id in the vocabulary.\n",
    "txt_tokens = tf.gather(en_vocab, token_batch)\n",
    "# Join with spaces.\n",
    "tf.strings.reduce_join(txt_tokens, separator=' ', axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "id": "4cd74ad5-5f62-4988-8a63-dab21888469b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(3,), dtype=string, numpy=\n",
       "array([b'in his briefing on economic development , al horner will give you details of programs we fund to foster partnerships between the private sector and first nations and inuit communities , in areas like resource development projects , for example .',\n",
       "       b'( b ) positive aspects',\n",
       "       b'crop insurance payments include only government crop insurance programs ; private hail insurance payments are excluded .'],\n",
       "      dtype=object)>"
      ]
     },
     "execution_count": 99,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "words = en_tokenizer.detokenize(token_batch)\n",
    "tf.strings.reduce_join(words, separator=' ', axis=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b50267a4-f240-4006-8f44-0dccfda31c05",
   "metadata": {},
   "source": [
    "### Customization and export"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0716fa1d-e189-4221-8569-133caa51539c",
   "metadata": {},
   "source": [
    "#### Custom tokenization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "id": "6eaf155c-637c-46e9-87b7-dc194059c652",
   "metadata": {},
   "outputs": [],
   "source": [
    "START = tf.argmax(tf.constant(reserved_tokens) == \"[START]\")\n",
    "END = tf.argmax(tf.constant(reserved_tokens) == \"[END]\")\n",
    "\n",
    "def add_start_end(ragged):\n",
    "  count = ragged.bounding_shape()[0]\n",
    "  starts = tf.fill([count,1], START)\n",
    "  ends = tf.fill([count,1], END)\n",
    "  return tf.concat([starts, ragged, ends], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "id": "ad663974-03f4-4532-a017-af8d461c3f8b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(3,), dtype=string, numpy=\n",
       "array([b'[START] in his briefing on economic development , al horner will give you details of programs we fund to foster partnerships between the private sector and first nations and inuit communities , in areas like resource development projects , for example . [END]',\n",
       "       b'[START] ( b ) positive aspects [END]',\n",
       "       b'[START] crop insurance payments include only government crop insurance programs ; private hail insurance payments are excluded . [END]'],\n",
       "      dtype=object)>"
      ]
     },
     "execution_count": 101,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "words = en_tokenizer.detokenize(add_start_end(token_batch))\n",
    "tf.strings.reduce_join(words, separator=' ', axis=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3d59e86-4611-407d-87d1-e5f4e90191da",
   "metadata": {},
   "source": [
    "#### Custom detokenization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "id": "d4f634e3-64ee-47ad-86a2-2ad5b867ea3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cleanup_text(reserved_tokens, token_txt):\n",
    "  # Drop the reserved tokens, except for \"[UNK]\".\n",
    "  bad_tokens = [re.escape(tok) for tok in reserved_tokens if tok != \"[UNK]\"]\n",
    "  bad_token_re = \"|\".join(bad_tokens)\n",
    "\n",
    "  bad_cells = tf.strings.regex_full_match(token_txt, bad_token_re)\n",
    "  result = tf.ragged.boolean_mask(token_txt, ~bad_cells)\n",
    "\n",
    "  # Join them into strings.\n",
    "  result = tf.strings.reduce_join(result, separator=' ', axis=-1)\n",
    "\n",
    "  return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "id": "dc94aa45-98a8-4609-a1cc-fa86706c80cd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([b'In his briefing on economic development, Al Horner will give you details of programs we fund to foster partnerships between the private sector and First Nations and Inuit communities, in areas like resource development projects, for example.',\n",
       "       b'(b) Positive aspects',\n",
       "       b'Crop insurance payments include only government crop insurance programs; private hail insurance payments are excluded.'],\n",
       "      dtype=object)"
      ]
     },
     "execution_count": 103,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "en_examples.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "id": "d582ce5a-1bb0-4f31-b8bd-ba04a2297429",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.RaggedTensor [[b'in', b'his', b'briefing', b'on', b'economic', b'development', b',',\n",
       "  b'al', b'horner', b'will', b'give', b'you', b'details', b'of',\n",
       "  b'programs', b'we', b'fund', b'to', b'foster', b'partnerships',\n",
       "  b'between', b'the', b'private', b'sector', b'and', b'first', b'nations',\n",
       "  b'and', b'inuit', b'communities', b',', b'in', b'areas', b'like',\n",
       "  b'resource', b'development', b'projects', b',', b'for', b'example', b'.'],\n",
       " [b'(', b'b', b')', b'positive', b'aspects'],\n",
       " [b'crop', b'insurance', b'payments', b'include', b'only', b'government',\n",
       "  b'crop', b'insurance', b'programs', b';', b'private', b'hail',\n",
       "  b'insurance', b'payments', b'are', b'excluded', b'.']                  ]>"
      ]
     },
     "execution_count": 104,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "token_batch = en_tokenizer.tokenize(en_examples).merge_dims(-2,-1)\n",
    "words = en_tokenizer.detokenize(token_batch)\n",
    "words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "id": "6ddb6fca-d7bc-4743-a444-d466eccf70a0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([b'in his briefing on economic development , al horner will give you details of programs we fund to foster partnerships between the private sector and first nations and inuit communities , in areas like resource development projects , for example .',\n",
       "       b'( b ) positive aspects',\n",
       "       b'crop insurance payments include only government crop insurance programs ; private hail insurance payments are excluded .'],\n",
       "      dtype=object)"
      ]
     },
     "execution_count": 105,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cleanup_text(reserved_tokens, words).numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c780c1e0-6025-41f4-b062-2020ffe89aa2",
   "metadata": {},
   "source": [
    "### Export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "id": "355eb800-b79f-4f3b-8f1d-2e710d40d5fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomTokenizer(tf.Module):\n",
    "  def __init__(self, reserved_tokens, vocab_path):\n",
    "    self.tokenizer = text.BertTokenizer(vocab_path, lower_case=True)\n",
    "    self._reserved_tokens = reserved_tokens\n",
    "    self._vocab_path = tf.saved_model.Asset(vocab_path)\n",
    "\n",
    "    vocab = pathlib.Path(vocab_path).read_text().splitlines()\n",
    "    self.vocab = tf.Variable(vocab)\n",
    "\n",
    "    ## Create the signatures for export:   \n",
    "\n",
    "    # Include a tokenize signature for a batch of strings. \n",
    "    self.tokenize.get_concrete_function(\n",
    "        tf.TensorSpec(shape=[None], dtype=tf.string))\n",
    "\n",
    "    # Include `detokenize` and `lookup` signatures for:\n",
    "    #   * `Tensors` with shapes [tokens] and [batch, tokens]\n",
    "    #   * `RaggedTensors` with shape [batch, tokens]\n",
    "    self.detokenize.get_concrete_function(\n",
    "        tf.TensorSpec(shape=[None, None], dtype=tf.int64))\n",
    "    self.detokenize.get_concrete_function(\n",
    "          tf.RaggedTensorSpec(shape=[None, None], dtype=tf.int64))\n",
    "\n",
    "    self.lookup.get_concrete_function(\n",
    "        tf.TensorSpec(shape=[None, None], dtype=tf.int64))\n",
    "    self.lookup.get_concrete_function(\n",
    "          tf.RaggedTensorSpec(shape=[None, None], dtype=tf.int64))\n",
    "\n",
    "    # These `get_*` methods take no arguments\n",
    "    self.get_vocab_size.get_concrete_function()\n",
    "    self.get_vocab_path.get_concrete_function()\n",
    "    self.get_reserved_tokens.get_concrete_function()\n",
    "\n",
    "  @tf.function\n",
    "  def tokenize(self, strings):\n",
    "    enc = self.tokenizer.tokenize(strings)\n",
    "    # Merge the `word` and `word-piece` axes.\n",
    "    enc = enc.merge_dims(-2,-1)\n",
    "    enc = add_start_end(enc)\n",
    "    return enc\n",
    "\n",
    "  @tf.function\n",
    "  def detokenize(self, tokenized):\n",
    "    words = self.tokenizer.detokenize(tokenized)\n",
    "    return cleanup_text(self._reserved_tokens, words)\n",
    "\n",
    "  @tf.function\n",
    "  def lookup(self, token_ids):\n",
    "    return tf.gather(self.vocab, token_ids)\n",
    "\n",
    "  @tf.function\n",
    "  def get_vocab_size(self):\n",
    "    return tf.shape(self.vocab)[0]\n",
    "\n",
    "  @tf.function\n",
    "  def get_vocab_path(self):\n",
    "    return self._vocab_path\n",
    "\n",
    "  @tf.function\n",
    "  def get_reserved_tokens(self):\n",
    "    return tf.constant(self._reserved_tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "id": "1ccf3402-53d1-4546-b231-e6d5a4983b71",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizers = tf.Module()\n",
    "tokenizers.fr = CustomTokenizer(reserved_tokens, 'fr_vocab.txt')\n",
    "tokenizers.en = CustomTokenizer(reserved_tokens, 'en_vocab.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "id": "d1443fc4-69eb-4941-9711-a186344ebcef",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'fr_en_tokenizer'\n",
    "tf.saved_model.save(tokenizers, model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "id": "60278f3c-e636-4a3d-a9af-d5ef370d643a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "7955"
      ]
     },
     "execution_count": 109,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reloaded_tokenizers = tf.saved_model.load(model_name)\n",
    "reloaded_tokenizers.en.get_vocab_size().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "id": "5f745277-c949-41ea-a2b4-5b5b5522a9f6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[   2,  429, 2033,  423, 2541, 7258,  668, 4688,    4,    3]])"
      ]
     },
     "execution_count": 110,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokens = reloaded_tokenizers.en.tokenize(['Hello TensorFlow!'])\n",
    "tokens.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "id": "679fb0fa-9233-42bb-b3cb-e6984b1e5ef7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.RaggedTensor [[b'[START]', b'##i', b'settlement', b'part', b'saint', b'##pelled',\n",
       "  b'result', b'##ew', b'!', b'[END]']]>"
      ]
     },
     "execution_count": 111,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text_tokens = reloaded_tokenizers.en.lookup(tokens)\n",
    "text_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "id": "2c2c93ec-5399-46f6-bb02-b177cc3f1181",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hello tensorflow !\n"
     ]
    }
   ],
   "source": [
    "round_trip = reloaded_tokenizers.en.detokenize(tokens)\n",
    "\n",
    "print(round_trip.numpy()[0].decode('utf-8'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "id": "41518d1e-86ad-4db2-be53-6f96f18247ad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  adding: fr_en_tokenizer/ (stored 0%)\n",
      "  adding: fr_en_tokenizer/variables/ (stored 0%)\n",
      "  adding: fr_en_tokenizer/variables/variables.data-00000-of-00001 (deflated 51%)\n",
      "  adding: fr_en_tokenizer/variables/variables.index (deflated 33%)\n",
      "  adding: fr_en_tokenizer/assets/ (stored 0%)\n",
      "  adding: fr_en_tokenizer/assets/en_vocab.txt (deflated 54%)\n",
      "  adding: fr_en_tokenizer/assets/fr_vocab.txt (deflated 57%)\n",
      "  adding: fr_en_tokenizer/saved_model.pb (deflated 91%)\n",
      "  adding: fr_en_tokenizer/fingerprint.pb (stored 0%)\n"
     ]
    }
   ],
   "source": [
    "!zip -r {model_name}.zip {model_name}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "911dd5fb-b846-4a1d-80db-4225ff4cfab6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}