BERT Embeddings (on GPU).ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import logging\n",
    "import pickle\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from math import ceil\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "from sklearn.linear_model import LinearRegression\n",
    "\n",
    "from bert_serving.client import BertClient\n",
    "from bert_serving.server.graph import optimize_graph\n",
    "from bert_serving.server.helper import get_args_parser\n",
    "from bert_serving.server.bert.tokenization import FullTokenizer\n",
    "from bert_serving.server.bert.extract_features import convert_lst_to_features\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow.python.estimator.estimator import Estimator\n",
    "from tensorflow.python.estimator.run_config import RunConfig\n",
    "from tensorflow.python.estimator.model_fn import EstimatorSpec\n",
    "from tensorflow.keras.utils import Progbar\n",
    "\n",
    "log = logging.getLogger('tensorflow')\n",
    "log.setLevel(logging.INFO)\n",
    "log.handlers = []"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BERT Embeddings on GPU\n",
    "\n",
    "This notebook uses `Bert-as-a-Service` and `TensorFlow` to extract embeddings from tokens and documents. This is part of a larger project to use these embeddings to calculate Semantic Similarity to clean search results in user-generated data. Please visit `Embedding Similarity Across Models` to see this done (along with other models)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_inference_speed(seq_len):\n",
    "    lr = LinearRegression()\n",
    "    lr.fit(X = np.array([4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]).reshape(-1, 1),\n",
    "           y = np.array([0.0214, 0.0236, 0.0257, 0.0277, 0.0289, 0.0316, 0.0344, 0.0355, 0.0422, 0.0429, 0.0454]).reshape(-1, 1)\n",
    "          );\n",
    "    return lr.predict(np.array([[seq_len]]))[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Hours: 9.826'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f'Hours: {round(((predict_inference_speed(seq_len = 22)) * 550000) / 60 / 60, 3)}'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build / Export Graph\n",
    "\n",
    " - **Models**:\n",
    "     - Token Model Seq Len: 4\n",
    "     - Doc Model Seq Len: 22 (99.9% coverage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_DIR = 'models/wwm_uncased_L-24_H-1024_A-16'\n",
    "GRAPH_DIR = 'models/graph'\n",
    "NUM_WORKER = \"8\"                                # concurrency\n",
    "POOL_LAYER = \"-2\"                               # second to last layer (suggested as last layer is too biased on target, see Bert-as-a-Service docs for more info)\n",
    "POOL_STRAT = 'REDUCE_MEAN'                      # averages vectors for all tokens in a sequence\n",
    "SEQ_LEN = 22                                    # in tokens, changing this value linearly affects inference speed\n",
    "GRAPH_OUT = f'extractor_seq-len-{SEQ_LEN}.pbtxt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.gfile.MkDir(GRAPH_DIR)\n",
    "\n",
    "parser = get_args_parser()\n",
    "carg = parser.parse_args(args=['-model_dir', MODEL_DIR,\n",
    "                               \"-graph_tmp_dir\", GRAPH_DIR,\n",
    "                               '-max_seq_len', str(SEQ_LEN),\n",
    "                               '-pooling_layer', str(POOL_LAYER),\n",
    "                               '-pooling_strategy', str(POOL_STRAT),\n",
    "                               '-num_worker', str(NUM_WORKER),\n",
    "                              ])\n",
    "\n",
    "tmpfi_name, config = optimize_graph(carg)\n",
    "graph_fout = os.path.join(GRAPH_DIR, GRAPH_OUT)\n",
    "\n",
    "tf.gfile.Rename(\n",
    "    tmpfi_name,\n",
    "    graph_fout,\n",
    "    overwrite=True\n",
    ")\n",
    "print(\"Serialized graph to {}\".format(graph_fout))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Feature Extractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "SEQ_LEN = int(SEQ_LEN)\n",
    "GRAPH_PATH = f'models/graph/extractor_seq-len-{SEQ_LEN}.pbtxt'\n",
    "VOCAB_PATH = 'models/wwm_uncased_L-24_H-1024_A-16/vocab.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "INPUT_NAMES = ['input_ids', 'input_mask', 'input_type_ids']\n",
    "bert_tokenizer = FullTokenizer(VOCAB_PATH)\n",
    "\n",
    "def build_feed_dict(texts):\n",
    "    \n",
    "    text_features = list(\n",
    "        convert_lst_to_features(\n",
    "            lst_str = texts,\n",
    "            max_seq_length = SEQ_LEN,\n",
    "            max_position_embeddings = SEQ_LEN,\n",
    "            tokenizer = bert_tokenizer,\n",
    "            logger = log,\n",
    "            is_tokenized = False,\n",
    "            mask_cls_sep = False\n",
    "    ))\n",
    "\n",
    "    target_shape = (len(texts), -1)\n",
    "\n",
    "    feed_dict = {}\n",
    "    for iname in INPUT_NAMES:\n",
    "        features_i = np.array([getattr(f, iname) for f in text_features])\n",
    "        features_i = features_i.reshape(target_shape)\n",
    "        features_i = features_i.astype(\"int32\")\n",
    "        feed_dict[iname] = features_i\n",
    "\n",
    "    return feed_dict\n",
    "\n",
    "def build_input_fn(container):\n",
    "    \n",
    "    def gen():\n",
    "        while True:\n",
    "            try:\n",
    "                yield build_feed_dict(container.get())\n",
    "            except:\n",
    "                yield build_feed_dict(container.get())\n",
    "    \n",
    "    def input_fn():\n",
    "        return tf.data.Dataset.from_generator(\n",
    "            gen,\n",
    "            output_types = {iname: tf.int32 for iname in INPUT_NAMES},\n",
    "            output_shapes = {iname: (None, None) for iname in INPUT_NAMES})\n",
    "    \n",
    "    return input_fn\n",
    "\n",
    "class DataContainer:\n",
    "    \n",
    "    def __init__(self):\n",
    "        self._texts = None\n",
    "    \n",
    "    def set(self, texts):\n",
    "        if type(texts) is str:\n",
    "            texts = [texts]\n",
    "        self._texts = texts\n",
    "        \n",
    "    def get(self):\n",
    "        return self._texts\n",
    "    \n",
    "def model_fn(features, mode):\n",
    "    \n",
    "    with tf.gfile.GFile(GRAPH_PATH, 'rb') as f:\n",
    "        graph_def = tf.GraphDef()\n",
    "        graph_def.ParseFromString(f.read())\n",
    "    \n",
    "    output = tf.import_graph_def(\n",
    "        graph_def,\n",
    "        input_map={k + ':0': features[k] for k in INPUT_NAMES},\n",
    "        return_elements=['final_encodes:0']\n",
    "    )\n",
    "    \n",
    "    return EstimatorSpec(mode=mode, predictions={'output': output[0]})\n",
    "\n",
    "def batch(iterable, n=1):\n",
    "    l = len(iterable)\n",
    "    for ndx in range(0, l, n):\n",
    "        yield iterable[ndx:min(ndx + n, l)]\n",
    "\n",
    "def build_vectorizer(_estimator, _input_fn_builder, batch_size=128):\n",
    "    container = DataContainer()\n",
    "    predict_fn = _estimator.predict(_input_fn_builder(container), yield_single_examples=False)\n",
    "    \n",
    "    def vectorize(text, verbose=False):\n",
    "        x = []\n",
    "        bar = Progbar(len(text))\n",
    "        \n",
    "        for text_batch in batch(text, batch_size):\n",
    "            container.set(text_batch)\n",
    "            x.append(next(predict_fn)['output'])\n",
    "            if verbose:\n",
    "                bar.add(len(text_batch))\n",
    "        \n",
    "        r = np.vstack(x)\n",
    "        return r\n",
    "    \n",
    "    return vectorize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# init and build embedding generator\n",
    "estimator = Estimator(model_fn = model_fn)\n",
    "bert_vectorizer = build_vectorizer(estimator, build_input_fn)\n",
    "_ = bert_vectorizer(['music']); del _"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get Embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('all_queries_and_plns.pkl', 'rb') as f:\n",
    "    all_queries_and_plns = pickle.load(f)\n",
    "\n",
    "all_tokens = list(set([t for tt in [token.split() for token in all_queries_and_plns] for t in tt]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "token_embeds = bert_vectorizer(all_tokens)\n",
    "\n",
    "token_embed_dict = {token: embed.reshape(1, -1) for token, embed in zip(all_tokens, token_embeds)}\n",
    "\n",
    "with open('bert_token_embeds.pkl', 'wb') as f:\n",
    "    pickle.dump(token_embed_dict, f, protocol = pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_emb_dict(word_list):\n",
    "    return {w: e.reshape(1, -1) for w, e in zip(word_list, bert_vectorizer(word_list))}\n",
    "\n",
    "def get_doc_embeds(batches_ = 10):\n",
    "    write_fname = 'bert_doc_embeds.pkl'\n",
    "    batch_size_ = ceil(len(all_queries_and_plns) / batches_)\n",
    "\n",
    "    saved_emb = {}\n",
    "\n",
    "    for i in range(batches_):\n",
    "\n",
    "        print(f'Batch: {i+1}')\n",
    "        print(f'Saved Emb Size: {len(saved_emb)}')\n",
    "        print()\n",
    "\n",
    "        batch_ = list(all_queries_and_plns[batch_size_*i:batch_size_*(i+1)])\n",
    "        saved_emb = {**saved_emb, **get_emb_dict(batch_)}\n",
    "\n",
    "        with open(write_fname, 'wb') as f:\n",
    "            pickle.dump(saved_emb, f, protocol = pickle.HIGHEST_PROTOCOL)\n",
    "    \n",
    "    return saved_emb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_embeds = get_doc_embeds()"
   ]
  },
  {
   "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.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}