StockEntityRecognitionUnmasked.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "treated-superintendent",
   "metadata": {},
   "source": [
    "# Stock Entity Recognition Unmasked"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "emerging-motel",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas\n",
    "import re\n",
    "import json\n",
    "import math\n",
    "import numpy\n",
    "import os\n",
    "import tensorflow as tf\n",
    "\n",
    "from itertools import chain\n",
    "from multiprocessing import Pool\n",
    "from functools import partial\n",
    "from transformers import TFBertForTokenClassification, BertTokenizerFast\n",
    "\n",
    "from sklearn.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "crucial-filing",
   "metadata": {},
   "source": [
    "## Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "authorized-aruba",
   "metadata": {},
   "outputs": [],
   "source": [
    "THREADS = 48\n",
    "\n",
    "DEVICE = '/cpu:0'\n",
    "\n",
    "TOKENS_AND_LABELS_TRAIN_LOC = 'data/ner/tokens_and_labels_train.parquet'\n",
    "TOKENS_AND_LABELS_TEST_LOC = 'data/ner/tokens_and_labels_test.parquet'\n",
    "MODEL_LOC = 'data/ner/unmasked/best.ckpt'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "headed-funeral",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "liked-weight",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>tokens</th>\n",
       "      <th>labels</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>29383</th>\n",
       "      <td>pretty sure USO is the closest to 1:1 with raw...</td>\n",
       "      <td>0 0 0 0 0 0 0 0 0 0 1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>109062</th>\n",
       "      <td>Everyone knows there’s 2 rules to follow: neve...</td>\n",
       "      <td>0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>117238</th>\n",
       "      <td>Ford has been mismanaged for years . They cost...</td>\n",
       "      <td>0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>36584</th>\n",
       "      <td>I own TERP . They will likely be bought out by...</td>\n",
       "      <td>0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>125747</th>\n",
       "      <td>Just did a MMM 162.5c 3/6 . More masks please!</td>\n",
       "      <td>0 0 0 1 0 0 0 0 0 0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>147873</th>\n",
       "      <td>Watch a deal happens , we moon , then JP gets ...</td>\n",
       "      <td>0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>161762</th>\n",
       "      <td>Look up [TEMPEST](https://youtu.be/APBSaJ5AA_c...</td>\n",
       "      <td>0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>72042</th>\n",
       "      <td>I had a bunch of SNDL I got cheap so I , too ,...</td>\n",
       "      <td>0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>138661</th>\n",
       "      <td>Taking financial advice from a convicted felon...</td>\n",
       "      <td>0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>67483</th>\n",
       "      <td>if you HAD TO choose , would you go with VOOor...</td>\n",
       "      <td>0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>142674 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                   tokens  \\\n",
       "29383   pretty sure USO is the closest to 1:1 with raw...   \n",
       "109062  Everyone knows there’s 2 rules to follow: neve...   \n",
       "117238  Ford has been mismanaged for years . They cost...   \n",
       "36584   I own TERP . They will likely be bought out by...   \n",
       "125747     Just did a MMM 162.5c 3/6 . More masks please!   \n",
       "...                                                   ...   \n",
       "147873  Watch a deal happens , we moon , then JP gets ...   \n",
       "161762  Look up [TEMPEST](https://youtu.be/APBSaJ5AA_c...   \n",
       "72042   I had a bunch of SNDL I got cheap so I , too ,...   \n",
       "138661  Taking financial advice from a convicted felon...   \n",
       "67483   if you HAD TO choose , would you go with VOOor...   \n",
       "\n",
       "                                                   labels  \n",
       "29383                               0 0 0 0 0 0 0 0 0 0 1  \n",
       "109062  0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 ...  \n",
       "117238  0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 ...  \n",
       "36584   0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 ...  \n",
       "125747                                0 0 0 1 0 0 0 0 0 0  \n",
       "...                                                   ...  \n",
       "147873  0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 ...  \n",
       "161762  0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 ...  \n",
       "72042   0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 ...  \n",
       "138661                0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0  \n",
       "67483                     0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0  \n",
       "\n",
       "[142674 rows x 2 columns]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokens_and_labels_train = pandas.read_parquet(TOKENS_AND_LABELS_TRAIN_LOC)\n",
    "tokens_and_labels_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "mounted-andrews",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')\n",
    "\n",
    "def encode_labels(max_length, chunk):\n",
    "    # loss function doesn't compute loss for labels if -100 so default to that\n",
    "    labels_encoded = numpy.ones((len(chunk), max_length), dtype=numpy.int8) * -100\n",
    "\n",
    "    # start filling in labels using the offset_mapping\n",
    "    for observation_i, (offset_mapping, labels_raw) in enumerate(chunk):\n",
    "        labels = numpy.array(labels_raw.split(' ')).astype(int)\n",
    "        label_i = 0\n",
    "        \n",
    "        for offset_i, offset in enumerate(offset_mapping):\n",
    "            if offset[0] == 0 and offset[1] != 0:\n",
    "                labels_encoded[observation_i][offset_i] = labels[label_i]\n",
    "                label_i += 1\n",
    "                \n",
    "    return labels_encoded\n",
    "\n",
    "def encode_df(df, max_length=256):\n",
    "    # encode everything\n",
    "    inputs_encoded = tokenizer(\n",
    "        # split ourselves so we can align with labels\n",
    "        list(map(lambda o: o.split(' '), df['tokens'])), \n",
    "        return_tensors=\"tf\",\n",
    "        is_split_into_words=True,\n",
    "        # offset mappings to align labels to the first word piece\n",
    "        return_offsets_mapping=True,\n",
    "        # make sure the same length across all encodings\n",
    "        max_length=max_length, \n",
    "        padding='max_length',\n",
    "        truncation=True\n",
    "    )\n",
    "    \n",
    "    offsets_with_labels = list(zip(inputs_encoded.offset_mapping.numpy(), df['labels']))\n",
    "    chunk_size = len(offsets_with_labels) / THREADS\n",
    "    offsets_with_labels_chunks = [offsets_with_labels[round(chunk_size  *i):round(chunk_size * (i + 1))] for i in range(0, THREADS)]\n",
    "    \n",
    "    with Pool(THREADS) as pool:\n",
    "        encoded_labels = pool.map(partial(encode_labels, max_length), offsets_with_labels_chunks)\n",
    "        \n",
    "    return inputs_encoded, numpy.stack(list(chain(*encoded_labels)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "medium-appliance",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoded = encode_df(tokens_and_labels_train)\n",
    "\n",
    "tf_train_dataset = tf.data.Dataset.from_tensor_slices(({\n",
    "    'input_ids': encoded[0]['input_ids'],\n",
    "    'token_type_ids': encoded[0]['token_type_ids'],\n",
    "    'attention_mask': encoded[0]['attention_mask']\n",
    "}, encoded[1])).batch(32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "looking-anthropology",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "All model checkpoint layers were used when initializing TFBertForTokenClassification.\n",
      "\n",
      "Some layers of TFBertForTokenClassification were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['classifier']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"tf_bert_for_token_classification\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "bert (TFBertMainLayer)       multiple                  108891648 \n",
      "_________________________________________________________________\n",
      "dropout_37 (Dropout)         multiple                  0         \n",
      "_________________________________________________________________\n",
      "classifier (Dense)           multiple                  1538      \n",
      "=================================================================\n",
      "Total params: 108,893,186\n",
      "Trainable params: 108,893,186\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n",
      "Epoch 1/5\n",
      "WARNING:tensorflow:The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "WARNING:tensorflow:AutoGraph could not transform <bound method Socket.send of <zmq.sugar.socket.Socket object at 0x7f26c953a280>> and will run it as-is.\n",
      "Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.\n",
      "Cause: module, class, method, function, traceback, frame, or code object was expected, got cython_function_or_method\n",
      "To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert\n",
      "WARNING: AutoGraph could not transform <bound method Socket.send of <zmq.sugar.socket.Socket object at 0x7f26c953a280>> and will run it as-is.\n",
      "Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.\n",
      "Cause: module, class, method, function, traceback, frame, or code object was expected, got cython_function_or_method\n",
      "To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert\n",
      "WARNING:tensorflow:The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "WARNING:tensorflow:From /home/zuyezheng/.pyenv/versions/reddit/lib/python3.8/site-packages/tensorflow/python/ops/array_ops.py:5043: calling gather (from tensorflow.python.ops.array_ops) with validate_indices is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "The `validate_indices` argument has no effect. Indices are always validated on CPU and never validated on GPU.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/zuyezheng/.pyenv/versions/reddit/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py:376: UserWarning: Using `-1` to mask the loss for the token is deprecated. Please use `-100` instead.\n",
      "  return py_builtins.overload_of(f)(*args)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "WARNING:tensorflow:The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n",
      "4459/4459 [==============================] - 1558s 347ms/step - loss: 0.0135 - accuracy: 0.2101\n",
      "\n",
      "Epoch 00001: saving model to data/ner/unmasked/best.ckpt\n",
      "Epoch 2/5\n",
      "4459/4459 [==============================] - 1505s 338ms/step - loss: 0.0026 - accuracy: 0.2109\n",
      "\n",
      "Epoch 00002: saving model to data/ner/unmasked/best.ckpt\n",
      "Epoch 3/5\n",
      "4459/4459 [==============================] - 1510s 339ms/step - loss: 0.0021 - accuracy: 0.2110\n",
      "\n",
      "Epoch 00003: saving model to data/ner/unmasked/best.ckpt\n",
      "Epoch 4/5\n",
      "4459/4459 [==============================] - 1508s 338ms/step - loss: 0.0020 - accuracy: 0.2110\n",
      "\n",
      "Epoch 00004: saving model to data/ner/unmasked/best.ckpt\n",
      "Epoch 5/5\n",
      "4459/4459 [==============================] - 1509s 338ms/step - loss: 0.0015 - accuracy: 0.2110\n",
      "\n",
      "Epoch 00005: saving model to data/ner/unmasked/best.ckpt\n",
      "CPU times: user 1h 24min 27s, sys: 10min 15s, total: 1h 34min 43s\n",
      "Wall time: 2h 6min 38s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.callbacks.History at 0x7f2190887bb0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "model = TFBertForTokenClassification.from_pretrained('bert-base-uncased', num_labels = 2)\n",
    "\n",
    "optimizer = tf.optimizers.Adam(learning_rate=5e-5)\n",
    "model.compile(optimizer=optimizer, loss=model.compute_loss, metrics=['accuracy'])\n",
    "model.summary()\n",
    "\n",
    "cp_callback = tf.keras.callbacks.ModelCheckpoint(\n",
    "    filepath=MODEL_LOC,\n",
    "    save_weights_only=True,\n",
    "    verbose=1\n",
    ")\n",
    "\n",
    "model.fit(tf_train_dataset, epochs=5, callbacks=[cp_callback])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "alert-construction",
   "metadata": {},
   "source": [
    "## Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "silver-bristol",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "All model checkpoint layers were used when initializing TFBertForTokenClassification.\n",
      "\n",
      "Some layers of TFBertForTokenClassification were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['classifier']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x7f23fa078f40>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load up the tokenizer and model\n",
    "tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')\n",
    "model = TFBertForTokenClassification.from_pretrained('bert-base-uncased', num_labels = 2)\n",
    "model.load_weights(MODEL_LOC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "alike-regulation",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(sentence, max_length=256):\n",
    "    tokens = sentence.split(' ')\n",
    "\n",
    "    test_encoding = tokenizer(\n",
    "        tokens,\n",
    "        return_tensors=\"tf\",\n",
    "        is_split_into_words=True,\n",
    "        max_length=max_length, \n",
    "        padding='max_length',\n",
    "        truncation=True,\n",
    "        return_offsets_mapping=True\n",
    "    )\n",
    "\n",
    "    # grab the offset mappings\n",
    "    offset_mapping = test_encoding.offset_mapping\n",
    "    del test_encoding['offset_mapping']\n",
    "    test_encoding\n",
    "\n",
    "    prediction = tf.argsort(model(test_encoding).logits[0])\n",
    "\n",
    "    token_predictions = []\n",
    "    \n",
    "    num_tokens = len(test_encoding.attention_mask[0][test_encoding.attention_mask[0] == 1])\n",
    "    token_i = 0\n",
    "    \n",
    "    for i in range(num_tokens):\n",
    "        offset = offset_mapping[0][i]\n",
    "        token_prediction = prediction[i]\n",
    "\n",
    "        if offset[0] == 0 and offset[1] != 0:\n",
    "            token_predictions.append([tokens[token_i], bool(token_prediction[1] == 1)])\n",
    "            token_i += 1\n",
    "            \n",
    "    return token_predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "moving-strip",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[['my', False], ['msft', False], ['is', False], ['crm', True]]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "with tf.device(DEVICE):\n",
    "    display(predict('my msft is crm'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "overall-texture",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "reddit",
   "language": "python",
   "name": "reddit"
  },
  "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": 5
}