99_ported_hlda_sampler.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ported code!\n",
    "\n",
    "> We are porting to Python 3.x the code for Hierarchical LDA available here: https://github.com/joewandy/hlda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp sampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export \n",
    "import csv\n",
    "from math import log\n",
    "import sys\n",
    "\n",
    "from numpy.random import RandomState\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "class NCRPNode(object):\n",
    "\n",
    "    # class variable to keep track of total nodes created so far\n",
    "    total_nodes = 0\n",
    "    last_node_id = 0\n",
    "\n",
    "    def __init__(self, num_levels, vocab, parent=None, level=0,\n",
    "                 random_state=None):\n",
    "\n",
    "        self.node_id = NCRPNode.last_node_id\n",
    "        NCRPNode.last_node_id += 1\n",
    "\n",
    "        self.customers = 0\n",
    "        self.parent = parent\n",
    "        self.children = []\n",
    "        self.level = level\n",
    "        self.total_words = 0\n",
    "        self.num_levels = num_levels\n",
    "\n",
    "        self.vocab = np.array(vocab)\n",
    "        self.word_counts = np.zeros(len(vocab))\n",
    "\n",
    "        if random_state is None:\n",
    "            self.random_state = RandomState()\n",
    "        else:\n",
    "            self.random_state = random_state\n",
    "\n",
    "    def __repr__(self):\n",
    "        parent_id = None\n",
    "        if self.parent is not None:\n",
    "            parent_id = self.parent.node_id\n",
    "        return 'Node=%d level=%d customers=%d total_words=%d parent=%s' % (self.node_id,\n",
    "            self.level, self.customers, self.total_words, parent_id)\n",
    "\n",
    "    def add_child(self):\n",
    "        ''' Adds a child to the next level of this node '''\n",
    "        node = NCRPNode(self.num_levels, self.vocab, parent=self, level=self.level+1)\n",
    "        self.children.append(node)\n",
    "        NCRPNode.total_nodes += 1\n",
    "        return node\n",
    "\n",
    "    def is_leaf(self):\n",
    "        ''' Check if this node is a leaf node '''\n",
    "        return self.level == self.num_levels-1\n",
    "\n",
    "    def get_new_leaf(self):\n",
    "        ''' Keeps adding nodes along the path until a leaf node is generated'''\n",
    "        node = self\n",
    "        for l in range(self.level, self.num_levels-1):\n",
    "            node = node.add_child()\n",
    "        return node\n",
    "\n",
    "    def drop_path(self):\n",
    "        ''' Removes a document from a path starting from this node '''\n",
    "        node = self\n",
    "        node.customers -= 1\n",
    "        if node.customers == 0:\n",
    "            node.parent.remove(node)\n",
    "        for level in range(1, self.num_levels): # skip the root\n",
    "            node = node.parent\n",
    "            node.customers -= 1\n",
    "            if node.customers == 0:\n",
    "                node.parent.remove(node)\n",
    "\n",
    "    def remove(self, node):\n",
    "        ''' Removes a child node '''\n",
    "        self.children.remove(node)\n",
    "        NCRPNode.total_nodes -= 1\n",
    "\n",
    "    def add_path(self):\n",
    "        ''' Adds a document to a path starting from this node '''\n",
    "        node = self\n",
    "        node.customers += 1\n",
    "        for level in range(1, self.num_levels):\n",
    "            node = node.parent\n",
    "            node.customers += 1\n",
    "\n",
    "    def select(self, gamma):\n",
    "        ''' Selects an existing child or create a new one according to the CRP '''\n",
    "\n",
    "        weights = np.zeros(len(self.children)+1)\n",
    "        weights[0] = float(gamma) / (gamma+self.customers)\n",
    "        i = 1\n",
    "        for child in self.children:\n",
    "            weights[i] = float(child.customers) / (gamma + self.customers)\n",
    "            i += 1\n",
    "\n",
    "        choice = self.random_state.multinomial(1, weights).argmax()\n",
    "        if choice == 0:\n",
    "            return self.add_child()\n",
    "        else:\n",
    "            return self.children[choice-1]\n",
    "\n",
    "    def get_top_words(self, n_words, with_weight):\n",
    "        ''' Get the top n words in this node '''\n",
    "\n",
    "        pos = np.argsort(self.word_counts)[::-1]\n",
    "        sorted_vocab = self.vocab[pos]\n",
    "        sorted_vocab = sorted_vocab[:n_words]\n",
    "        sorted_weights = self.word_counts[pos]\n",
    "        sorted_weights = sorted_weights[:n_words]\n",
    "\n",
    "        output = ''\n",
    "        for word, weight in zip(sorted_vocab, sorted_weights):\n",
    "            if with_weight:\n",
    "                output += '%s (%d), ' % (word, weight)\n",
    "            else:\n",
    "                output += '%s, ' % word\n",
    "        return output\n",
    "\n",
    "class HierarchicalLDA(object):\n",
    "\n",
    "    def __init__(self, corpus, vocab,\n",
    "                 alpha=10.0, gamma=1.0, eta=0.1,\n",
    "                 seed=0, verbose=True, num_levels=3):\n",
    "\n",
    "        NCRPNode.total_nodes = 0\n",
    "        NCRPNode.last_node_id = 0\n",
    "\n",
    "        self.corpus = corpus\n",
    "        self.vocab = vocab\n",
    "        self.alpha = alpha  # smoothing on doc-topic distributions\n",
    "        self.gamma = gamma  # \"imaginary\" customers at the next, as yet unused table\n",
    "        self.eta = eta      # smoothing on topic-word distributions\n",
    "\n",
    "        self.seed = seed\n",
    "        self.random_state = RandomState(seed)\n",
    "        self.verbose = verbose\n",
    "\n",
    "        self.num_levels = num_levels\n",
    "        self.num_documents = len(corpus)\n",
    "        self.num_types = len(vocab)\n",
    "        self.eta_sum = eta * self.num_types\n",
    "\n",
    "        # if self.verbose:\n",
    "        #     for d in range(len(self.corpus)):\n",
    "        #         doc = self.corpus[d]\n",
    "        #         words = ' '.join([self.vocab[n] for n in doc])\n",
    "        #         print 'doc_%d = %s' % (d, words)\n",
    "\n",
    "        # initialise a single path\n",
    "        path = np.zeros(self.num_levels, dtype=np.object)\n",
    "\n",
    "        # initialize and fill the topic pointer arrays for\n",
    "        # every document. Set everything to the single path that\n",
    "        # we added earlier.\n",
    "        self.root_node = NCRPNode(self.num_levels, self.vocab)\n",
    "        self.document_leaves = {}                                   # currently selected path (ie leaf node) through the NCRP tree\n",
    "        self.levels = np.zeros(self.num_documents, dtype=np.object) # indexed < doc, token >\n",
    "        for d in range(len(self.corpus)):\n",
    "\n",
    "            # populate nodes into the path of this document\n",
    "            doc = self.corpus[d]\n",
    "            doc_len = len(doc)\n",
    "            path[0] = self.root_node\n",
    "            self.root_node.customers += 1 # always add to the root node first\n",
    "            for level in range(1, self.num_levels):\n",
    "                # at each level, a node is selected by its parent node based on the CRP prior\n",
    "                parent_node = path[level-1]\n",
    "                level_node = parent_node.select(self.gamma)\n",
    "                level_node.customers += 1\n",
    "                path[level] = level_node\n",
    "\n",
    "            # set the leaf node for this document\n",
    "            leaf_node = path[self.num_levels-1]\n",
    "            self.document_leaves[d] = leaf_node\n",
    "\n",
    "            # randomly assign each word in the document to a level (node) along the path\n",
    "            self.levels[d] = np.zeros(doc_len, dtype=np.int)\n",
    "            for n in range(doc_len):\n",
    "                w = doc[n]\n",
    "                random_level = self.random_state.randint(self.num_levels)\n",
    "                random_node = path[random_level]\n",
    "                random_node.word_counts[w] += 1\n",
    "                random_node.total_words += 1\n",
    "                self.levels[d][n] = random_level\n",
    "\n",
    "    def estimate(self, num_samples, display_topics=50, n_words=5, with_weights=True):\n",
    "\n",
    "        print('HierarchicalLDA sampling')\n",
    "        for s in range(num_samples):\n",
    "\n",
    "            sys.stdout.write('.')\n",
    "\n",
    "            for d in range(len(self.corpus)):\n",
    "                self.sample_path(d)\n",
    "\n",
    "            for d in range(len(self.corpus)):\n",
    "                self.sample_topics(d)\n",
    "\n",
    "            if (s > 0) and ((s+1) % display_topics == 0):\n",
    "                print(\" %d\" % (s+1))\n",
    "                self.print_nodes(n_words, with_weights)\n",
    "\n",
    "    def sample_path(self, d):\n",
    "\n",
    "        # define a path starting from the leaf node of this doc\n",
    "        path = np.zeros(self.num_levels, dtype=np.object)\n",
    "        node = self.document_leaves[d]\n",
    "        for level in range(self.num_levels-1, -1, -1): # e.g. [3, 2, 1, 0] for num_levels = 4\n",
    "            path[level] = node\n",
    "            node = node.parent\n",
    "\n",
    "        # remove this document from the path, deleting empty nodes if necessary\n",
    "        self.document_leaves[d].drop_path()\n",
    "\n",
    "        ############################################################\n",
    "        # calculates the prior p(c_d | c_{-d}) in eq. (4)\n",
    "        ############################################################\n",
    "\n",
    "        node_weights = {}\n",
    "        self.calculate_ncrp_prior(node_weights, self.root_node, 0.0)\n",
    "\n",
    "        ############################################################\n",
    "        # calculates the likelihood p(w_d | c, w_{-d}, z) in eq. (4)\n",
    "        ############################################################\n",
    "\n",
    "        level_word_counts = {}\n",
    "        for level in range(self.num_levels):\n",
    "            level_word_counts[level] = {}\n",
    "        doc_levels = self.levels[d]\n",
    "        doc = self.corpus[d]\n",
    "\n",
    "        # remove doc from path\n",
    "        for n in range(len(doc)): # for each word in the doc\n",
    "\n",
    "            # count the word at each level\n",
    "            level = doc_levels[n]\n",
    "            w = doc[n]\n",
    "            if w not in level_word_counts[level]:\n",
    "                level_word_counts[level][w] = 1\n",
    "            else:\n",
    "                level_word_counts[level][w] += 1\n",
    "\n",
    "            # remove word count from the node at that level\n",
    "            level_node = path[level]\n",
    "            level_node.word_counts[w] -= 1\n",
    "            level_node.total_words -= 1\n",
    "            assert level_node.word_counts[w] >= 0\n",
    "            assert level_node.total_words >= 0\n",
    "\n",
    "        self.calculate_doc_likelihood(node_weights, level_word_counts)\n",
    "\n",
    "        ############################################################\n",
    "        # pick a new path\n",
    "        ############################################################\n",
    "\n",
    "        nodes = np.array(list(node_weights.keys()))\n",
    "        weights = np.array([node_weights[node] for node in nodes])\n",
    "        weights = np.exp(weights - np.max(weights)) # normalise so the largest weight is 1\n",
    "        weights = weights / np.sum(weights)\n",
    "\n",
    "        choice = self.random_state.multinomial(1, weights).argmax()\n",
    "        node = nodes[choice]\n",
    "\n",
    "        # if we picked an internal node, we need to add a new path to the leaf\n",
    "        if not node.is_leaf():\n",
    "            node = node.get_new_leaf()\n",
    "\n",
    "        # add the doc back to the path\n",
    "        node.add_path()                     # add a customer to the path\n",
    "        self.document_leaves[d] = node      # store the leaf node for this doc\n",
    "\n",
    "        # add the words\n",
    "        for level in range(self.num_levels-1, -1, -1): # e.g. [3, 2, 1, 0] for num_levels = 4\n",
    "            word_counts = level_word_counts[level]\n",
    "            for w in word_counts:\n",
    "                node.word_counts[w] += word_counts[w]\n",
    "                node.total_words += word_counts[w]\n",
    "            node = node.parent\n",
    "\n",
    "    def calculate_ncrp_prior(self, node_weights, node, weight):\n",
    "        ''' Calculates the prior on the path according to the nested CRP '''\n",
    "\n",
    "        for child in node.children:\n",
    "            child_weight = log( float(child.customers) / (node.customers + self.gamma) )\n",
    "            self.calculate_ncrp_prior(node_weights, child, weight + child_weight)\n",
    "\n",
    "        node_weights[node] = weight + log( self.gamma / (node.customers + self.gamma))\n",
    "\n",
    "    def calculate_doc_likelihood(self, node_weights, level_word_counts):\n",
    "\n",
    "        # calculate the weight for a new path at a given level\n",
    "        new_topic_weights = np.zeros(self.num_levels)\n",
    "        for level in range(1, self.num_levels):  # skip the root\n",
    "\n",
    "            word_counts = level_word_counts[level]\n",
    "            total_tokens = 0\n",
    "\n",
    "            for w in word_counts:\n",
    "                count = word_counts[w]\n",
    "                for i in range(count):  # why ?????????\n",
    "                    new_topic_weights[level] += log((self.eta + i) / (self.eta_sum + total_tokens))\n",
    "                    total_tokens += 1\n",
    "\n",
    "        self.calculate_word_likelihood(node_weights, self.root_node, 0.0, level_word_counts, new_topic_weights, 0)\n",
    "\n",
    "    def calculate_word_likelihood(self, node_weights, node, weight, level_word_counts, new_topic_weights, level):\n",
    "\n",
    "        # first calculate the likelihood of the words at this level, given this topic\n",
    "        node_weight = 0.0\n",
    "        word_counts = level_word_counts[level]\n",
    "        total_words = 0\n",
    "\n",
    "        for w in word_counts:\n",
    "            count = word_counts[w]\n",
    "            for i in range(count): # why ?????????\n",
    "                node_weight += log( (self.eta + node.word_counts[w] + i) /\n",
    "                                    (self.eta_sum + node.total_words + total_words) )\n",
    "                total_words += 1\n",
    "\n",
    "        # propagate that weight to the child nodes\n",
    "        for child in node.children:\n",
    "            self.calculate_word_likelihood(node_weights, child, weight + node_weight,\n",
    "                                           level_word_counts, new_topic_weights, level+1)\n",
    "\n",
    "        # finally if this is an internal node, add the weight of a new path\n",
    "        level += 1\n",
    "        while level < self.num_levels:\n",
    "            node_weight += new_topic_weights[level]\n",
    "            level += 1\n",
    "\n",
    "        node_weights[node] += node_weight\n",
    "\n",
    "    def sample_topics(self, d):\n",
    "\n",
    "        doc = self.corpus[d]\n",
    "\n",
    "        # initialise level counts\n",
    "        doc_levels = self.levels[d]\n",
    "        level_counts = np.zeros(self.num_levels, dtype=np.int)\n",
    "        for c in doc_levels:\n",
    "            level_counts[c] += 1\n",
    "\n",
    "        # get the leaf node and populate the path\n",
    "        path = np.zeros(self.num_levels, dtype=np.object)\n",
    "        node = self.document_leaves[d]\n",
    "        for level in range(self.num_levels-1, -1, -1): # e.g. [3, 2, 1, 0] for num_levels = 4\n",
    "            path[level] = node\n",
    "            node = node.parent\n",
    "\n",
    "        # sample a new level for each word\n",
    "        level_weights = np.zeros(self.num_levels)\n",
    "        for n in range(len(doc)):\n",
    "\n",
    "            w = doc[n]\n",
    "            word_level = doc_levels[n]\n",
    "\n",
    "            # remove from model\n",
    "            level_counts[word_level] -= 1\n",
    "            node = path[word_level]\n",
    "            node.word_counts[w] -= 1\n",
    "            node.total_words -= 1\n",
    "\n",
    "            # pick new level\n",
    "            for level in range(self.num_levels):\n",
    "                level_weights[level] = (self.alpha + level_counts[level]) *                     \\\n",
    "                    (self.eta + path[level].word_counts[w]) /                                   \\\n",
    "                    (self.eta_sum + path[level].total_words)\n",
    "            level_weights = level_weights / np.sum(level_weights)\n",
    "            level = self.random_state.multinomial(1, level_weights).argmax()\n",
    "\n",
    "            # put the word back into the model\n",
    "            doc_levels[n] = level\n",
    "            level_counts[level] += 1\n",
    "            node = path[level]\n",
    "            node.word_counts[w] += 1\n",
    "            node.total_words += 1\n",
    "\n",
    "    def print_nodes(self, n_words, with_weights):\n",
    "        self.print_node(self.root_node, 0, n_words, with_weights)\n",
    "\n",
    "    def print_node(self, node, indent, n_words, with_weights):\n",
    "        out = '    ' * indent\n",
    "        out += 'topic=%d level=%d (documents=%d): ' % (node.node_id, node.level, node.customers)\n",
    "        out += node.get_top_words(n_words, with_weights)\n",
    "        print(out)\n",
    "        for child in node.children:\n",
    "            self.print_node(child, indent+1, n_words, with_weights)\n",
    "\n",
    "def load_vocab(file_name):\n",
    "    with open(file_name, 'rb') as f:\n",
    "        vocab = []\n",
    "        reader = csv.reader(f)\n",
    "        for row in reader:\n",
    "            idx, word = row\n",
    "            stripped = word.strip()\n",
    "            vocab.append(stripped)\n",
    "        return vocab\n",
    "\n",
    "def load_corpus(file_name):\n",
    "    with open(file_name, 'rb') as f:\n",
    "        corpus = []\n",
    "        reader = csv.reader(f)\n",
    "        for row in reader:\n",
    "            doc = []\n",
    "            for idx_and_word in row:\n",
    "                stripped = idx_and_word.strip()\n",
    "                tokens = stripped.split(' ')\n",
    "                if len(tokens) == 2:\n",
    "                    idx, word = tokens\n",
    "                    doc.append(int(idx))\n",
    "            corpus.append(doc)\n",
    "        return corpus"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.2 64-bit ('risotto': venv)",
   "language": "python",
   "name": "python38264bitrisottovenv582be2314c5f42eab2123d4127b9e107"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}