Get most likely topic per document in pandas dataframe using gensim

Question

I am using gensim LDA to build a topic model for a bunch of documents that I have stored in a pandas data frame. Once the model is built, I can call model.get_document_topics(model_corpus) to get a list of list of tuples showing the topic distribution for each document. For example, when I am working with 20 topics, I might get the following for the first three documents in my data frame:

[(5, 0.11253482), (7, 0.75876033)]
[(19, 0.96343607)]
[(0, 0.010002977),
 (1, 0.010002977),
 (2, 0.010002977),
 (3, 0.010002979),
 (4, 0.8099435),
 (5, 0.010002977),
 (6, 0.010002977),
 (7, 0.010002977),
 (8, 0.010002977),
 (9, 0.010002977),
 (10, 0.010002977),
 (11, 0.010002977),
 (12, 0.010002977),
 (13, 0.010002977),
 (14, 0.010002977),
 (15, 0.010002977),
 (16, 0.010002977),
 (17, 0.010002977),
 (18, 0.010002977),
 (19, 0.010002977)]

This means that the most likely topic for document_1 is 7, for document_2 is 19, and for document_3 is 4. The primary output that I would like to see is simply this most likely topic for each document. The way I'm doing this now is using a loop:

import numpy as np
import pandas as pd

def get_max(doc):
        idx,l = zip(*doc)
        return idx[np.argmax(l)]

data['doc_topic'] = [get_max(doc) for doc in model.get_document_topics(model_corpus)]

I have around 80k documents in my data frame, so this code takes about 45 seconds to execute. But since gensim has already done all the computations, I keep thinking that that 45 seconds of computational time is simply spent on reorganizing data, so there must be a more efficient way of doing this.

If possible, a secondary output that would be nice to have is the document-topic matrix, such that each row corresponds to a document in my data frame, and each column represents the probability (or similarity) of the document to the topic. So this would yield a DxT matrix, where D is the number of documents, and T is the number of topics.

Answer 1

Timing your get_max function on the largest list you provided gives a time of a little over 9 us which seems quite efficient. I have, however, written another function that does the same but takes a little more than 3 us, so almost 3x as fast. For smaller lists this speedup will become even greater (~10x speedup for lists of size 1 or 2).

doc = [(0, 0.010002977),
 (1, 0.010002977),
 (2, 0.010002977),
 (3, 0.010002979),
 (4, 0.8099435),
 (5, 0.010002977),
 (6, 0.010002977),
 (7, 0.010002977),
 (8, 0.010002977),
 (9, 0.010002977),
 (10, 0.010002977),
 (11, 0.010002977),
 (12, 0.010002977),
 (13, 0.010002977),
 (14, 0.010002977),
 (15, 0.010002977),
 (16, 0.010002977),
 (17, 0.010002977),
 (18, 0.010002977),
 (19, 0.010002977)]

def get_max(doc):
    idx,l = zip(*doc)
    return idx[np.argmax(l)]

def get_max_new(doc):
    return sorted(doc, key=lambda x: x[1], reverse=True)[0][0]

%timeit get_max(doc) # 9.14 us

%timeit get_max_new(doc) # 3.53 us

There may be even more speedups that can be achieved, however that is more difficult for me to test at the moment.