tfjs-convert.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflowjs as tfjs\n",
    "import numpy as np\n",
    "import os, time, json, dill, textwrap, re, string, shutil\n",
    "from tqdm import tqdm\n",
    "from MyModel import MyModel, OneStep\n",
    "from tokenizer import tokenizer, detokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Compute dtype: float32\n",
      "Variable dtype: float32\n"
     ]
    }
   ],
   "source": [
    "policy = tf.keras.mixed_precision.Policy('float32')\n",
    "tf.keras.mixed_precision.set_global_policy(policy)\n",
    "\n",
    "print('Compute dtype: %s' % policy.compute_dtype)\n",
    "print('Variable dtype: %s' % policy.variable_dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# gpu growth\n",
    "gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "for gpu in gpus:\n",
    "    tf.config.experimental.set_memory_growth(gpu, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load model\n",
    "model_name = \"models/model1-test2\"\n",
    "meta = json.load(open(f\"{model_name}/meta.json\"))\n",
    "vocab = json.load(open(f\"{model_name}/vocab.json\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OneStep.__init__\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(1, 1, 5000), dtype=float32, numpy=\n",
       "array([[[-0.00034489,  0.00017292,  0.00110298, ..., -0.00138007,\n",
       "         -0.00050311,  0.00085747]]], dtype=float32)>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = MyModel(vocab_size=meta[\"vocab_size\"], embedding_dim=meta[\"embedding_dim\"], rnn_units=meta[\"rnn_units\"])\n",
    "\n",
    "one_step_model = OneStep(model)\n",
    "model(tf.zeros([1, 1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_weights(f\"{model_name}/weights.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor([  2   7 253 183  15 142  18], shape=(7,), dtype=int32)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [00:01<00:00, 52.77it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================================\n",
      "I recently published my thesis for a year before. I don't know if I should mention my\n",
      "paper in this paper before I handed it out to the authors. Would it be appropriate to\n",
      "thank the editor based on the blog suggestion which I want to publish in the journal\n",
      "because my paper is likely that he has 2 good reviews?\n",
      "\n",
      "Anyway, it seems like time was\n",
      "that he has contributed more to my work, and he took two of his data analysis papers. The\n",
      "book is still widely in its\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "def generate_text(starting, num_generate, onestepmodel=None):\n",
    "    if onestepmodel is None:\n",
    "        onestepmodel = one_step_model\n",
    "    states = None\n",
    "    # next_char = text_vectorization(tokenizer(starting))\n",
    "    next_char = tf.constant(\n",
    "        [(vocab.index(c) if c in vocab else vocab.index(\"[UNK]\")) for c in tokenizer(starting)]\n",
    "    )\n",
    "    print(next_char)\n",
    "    result = next_char.numpy()\n",
    "\n",
    "    for _ in tqdm(range(num_generate)):\n",
    "        next_char = tf.expand_dims(next_char, 0)\n",
    "        next_char, states = one_step_model.generate_one_step(next_char, states=states, temperature=0.9)\n",
    "        # while True:\n",
    "        #     next_char1, states1 = one_step_model.generate_one_step(next_char, states=states, temperature=0.1)\n",
    "        #     if not (next_char1[0] == 1):\n",
    "        #         next_char, states = next_char1, states1\n",
    "        #         break\n",
    "        if next_char[0] == vocab.index(\"<SPECIAL:END>\"):\n",
    "            break\n",
    "\n",
    "        result = np.append(result, next_char.numpy(), axis=0)\n",
    "    result = detokenizer([vocab[r] for r in result])\n",
    "    w = textwrap.TextWrapper(width=90, break_long_words=True, replace_whitespace=False)\n",
    "    print(\"=\" * 50)\n",
    "    print(w.fill(result))\n",
    "\n",
    "\n",
    "generate_text(\"I recently published my thesis for\", 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Skipping full serialization of Keras layer <MyModel.OneStep object at 0x0000026A797EBE20>, because it is not built.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:Found untraced functions such as lstm_cell_layer_call_fn, lstm_cell_layer_call_and_return_conditional_losses while saving (showing 2 of 2). These functions will not be directly callable after loading.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: tmp/model2-test/saved_model\\assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: tmp/model2-test/saved_model\\assets\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor([  2   7 253 183  15 142  18], shape=(7,), dtype=int32)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 76%|███████▌  | 76/100 [00:00<00:00, 168.89it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================================\n",
      "I recently published my thesis for a (new) Phd.\n",
      "\n",
      "I would like to ask whether this is the\n",
      "case, should I include a note that an author is he / she will be to contact the paper on\n",
      "his / her Phd negatively or assume that I am the sole author of the paper, or would there\n",
      "be any potential ethical in which a claim is not working on?\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# save model as saved_model with signature\n",
    "\n",
    "NEW_MODEL_NAME = \"tmp/model2-test\"\n",
    "\n",
    "tf.saved_model.save(\n",
    "    one_step_model,\n",
    "    f\"{NEW_MODEL_NAME}/saved_model\",\n",
    "    signatures=one_step_model.generate_one_step.get_concrete_function(\n",
    "        tf.TensorSpec(shape=[None, None], dtype=tf.int32, name=\"input_ids\"),\n",
    "        tf.TensorSpec(shape=[2, None, None], dtype=tf.float32, name=\"states\"),\n",
    "        tf.TensorSpec(shape=[], dtype=tf.float32, name=\"temperature\"),\n",
    "    ),\n",
    ")\n",
    "\n",
    "one_step_model2 = tf.saved_model.load(f\"{NEW_MODEL_NAME}/saved_model\")\n",
    "\n",
    "\n",
    "generate_text(\"I recently published my thesis for\", 100, onestepmodel=one_step_model2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor([  2  15 142  14 104   5  28], shape=(7,), dtype=int32)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [00:00<00:00, 178.73it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================================\n",
      "My thesis is good, but unfortunately I have a strong inclination to reject the paper\n",
      "without being asked out.\n",
      "\n",
      "For my first few days, my advisor wanted to take the draft of\n",
      "the second paper online, but the people just found my thesis.\n",
      "\n",
      "My Phd advisor told me to\n",
      "not mention that I am on track committee (this process of mine).\n",
      "\n",
      "I have been struggling\n",
      "to contact all the talks from the world and physical (I see it as a very bad teacher)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "generate_text(\"My thesis is good, but\", 100, onestepmodel=one_step_model2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing weight file tmp/model2-test/tfjs\\model.json...\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'tmp/model2-test/tfjs/tokenizer.py'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tfjs.converters.convert_tf_saved_model(\n",
    "    f\"{NEW_MODEL_NAME}/saved_model\", f\"{NEW_MODEL_NAME}/tfjs\", quantization_dtype_map={\"float16\": \"*\"}\n",
    ")\n",
    "\n",
    "# save vocab\n",
    "with open(f\"{NEW_MODEL_NAME}/tfjs/vocab.json\", \"w\") as f:\n",
    "    json.dump(vocab, f)\n",
    "\n",
    "# copy tokenizer.py to tfjs folder\n",
    "shutil.copy(\"tokenizer.py\", f\"{NEW_MODEL_NAME}/tfjs/tokenizer.py\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "ename": "Exception",
     "evalue": "stop here",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mException\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-17-3799461a9b8d>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[1;32m----> 1\u001b[1;33m \u001b[1;32mraise\u001b[0m \u001b[0mException\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m\"stop here\"\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;31mException\u001b[0m: stop here"
     ]
    }
   ],
   "source": [
    "raise Exception(\"stop here\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "converter = tf.lite.TFLiteConverter.from_saved_model(f\"{NEW_MODEL_NAME}/saved_model\")\n",
    "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
    "converter.target_spec.supported_ops = [\n",
    "  tf.lite.OpsSet.TFLITE_BUILTINS, # enable TensorFlow Lite ops.\n",
    "  tf.lite.OpsSet.SELECT_TF_OPS # enable TensorFlow ops.\n",
    "]\n",
    "tflite_quant_model = converter.convert()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'tmp/model2-test/tflite/tokenizer.py'"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# make folder\n",
    "os.makedirs(f\"{NEW_MODEL_NAME}/tflite\", exist_ok=True)\n",
    "\n",
    "# save tflite model\n",
    "with open(f\"{NEW_MODEL_NAME}/tflite/model.tflite\", \"wb\") as f:\n",
    "    f.write(tflite_quant_model)\n",
    "\n",
    "# save vocab\n",
    "with open(f\"{NEW_MODEL_NAME}/tflite/vocab.json\", \"w\") as f:\n",
    "    json.dump(vocab, f)\n",
    "\n",
    "# copy tokenizer.py to tflite folder\n",
    "shutil.copy(\"tokenizer.py\", f\"{NEW_MODEL_NAME}/tflite/tokenizer.py\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tf",
   "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.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}