convert_pytorch_distill.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from tf2gpt.model import GPT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mp_rank_00_model_states.pt  mp_rank_01_model_states.pt\r\n"
     ]
    }
   ],
   "source": [
    "!ls 'CPM-distill/310000'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "m0 = torch.load('./CPM-distill/310000/mp_rank_00_model_states.pt', map_location='cpu')\n",
    "m1 = torch.load('./CPM-distill/310000/mp_rank_01_model_states.pt', map_location='cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_weight(model, name):\n",
    "    for n, w in model['module'].items():\n",
    "        if name == n:\n",
    "            return w, list(w.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "word_embeddings.weight torch.Size([15000, 768])\n",
      "position_embeddings.weight torch.Size([1024, 768])\n",
      "transformer.layers.0.input_layernorm.weight torch.Size([768])\n",
      "transformer.layers.0.input_layernorm.bias torch.Size([768])\n",
      "transformer.layers.0.attention.query_key_value.weight torch.Size([1152, 768])\n",
      "transformer.layers.0.attention.query_key_value.bias torch.Size([1152])\n",
      "transformer.layers.0.attention.dense.weight torch.Size([768, 384])\n",
      "transformer.layers.0.attention.dense.bias torch.Size([768])\n",
      "transformer.layers.0.post_attention_layernorm.weight torch.Size([768])\n",
      "transformer.layers.0.post_attention_layernorm.bias torch.Size([768])\n",
      "transformer.layers.0.mlp.dense_h_to_4h.weight torch.Size([1536, 768])\n",
      "transformer.layers.0.mlp.dense_h_to_4h.bias torch.Size([1536])\n",
      "transformer.layers.0.mlp.dense_4h_to_h.weight torch.Size([768, 1536])\n",
      "transformer.layers.0.mlp.dense_4h_to_h.bias torch.Size([768])\n",
      "transformer.final_layernorm.weight torch.Size([768])\n",
      "transformer.final_layernorm.bias torch.Size([768])\n"
     ]
    }
   ],
   "source": [
    "for n, w in m0['module'].items():\n",
    "    if '.layers.' in n:\n",
    "        if '.layers.0.' in n:\n",
    "            print(n, w.shape)\n",
    "    else:\n",
    "        print(n, w.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([15000, 768])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "find_weight(m0, 'word_embeddings.weight')[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1024, 768])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "find_weight(m0, 'position_embeddings.weight')[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# list(m0['module'].keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpt = GPT(\n",
    "    vocab_size=30_000,\n",
    "    layer_size=12,\n",
    "    block_size=1024,  # position embedding\n",
    "    embedding_dropout=0.0,\n",
    "    embedding_size=768,\n",
    "    num_attention_heads=12,\n",
    "    attention_dropout=0.0,\n",
    "    residual_dropout=0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 1, 30000)\n"
     ]
    }
   ],
   "source": [
    "print(gpt(tf.constant([[1]])).shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for x in gpt.weights:\n",
    "#     assert x.dtype == tf.float16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gpt/embedding/embeddings:0 (30000, 768)\n",
      "position_embeddings:0 (1024, 768)\n",
      "gpt/layer00/attention/query_layer/kernel:0 (768, 768)\n",
      "gpt/layer00/attention/query_layer/bias:0 (768,)\n",
      "gpt/layer00/attention/key_layer/kernel:0 (768, 768)\n",
      "gpt/layer00/attention/key_layer/bias:0 (768,)\n",
      "gpt/layer00/attention/value_layer/kernel:0 (768, 768)\n",
      "gpt/layer00/attention/value_layer/bias:0 (768,)\n",
      "gpt/layer00/attention/context_projection_layer/kernel:0 (768, 768)\n",
      "gpt/layer00/attention/context_projection_layer/bias:0 (768,)\n",
      "gpt/layer00/LayerNorm_mlp_ln0/gamma:0 (768,)\n",
      "gpt/layer00/LayerNorm_mlp_ln0/beta:0 (768,)\n",
      "gpt/layer00/LayerNorm_mlp_ln1/gamma:0 (768,)\n",
      "gpt/layer00/LayerNorm_mlp_ln1/beta:0 (768,)\n",
      "gpt/layer00/intermediate/kernel:0 (768, 3072)\n",
      "gpt/layer00/intermediate/bias:0 (3072,)\n",
      "gpt/layer00/output/kernel:0 (3072, 768)\n",
      "gpt/layer00/output/bias:0 (768,)\n",
      "gpt/LayerNorm_final_norm/gamma:0 (768,)\n",
      "gpt/LayerNorm_final_norm/beta:0 (768,)\n"
     ]
    }
   ],
   "source": [
    "for x in gpt.weights:\n",
    "    if 'gpt/layer' in x.name:\n",
    "        if 'gpt/layer00' in x.name:\n",
    "            print(x.name, x.shape)\n",
    "    else:\n",
    "        print(x.name, x.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1536, 768])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "n_layer=0\n",
    "find_weight(m0, f'transformer.layers.{n_layer}.mlp.dense_h_to_4h.weight')[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "new_weights = []\n",
    "\n",
    "for x in gpt.weights:\n",
    "    xs = list(x.shape)\n",
    "\n",
    "    if 'gpt/embedding/embeddings:' in x.name:\n",
    "        pname = 'word_embeddings.weight'\n",
    "        w0, ws0 = find_weight(m0, pname)\n",
    "        w1, ws1 = find_weight(m1, pname)\n",
    "        assert ws0 == [15_000, 768]\n",
    "        assert ws1 == [15_000, 768]\n",
    "        w = np.concatenate([w0.numpy(), w1.numpy()])\n",
    "        assert w.shape == (3_0000, 768)\n",
    "        new_weights.append((x.name, xs, pname, w))\n",
    "\n",
    "    elif 'position_embeddings' in x.name:\n",
    "        pname = 'position_embeddings.weight'\n",
    "        w0, ws0 = find_weight(m0, pname)\n",
    "        assert xs == ws0\n",
    "        w = w0.numpy()\n",
    "        new_weights.append((x.name, xs, pname, w))\n",
    "\n",
    "    elif 'gpt/layer' in x.name:\n",
    "        n_layer = int(x.name[len('gpt/layer'):][:2])\n",
    "        if 'query_layer' in x.name or 'key_layer' in x.name or 'value_layer' in x.name:\n",
    "\n",
    "            if '/kernel' in x.name:\n",
    "                pname = f'transformer.layers.{n_layer}.attention.query_key_value.weight'\n",
    "                w0, ws0 = find_weight(m0, pname)\n",
    "                w1, ws1 = find_weight(m1, pname)\n",
    "                assert ws0 == [1152, 768]\n",
    "                assert ws1 == [1152, 768]\n",
    "                w = np.concatenate([w0.numpy(), w1.numpy()])\n",
    "                w = np.transpose(w)\n",
    "                if 'query_layer' in x.name:\n",
    "                    w = np.concatenate([w0.numpy()[:384, :], w1.numpy()[:384, :]])\n",
    "                elif 'key_layer' in x.name:\n",
    "                    w = np.concatenate([w0.numpy()[384:384*2, :], w1.numpy()[384:384*2, :]])\n",
    "                elif 'value_layer' in x.name:\n",
    "                    w = np.concatenate([w0.numpy()[384*2:, :], w1.numpy()[384*2:, :]])\n",
    "                w = np.transpose(w)\n",
    "                assert w.shape == (768, 768)\n",
    "                new_weights.append((x.name, xs, pname, w))\n",
    "\n",
    "            elif '/bias' in x.name:\n",
    "                pname = f'transformer.layers.{n_layer}.attention.query_key_value.bias'\n",
    "                w0, ws0 = find_weight(m0, pname)\n",
    "                w1, ws1 = find_weight(m1, pname)\n",
    "                assert ws0 == [1152,]\n",
    "                assert ws1 == [1152,]\n",
    "                w = np.concatenate([w0.numpy(), w1.numpy()])\n",
    "                if 'query_layer' in x.name:\n",
    "                    w = np.concatenate([w0.numpy()[:384], w1.numpy()[:384]])\n",
    "                elif 'key_layer' in x.name:\n",
    "                    w = np.concatenate([w0.numpy()[384:384*2], w1.numpy()[384:384*2]])\n",
    "                elif 'value_layer' in x.name:\n",
    "                    w = np.concatenate([w0.numpy()[384*2:], w1.numpy()[384*2:]])\n",
    "                assert w.shape == (768,)\n",
    "                new_weights.append((x.name, xs, pname, w))\n",
    "\n",
    "        elif 'attention/context_projection_layer/kernel' in x.name:\n",
    "            pname = f'transformer.layers.{n_layer}.attention.dense.weight'\n",
    "            w0, ws0 = find_weight(m0, pname)\n",
    "            w1, ws1 = find_weight(m1, pname)\n",
    "            w = np.concatenate([w0.numpy(), w1.numpy()], axis=1)\n",
    "            w = np.transpose(w)\n",
    "            assert w.shape == (768, 768)\n",
    "            new_weights.append((x.name, xs, pname, w))\n",
    "\n",
    "        elif 'attention/context_projection_layer/bias' in x.name:\n",
    "            pname = f'transformer.layers.{n_layer}.attention.dense.bias'\n",
    "            w0, ws0 = find_weight(m0, pname)\n",
    "            w = w0.numpy()\n",
    "            assert w.shape == (768,)\n",
    "            new_weights.append((x.name, xs, pname, w))\n",
    "\n",
    "        elif 'LayerNorm_mlp_ln0/gamma' in x.name:\n",
    "            pname = f'transformer.layers.{n_layer}.input_layernorm.weight'\n",
    "            w0, ws0 = find_weight(m0, pname)\n",
    "            w = w0.numpy()\n",
    "            assert ws0 == xs\n",
    "            new_weights.append((x.name, x.shape, pname, w))\n",
    "\n",
    "        elif 'LayerNorm_mlp_ln1/gamma' in x.name:\n",
    "            pname = f'transformer.layers.{n_layer}.post_attention_layernorm.weight'\n",
    "            w0, ws0 = find_weight(m0, pname)\n",
    "            w = w0.numpy()\n",
    "            assert ws0 == xs\n",
    "            new_weights.append((x.name, xs, pname, w))\n",
    "\n",
    "        elif 'LayerNorm_mlp_ln0/beta' in x.name:\n",
    "            pname = f'transformer.layers.{n_layer}.input_layernorm.bias'\n",
    "            w0, ws0 = find_weight(m0, pname)\n",
    "            w = w0.numpy()\n",
    "            assert ws0 == xs\n",
    "            new_weights.append((x.name, xs, pname, w))\n",
    "\n",
    "        elif 'LayerNorm_mlp_ln1/beta' in x.name:\n",
    "            pname = f'transformer.layers.{n_layer}.post_attention_layernorm.bias'\n",
    "            w0, ws0 = find_weight(m0, pname)\n",
    "            w = w0.numpy()\n",
    "            assert ws0 == xs\n",
    "            new_weights.append((x.name, xs, pname, w))\n",
    "\n",
    "        elif 'intermediate/kernel' in x.name:\n",
    "            pname = f'transformer.layers.{n_layer}.mlp.dense_h_to_4h.weight'\n",
    "            w0, ws0 = find_weight(m0, pname)\n",
    "            w1, ws1 = find_weight(m1, pname)\n",
    "            w = np.concatenate([w0.numpy(), w1.numpy()], axis=0)\n",
    "            w = np.transpose(w)\n",
    "            assert w.shape == (768, 3072)\n",
    "            new_weights.append((x.name, xs, pname, w))\n",
    "\n",
    "        elif 'intermediate/bias' in x.name:\n",
    "            pname = f'transformer.layers.{n_layer}.mlp.dense_h_to_4h.bias'\n",
    "            w0, ws0 = find_weight(m0, pname)\n",
    "            w1, ws1 = find_weight(m1, pname)\n",
    "            w = np.concatenate([w0.numpy(), w1.numpy()], axis=0)\n",
    "            w = np.transpose(w)\n",
    "            assert w.shape == (3072,)\n",
    "            new_weights.append((x.name, xs, pname, w))\n",
    "\n",
    "        elif '/output/kernel' in x.name:\n",
    "            pname = f'transformer.layers.{n_layer}.mlp.dense_4h_to_h.weight'\n",
    "            w0, ws0 = find_weight(m0, pname)\n",
    "            w1, ws1 = find_weight(m1, pname)\n",
    "            w = np.concatenate([w0.numpy(), w1.numpy()], axis=1)\n",
    "            w = np.transpose(w)\n",
    "            assert w.shape == (3072, 768)\n",
    "            new_weights.append((x.name, xs, pname, w))\n",
    "\n",
    "        elif '/output/bias' in x.name:\n",
    "            pname = f'transformer.layers.{n_layer}.mlp.dense_4h_to_h.bias'\n",
    "            w0, ws0 = find_weight(m0, pname)\n",
    "            w = w0.numpy()\n",
    "            assert w.shape == (768,)\n",
    "            new_weights.append((x.name, xs, pname, w))\n",
    "\n",
    "        else:\n",
    "            print('BAD', x.name, xs)\n",
    "            break\n",
    "    elif 'gpt/LayerNorm_final_norm/gamma' in x.name:\n",
    "        pname = 'transformer.final_layernorm.weight'\n",
    "        w0, ws0 = find_weight(m0, pname)\n",
    "        w = w0.numpy()\n",
    "        assert ws0 == xs\n",
    "        new_weights.append((x.name, xs, pname, w))\n",
    "    elif 'gpt/LayerNorm_final_norm/beta' in x.name:\n",
    "        pname = 'transformer.final_layernorm.bias'\n",
    "        w0, ws0 = find_weight(m0, pname)\n",
    "        w = w0.numpy()\n",
    "        assert ws0 == xs\n",
    "        new_weights.append((x.name, xs, pname, w))\n",
    "    else:\n",
    "        print('BAD', x.name, xs)\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(new_weights) == len(gpt.weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "for x in new_weights:\n",
    "    assert tuple(x[1]) == x[-1].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpt.set_weights([x[-1] for x in new_weights])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gpt2_tokenizer import GPT2Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "cbpe = GPT2Tokenizer(\n",
    "    'CPM-Generate/bpe_3w_new/vocab.json',\n",
    "    'CPM-Generate/bpe_3w_new/merges.txt',\n",
    "    model_file='CPM-Generate/bpe_3w_new/chinese_vocab.model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Building prefix dict from the default dictionary ...\n",
      "Loading model from cache /tmp/jieba.cache\n",
      "Loading model cost 0.524 seconds.\n",
      "Prefix dict has been built successfully.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 今天天气 不错 \n",
      "1 今天天气 不错 ,\n",
      "2 今天天气 不错 , 我\n",
      "3 今天天气 不错 , 我 就\n",
      "4 今天天气 不错 , 我 就 去\n",
      "5 今天天气 不错 , 我 就 去 了\n",
      "6 今天天气 不错 , 我 就 去 了 \n",
      "7 今天天气 不错 , 我 就 去 了 。\n",
      "8 今天天气 不错 , 我 就 去 了 。 \n",
      "9 今天天气 不错 , 我 就 去 了 。  \n"
     ]
    }
   ],
   "source": [
    "ids = cbpe.encode('今天天气不错')\n",
    "\n",
    "for i in range(10):\n",
    "    output = gpt(tf.constant([ids]))\n",
    "    nid = np.argmax(output[0, -1])\n",
    "    ids += [nid]\n",
    "    print(i, cbpe.decode(ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 243,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def batch_gather(a, b):\n",
    "    return tf.gather(a, b, batch_dims=1)\n",
    "\n",
    "\n",
    "@tf.function\n",
    "def top_k_top_p_sample(logits, num_samples=1, top_k=0, p=0.95):\n",
    "    batch_size, vocab_size = logits.shape\n",
    "    probs = tf.nn.softmax(logits, axis=-1)\n",
    "    \n",
    "    # [batch_size, vocab_perm]\n",
    "    indices = tf.argsort(probs, direction='DESCENDING')\n",
    "    logits_to_use = batch_gather(logits, indices)\n",
    "    cumulative_probabilities = tf.math.cumsum(batch_gather(probs, indices), axis=-1, exclusive=False)\n",
    "\n",
    "    # find the top pth index to cut off. careful we don't want to cutoff everything!\n",
    "    # result will be [batch_size, vocab_perm]\n",
    "    if p > 0.0:\n",
    "        exclude_mask = tf.logical_not(\n",
    "            tf.logical_or(cumulative_probabilities < p, tf.range(vocab_size)[None] < 1))\n",
    "        # OPTION A - sample in the sorted space, then unsort.\n",
    "        logits_to_use = logits_to_use - tf.cast(exclude_mask, tf.float32) * 1e10\n",
    "    \n",
    "    if top_k > 0:\n",
    "        logits_to_use = logits_to_use - tf.cast(\n",
    "            tf.argsort(logits_to_use, direction='DESCENDING') >= top_k,\n",
    "            dtype=tf.float32\n",
    "        ) * 1e10\n",
    "    \n",
    "    sample_perm = tf.random.categorical(logits=logits_to_use, num_samples=num_samples)\n",
    "    sample = batch_gather(indices, sample_perm)\n",
    "\n",
    "    return tf.cast(sample, tf.int64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 244,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def serve(inputs):\n",
    "    return gpt(inputs, kv_cache=None, use_cache=True)\n",
    "\n",
    "\n",
    "@tf.function\n",
    "def serve_cache(inputs, kv_cache):\n",
    "    return gpt(inputs, kv_cache=kv_cache, use_cache=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 245,
   "metadata": {},
   "outputs": [],
   "source": [
    "serve_concrete = serve.get_concrete_function(\n",
    "    tf.TensorSpec(shape=[None, None], dtype=tf.int64, name=\"inp\")\n",
    ")\n",
    "\n",
    "layer_size = 12\n",
    "attention_head = 12\n",
    "embedding_size = 768\n",
    "\n",
    "serve_cache_concrete = serve_cache.get_concrete_function(\n",
    "    tf.TensorSpec(shape=[None, None], dtype=tf.int64, name=\"inp\"),\n",
    "    tf.TensorSpec(shape=[\n",
    "        layer_size, None, 2, attention_head,\n",
    "        None, embedding_size // attention_head\n",
    "    ], dtype=tf.float32, name=\"kv_cache\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 246,
   "metadata": {},
   "outputs": [],
   "source": [
    "r = serve_concrete(\n",
    "    tf.constant([[1]], tf.int64)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 247,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 1, 30000) (12, 1, 2, 12, 1, 64)\n"
     ]
    }
   ],
   "source": [
    "print(r[0].shape, r[1].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 248,
   "metadata": {},
   "outputs": [],
   "source": [
    "r2 = serve_cache_concrete(\n",
    "    tf.constant([[1]], tf.int64),\n",
    "    r[1]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 249,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 1, 30000) (12, 1, 2, 12, 1, 64)\n"
     ]
    }
   ],
   "source": [
    "print(r2[0].shape, r2[1].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 250,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def sample(initial_inputs, length, top_k, top_p, temperature):\n",
    "    layer_size = 12\n",
    "    embedding_size = 768\n",
    "    attention_head = 12\n",
    "\n",
    "    i = tf.constant(0, dtype=tf.int64)\n",
    "    initial_logits, kv_cache = serve(initial_inputs)\n",
    "    logits_with_temperature = initial_logits[:, -1, :]\n",
    "    if temperature > 0.0:\n",
    "        logits_with_temperature /= temperature\n",
    "    inputs = top_k_top_p_sample(logits_with_temperature, 1, top_k, top_p)\n",
    "    stores = tf.concat([initial_inputs, inputs], axis=1)\n",
    "\n",
    "    def _cond(i, inputs, kv_cache, stores):\n",
    "        return i < length\n",
    "\n",
    "    def _body(i, inputs, kv_cache, stores):\n",
    "        new_logits, new_kv_cache = serve_cache(inputs, kv_cache)\n",
    "        logits_with_temperature = new_logits[:, -1, :]\n",
    "        if temperature > 0.0:\n",
    "            logits_with_temperature /= temperature\n",
    "        new_inputs = top_k_top_p_sample(logits_with_temperature, 1, top_k, top_p)\n",
    "        new_stores = tf.concat([stores, new_inputs], axis=-1)\n",
    "        new_kv_cache = tf.concat([\n",
    "            kv_cache,\n",
    "            new_kv_cache\n",
    "        ], axis=-2)\n",
    "        new_i = i + 1\n",
    "        return [new_i, new_inputs, new_kv_cache, new_stores]\n",
    "\n",
    "    result = tf.while_loop(\n",
    "        _cond, _body,\n",
    "        loop_vars=[i, inputs, kv_cache, stores],\n",
    "        shape_invariants=[\n",
    "            tf.TensorShape(None),\n",
    "            tf.TensorShape([None, None]),\n",
    "            tf.TensorShape([\n",
    "                layer_size, None, 2,\n",
    "                attention_head, None,\n",
    "                embedding_size // attention_head\n",
    "            ]),\n",
    "            tf.TensorShape([\n",
    "                None, None\n",
    "            ])\n",
    "        ]\n",
    "    )\n",
    "    return result[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 263,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[837 259 497 788   8   9  16  29  91  14   8  12   8  10  16  18   8 616\n",
      "   89 219]], shape=(1, 20), dtype=int64)\n",
      "今天天气 不错 , 我 就 去 了 。   我 在 院子 里\n"
     ]
    }
   ],
   "source": [
    "ids = cbpe.encode('今天天气不错')\n",
    "\n",
    "ret = sample(\n",
    "    tf.constant([ids], dtype=tf.int64),\n",
    "    tf.constant(15, dtype=tf.int64),\n",
    "    tf.constant(1, dtype=tf.int32),\n",
    "    tf.constant(1.0, dtype=tf.float32),\n",
    "    tf.constant(1.0, dtype=tf.float32)\n",
    ")\n",
    "print(ret)\n",
    "print(cbpe.decode(ret.numpy().tolist()[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 270,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[ 837  259  497  788    8    9   51   35   25  182   11 4588   97  124\n",
      "    46   63    8    9  148   16]], shape=(1, 20), dtype=int64)\n",
      "今天天气 不错 , 但是 这次 的 天气 并 不是 很 好 , 所以 我\n"
     ]
    }
   ],
   "source": [
    "ids = cbpe.encode('今天天气不错')\n",
    "\n",
    "ret = sample(\n",
    "    tf.constant([ids], dtype=tf.int64),\n",
    "    tf.constant(15, dtype=tf.int64),\n",
    "    tf.constant(10, dtype=tf.int32),\n",
    "    tf.constant(0.95, dtype=tf.float32),\n",
    "    tf.constant(1.0, dtype=tf.float32)\n",
    ")\n",
    "print(ret)\n",
    "print(cbpe.decode(ret.numpy().tolist()[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 271,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:Found untraced functions such as attention_layer_call_fn, attention_layer_call_and_return_conditional_losses, LayerNorm_mlp_ln0_layer_call_fn, LayerNorm_mlp_ln0_layer_call_and_return_conditional_losses, LayerNorm_mlp_ln1_layer_call_fn while saving (showing 5 of 780). These functions will not be directly callable after loading.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: ./cpm-distill-tf2/assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: ./cpm-distill-tf2/assets\n"
     ]
    }
   ],
   "source": [
    "gpt.save('./cpm-distill-tf2', include_optimizer=False, signatures={\n",
    "    'serving_default': sample.get_concrete_function(\n",
    "        tf.TensorSpec(shape=[None, None], dtype=tf.int64, name=\"inp\"),\n",
    "        tf.TensorSpec(shape=[None,], dtype=tf.int64, name=\"length\"),\n",
    "        tf.TensorSpec(shape=[None,], dtype=tf.int32, name=\"top_k\"),\n",
    "        tf.TensorSpec(shape=[None,], dtype=tf.float32, name=\"top_p\"),\n",
    "        tf.TensorSpec(shape=[None,], dtype=tf.float32, name=\"temperature\")\n",
    "    )\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 272,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "424M\t./cpm-distill-tf2\r\n"
     ]
    }
   ],
   "source": [
    "!du -sh ./cpm-distill-tf2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}