huffman.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-20T14:05:21.752713Z",
     "start_time": "2020-08-20T14:05:20.626026Z"
    }
   },
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "from heapq import *\n",
    "from collections import namedtuple\n",
    "import numpy as np\n",
    "import torch.nn.functional as F\n",
    "import torch\n",
    "import re\n",
    "\n",
    "# load some text\n",
    "with open('shakespeare.txt') as f:\n",
    "    text = f.read()\n",
    "\n",
    "# get rid of a few characters that break test coverage\n",
    "text = re.sub('[$&3]', '', text)\n",
    "test_split_idx = len(text)//3 * 2\n",
    "train_text = text[:test_split_idx]\n",
    "test_text = text[test_split_idx:]\n",
    "assert set(train_text) == set(test_text)\n",
    "\n",
    "# something short for testing\n",
    "short_test_text = test_text[13:1000]\n",
    "\n",
    "\n",
    "stoi = {char: i for i,char in enumerate(set(train_text))}\n",
    "# remove this\n",
    "stoi['>'] = len(stoi)\n",
    "itos = {i: char for char, i in stoi.items()}\n",
    "\n",
    "\n",
    "n_char_train = len(train_text)\n",
    "\n",
    "abs_frequencies = Counter(train_text)\n",
    "frequencies = {k: v/n_char_train for k,v in abs_frequencies.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build Huffman code\n",
    "This is done using a priority queue. The two elements with the lowest counts are merged and inserted back into the queue with added counts.\n",
    "The most frequent characters will get the shortest codes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-20T14:05:22.760277Z",
     "start_time": "2020-08-20T14:05:22.750728Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'F': '011011011', 'i': '10111', 'r': '0000', 's': '0001', 't': '1000', ' ': '110', 'C': '01101100', 'z': '111101010110', 'e': '1110', 'n': '11111', ':': '1001111', '\\n': '10100', 'B': '101100110', 'f': '011010', 'o': '0111', 'w': '100110', 'p': '1011000', 'c': '010111', 'd': '01100', 'a': '0100', 'y': '101011', 'u': '00110', 'h': '0010', ',': '101010', 'm': '101101', 'k': '0101001', '.': '0101011', 'A': '0101010', 'l': '10010', 'S': '01101111', 'Y': '011011010', 'v': '0101101', '?': '100111001', 'R': '11110100', 'M': '111101011', 'W': '01011001', \"'\": '11110110', 'L': '01010000', 'I': '1111001', 'N': '10110010', 'g': '001111', ';': '01010001', 'b': '1111000', '!': '011011101', 'O': '10011101', 'j': '11110101001', 'V': '11110101010', '-': '010110001', 'T': '0011100', 'H': '00111010', 'E': '11110111', 'U': '00111011', 'D': '100111000', 'P': '0110111000', 'q': '01101110011', 'x': '01101110010', 'J': '111101010111', 'G': '101100111', 'K': '010110000', 'Q': '111101010001', 'Z': '1111010100001', 'X': '1111010100000'}\n"
     ]
    }
   ],
   "source": [
    "Node = namedtuple('Node', ['left', 'right', 'chars'])\n",
    "\n",
    "def huffman_coding(frequencies):\n",
    "    \"\"\"\n",
    "    Builds a Huffman tree:\n",
    "    Args:\n",
    "        -frequences: Dictionary of characters and their frequencies\n",
    "    Returns:\n",
    "        codes: dictonary from characters to their binary codes for encoding\n",
    "        root: root node of the Huffman tree, for decoding\n",
    "    \"\"\"\n",
    "    codes = {k: '' for k in frequencies}\n",
    "\n",
    "    h = []\n",
    "    for char, freq in frequencies.items():\n",
    "        heappush(h, (freq, Node(None, None, char)))\n",
    "\n",
    "    while True:\n",
    "        freq1, node1 = heappop(h)\n",
    "        if not h:\n",
    "            break\n",
    "        freq2, node2 = heappop(h)\n",
    "\n",
    "        for char in node1.chars:\n",
    "            codes[char] += ('0')\n",
    "        for char in node2.chars:\n",
    "            codes[char] += ('1')\n",
    "\n",
    "        parent = Node(node1, node2, node1.chars+node2.chars)\n",
    "        parent_freq = freq1+freq2\n",
    "        heappush(h, (parent_freq, parent))\n",
    "\n",
    "    codes = {k: v[::-1] for k,v in codes.items()}\n",
    "    root = parent\n",
    "\n",
    "    return codes, root\n",
    "\n",
    "codes, root = huffman_coding(frequencies)\n",
    "print(codes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-20T13:48:43.074805Z",
     "start_time": "2020-08-20T13:48:43.068586Z"
    }
   },
   "source": [
    "## One-gram model\n",
    "Here we see that the Huffman code gives us a bit of advantage over constant length coding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-20T14:05:23.899210Z",
     "start_time": "2020-08-20T14:05:23.892551Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Entropy:  4.7833957052643425\n",
      "Expected code length:  4.82071185920971\n",
      "Constant code length:  5.954196310386875\n"
     ]
    }
   ],
   "source": [
    "expected_code_length = sum(len(codes[k]) * frequencies[k] for k in frequencies)\n",
    "expected_code_length\n",
    "\n",
    "frequencies_np = np.array(list(frequencies.values()))\n",
    "entropy = -(np.log2(frequencies_np) * frequencies_np).sum()\n",
    "\n",
    "print('Entropy: ', entropy)\n",
    "print('Expected code length: ', expected_code_length)\n",
    "print('Constant code length: ', np.log2(len(frequencies)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-20T14:05:24.717080Z",
     "start_time": "2020-08-20T14:05:24.710765Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Encoded length:  4699\n",
      "0101010000111010110000100000100011110000111001011101111001001110011000001011111111001111010101110100 ...\n"
     ]
    }
   ],
   "source": [
    "# encoding is easy here. Just join the codes for each character in the text.\n",
    "def encode(text, codes):\n",
    "    '''Encode text using the code we created\n",
    "    '''\n",
    "    return ''.join(codes[k] for k in text)\n",
    "\n",
    "encoded_text = encode(short_test_text, codes)\n",
    "\n",
    "# assert that encoded text is binary\n",
    "assert len(set(encoded_text)) == 2\n",
    "print(f'Encoded length: ', len(encoded_text))\n",
    "print(encoded_text[:100], '...')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-20T14:05:26.335742Z",
     "start_time": "2020-08-20T14:05:26.331076Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BPC:  4.760891590678825\n"
     ]
    }
   ],
   "source": [
    "def decode(encoded_text, root):\n",
    "    \"\"\"Decoding encoded text given a huffman tree. \n",
    "    Goes down the tree till a leave is reached\"\"\"\n",
    "    decoded_text = ''\n",
    "    \n",
    "    current_node = root\n",
    "    for bit in encoded_text:\n",
    "        if current_node.left is None:\n",
    "            decoded_text += current_node.chars\n",
    "            current_node = root\n",
    "        if bit == '0':\n",
    "            current_node = current_node.left\n",
    "        else:\n",
    "            current_node = current_node.right\n",
    "    \n",
    "    decoded_text += current_node.chars\n",
    "    return decoded_text\n",
    "\n",
    "# now we decode the text and check if we can reconstruct the input\n",
    "decoded_text = decode(encoded_text, root)\n",
    "assert decoded_text == short_test_text\n",
    "\n",
    "# this should be in the range of the entropy\n",
    "print('BPC: ', len(encoded_text) / len(short_test_text))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LSTM encoding/decoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-20T14:29:04.724091Z",
     "start_time": "2020-08-20T14:29:04.718955Z"
    }
   },
   "outputs": [],
   "source": [
    "# define the charRNN\n",
    "# TODO: test set, make it actually generalize\n",
    "\n",
    "class CharRNN(torch.nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, num_layers, batch_first=True):\n",
    "        super(CharRNN, self).__init__()\n",
    "        self.lstm = torch.nn.LSTM(input_size, hidden_size, num_layers, batch_first=batch_first)\n",
    "        self.linear = torch.nn.Linear(hidden_size, input_size)\n",
    "    \n",
    "    def forward(self, inputs, hidden=None):\n",
    "        x = F.one_hot(inputs, num_classes=len(stoi)).float()\n",
    "        x, hidden = self.lstm(x, hidden)\n",
    "        x = self.linear(x)\n",
    "        return x, hidden\n",
    "\n",
    "n_block = 64\n",
    "net = CharRNN(len(stoi), 64, 2)\n",
    "opt = torch.optim.Adam(net.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-20T14:41:32.776125Z",
     "start_time": "2020-08-20T14:40:22.656304Z"
    },
    "collapsed": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.6810500990338373\r"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-66-f3df025ca487>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     12\u001b[0m     \u001b[0mlogits\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhidden\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnet\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m     \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcross_entropy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlogits\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mview\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlogits\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtargets\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreshape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m     \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     15\u001b[0m     \u001b[0mopt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m     \u001b[0;31m# convert loss to bits\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/torch/tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph)\u001b[0m\n\u001b[1;32m    148\u001b[0m                 \u001b[0mproducts\u001b[0m\u001b[0;34m.\u001b[0m \u001b[0mDefaults\u001b[0m \u001b[0mto\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    149\u001b[0m         \"\"\"\n\u001b[0;32m--> 150\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    151\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    152\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mregister_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables)\u001b[0m\n\u001b[1;32m     97\u001b[0m     Variable._execution_engine.run_backward(\n\u001b[1;32m     98\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_tensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 99\u001b[0;31m         allow_unreachable=True)  # allow_unreachable flag\n\u001b[0m\u001b[1;32m    100\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    101\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# train it\n",
    "for i in range(5000):\n",
    "    indices = np.random.randint(n_char_train - n_block, size=128)\n",
    "    batch_clear = [train_text[i:i+n_block] for i in indices]\n",
    "    batch_input = [sample[:-1] for sample in batch_clear]\n",
    "    batch_target = [sample[1:] for sample in batch_clear]\n",
    "    \n",
    "    inputs = torch.tensor([[stoi[c] for c in sample] for sample in batch_input])\n",
    "    targets = torch.tensor([[stoi[c] for c in sample] for sample in batch_target])\n",
    "    \n",
    "    opt.zero_grad()\n",
    "    logits, hidden = net(inputs)\n",
    "    loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.reshape(-1))\n",
    "    loss.backward()\n",
    "    opt.step()\n",
    "    # convert loss to bits\n",
    "    print(loss.item() / np.log(2), end='\\r')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-20T14:41:35.334698Z",
     "start_time": "2020-08-20T14:41:35.331802Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.6341760555585276\r"
     ]
    }
   ],
   "source": [
    "print(loss.item() / np.log(2), end='\\r')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-20T14:41:41.893156Z",
     "start_time": "2020-08-20T14:41:36.118742Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2961"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoded_text = ''\n",
    "# first letter must be encoded without context\n",
    "codes, root = huffman_coding(frequencies)\n",
    "encoded_text += codes[short_test_text[0]]\n",
    "\n",
    "text_num = torch.tensor([[stoi[c] for c in short_test_text[:-1]]])\n",
    "logits, _ = net(text_num)\n",
    "probs = F.softmax(logits, -1).squeeze()\n",
    "\n",
    "for i in range(probs.shape[0]):\n",
    "    frequencies_step = probs[i]\n",
    "    frequencies_step = {itos[j]: frequencies_step[j] for j in range(frequencies_step.shape[0])}\n",
    "    codes, root = huffman_coding(frequencies_step)\n",
    "    encoded_text += codes[short_test_text[i+1]]\n",
    "    \n",
    "len(encoded_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-20T14:41:42.609822Z",
     "start_time": "2020-08-20T14:41:41.894411Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "As passes colouring.\n",
      "Dear gentlewoman,\n",
      "How fares our gracious lady?\n",
      "\n",
      "EMILIA:\n",
      "As well as one so great\n",
      "3.0\n"
     ]
    }
   ],
   "source": [
    "def decode_char(encoded_text, root):\n",
    "    current_node = root\n",
    "    for bit in encoded_text:\n",
    "        if bit == '0':\n",
    "            current_node = current_node.left\n",
    "        elif bit == '1':\n",
    "            current_node = current_node.right\n",
    "        else:\n",
    "            raise ValueError\n",
    "        \n",
    "        if current_node.left is None:\n",
    "            return current_node.chars\n",
    "    return '!'\n",
    "# decode first letter\n",
    "index = 0\n",
    "decoded_text = ''\n",
    "codes, root = huffman_coding(frequencies)\n",
    "current_char = decode_char(encoded_text[index:], root)\n",
    "\n",
    "decoded_text += current_char\n",
    "index += len(codes[current_char])\n",
    "\n",
    "hidden = None\n",
    "while True:\n",
    "    batch_num = torch.tensor(stoi[current_char]).view(1,1)\n",
    "    logits, hidden = net(batch_num, hidden)\n",
    "    probs = F.softmax(logits, -1).squeeze()\n",
    "    frequencies_step = {itos[j]: probs[j].item() for j in range(probs.shape[0])}\n",
    "    codes, root = huffman_coding(frequencies_step)\n",
    "    current_char = decode_char(encoded_text[index:], root)\n",
    "    \n",
    "    decoded_text += current_char\n",
    "    index += len(codes[current_char])\n",
    "    \n",
    "    if index >= len(encoded_text):\n",
    "        break\n",
    "\n",
    "print(decoded_text[:100])\n",
    "\n",
    "print(len(encoded_text) / len(short_test_text))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-21T07:53:40.517129Z",
     "start_time": "2020-08-21T07:53:40.513949Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6.314083080040527"
      ]
     },
     "execution_count": 73,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "779 * 8 / len(short_test_text)"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}