examples/bert.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BERT representation storage and loading (with sharding)\n",
    "\n",
    "In other examples, the data stored with `seqp` were discrete symbols from a closed set, like text tokens (with finite word-based vocabulary) and DNA nucleotide strings.\n",
    "\n",
    "In this notebook, we are going to store floating point numbers. More specifically, we are going to use [BERT](https://github.com/huggingface/pytorch-pretrained-BERT) to encode pieces of text as sequences of contextual token representations.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interface to BERT\n",
    "\n",
    "We will be using Hugging Face's port of BERT to Pytorch. Nevertheless, we will prepare a convenient wrapper to easily interface BERT. For this wrapper, we will use `seqp`'s `TextCodec` as base class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Better speed can be achieved with apex installed from https://www.github.com/nvidia/apex.\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from pytorch_pretrained_bert import BertTokenizer, BertForMaskedLM\n",
    "from seqp.encoding import TextCodec\n",
    "import torch\n",
    "from typing import Union, List, Optional\n",
    "\n",
    "\n",
    "DEFAULT_BERT_WEIGHTS = 'bert-base-multilingual-cased'\n",
    "\n",
    "class BertInterface(TextCodec):\n",
    "    def __init__(self, use_gpu=False):\n",
    "        self.tokenizer = BertTokenizer.from_pretrained(DEFAULT_BERT_WEIGHTS)\n",
    "        self.model = BertForMaskedLM.from_pretrained(DEFAULT_BERT_WEIGHTS)\n",
    "        self.model.eval()\n",
    "        use_gpu = use_gpu and torch.cuda.is_available()\n",
    "        self.device = torch.device(\"cuda\" if use_gpu else \"cpu\")\n",
    "        self.model.to(self.device)\n",
    "\n",
    "    def decode(self, embedded: Union[np.ndarray, torch.Tensor]) -> List[str]:\n",
    "        if isinstance(embedded, np.ndarray):\n",
    "            if len(embedded.shape) == 2:  # seq_length x emb_dim\n",
    "                embedded = np.expand_dims(embedded, 0)  # add batch dimension\n",
    "            assert len(embedded.shape) == 3\n",
    "            embedded = torch.from_numpy(embedded).to(self.device)\n",
    "        predictions = self.model.cls(embedded)\n",
    "        predicted_indexes = torch.argmax(predictions, dim=2).cpu().numpy()\n",
    "        predicted_tokens = self.tokenizer.convert_ids_to_tokens(predicted_indexes[0].tolist())\n",
    "        return predicted_tokens\n",
    "\n",
    "    def detokenize(self, tokens: List[str]) -> str:\n",
    "        return \" \".join(tokens).replace(\" ##\", \"\")\n",
    "\n",
    "    def tokenize(self, sentence: str) -> List[str]:\n",
    "        return self.tokenizer.tokenize(sentence)\n",
    "\n",
    "    def encode(self, tokens: List[str]) -> Optional[np.ndarray]:\n",
    "        tokenized_text = ['[CLS]'] + tokens\n",
    "        if len(tokenized_text) > self.tokenizer.max_len:\n",
    "            return None\n",
    "\n",
    "        indexed_tokens = self.tokenizer.convert_tokens_to_ids(tokenized_text)\n",
    "        tokens_tensor = torch.LongTensor([indexed_tokens]).to(self.device)\n",
    "        sequence_output, _ = self.model.bert(tokens_tensor, output_all_encoded_layers=False)\n",
    "        return sequence_output.detach()[0].cpu().numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are many details about BERT and Pytorch that we are encapsulated in this wrapper. They are not important for this example, but if you want to know more and don't know how to start with:\n",
    "- About BERT: please read [the original article](https://arxiv.org/abs/1810.04805) or one of the many tutorials about it, like [The Illustrated BERT](http://jalammar.github.io/illustrated-bert/), and please have a look at [Hugging Face's wonderful port](https://github.com/huggingface/pytorch-pretrained-BERT).\n",
    "- About Pytorch: I recommend the book [Natural Language Processing with PyTorch\n",
    "](http://shop.oreilly.com/product/0636920063445.do).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Encode BERT representations and persist them to sharded HDF5 files\n",
    "\n",
    "First we will download a text file (the Universal Declaration of Human Rights):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -q http://research.ics.aalto.fi/cog/data/udhr/txt/eng.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we will iterate over the lines in the file, encoding the text as contextual vector representations of each token, and we will save the sequences as `seqp` records:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from seqp.hdf5 import Hdf5RecordWriter\n",
    "from seqp.record import ShardedWriter\n",
    "\n",
    "input_file = 'eng.txt'\n",
    "output_file_template = \"bert_example_{:02d}.hdf5\"\n",
    "\n",
    "bert = BertInterface()\n",
    "\n",
    "with ShardedWriter(Hdf5RecordWriter,\n",
    "                   output_file_template,\n",
    "                   max_records_per_shard=10) as writer, open(input_file) as f:\n",
    "    for line in f:\n",
    "        line = line.strip()\n",
    "        tokens = bert.tokenize(line)\n",
    "        ctx_representations = bert.encode(tokens)\n",
    "        writer.write(ctx_representations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read representations back and decode them\n",
    "\n",
    "Now, we will use a `Hdf5RecordReader` to read back the token vector representations and will decode them back into tokens by means of our BERT interface."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vector representation shape: (338, 768)\n",
      "Sentence: . .amble and recognition of the inherent dignity and of the equal and inalienable rights of all members of the human family is the foundation of freedom , justice and peace in the world , whereas disregard and contempt for human rights have resulted in barbarous acts which have outraged the conscience of mankind , and the advent of a world in which human beings shall enjoy freedom of speech and belief and freedom from fear and want has been proclaimed as the highest aspiration of the common people , whereas it is essential , if man is not to be compelled to have recourse , as a last resort , to rebellion against tyranny and oppression , that human rights should be protected by the rule of law , whereas it is essential to promote the development of friendly relations between nations , whereas the peoples of the united nations have in the charter reaffirmed their faith in fundamental human rights , in the dignity and worth of the human person and in the equal rights of men and women and have determined to promote social progress and better standards of life in larger freedom , whereas member states have pledged themselves to achieve , in cooperation with the united nations , the promotion of universal respect for and observance of human rights and fundamental freedoms , whereas a common understanding of these rights and freedoms is of the greatest importance for the full realization of this pledge , now , therefore , the general assembly , proclaims this universal declaration of human rights as a common standard of achievement for all peoples and all nations , to the end that every individual and every organ of society , keeping this declaration constantly in mind , shall strive by\n"
     ]
    }
   ],
   "source": [
    "from glob import glob\n",
    "from seqp.hdf5 import Hdf5RecordReader\n",
    "\n",
    "with Hdf5RecordReader(glob('bert_example_*.hdf5')) as reader:\n",
    "    indexes = list(reader.indexes())\n",
    "    ctx_rep = reader.retrieve(indexes[9])\n",
    "    print(\"Vector representation shape: {}\".format(ctx_rep.shape))\n",
    "    sentence = bert.detokenize(bert.decode(ctx_rep))\n",
    "    print(\"Sentence: {}\".format(sentence))\n",
    "    "
   ]
  }
 ],
 "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.6.8"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "source": [],
    "metadata": {
     "collapsed": false
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}