# Get most likely topic per document in pandas dataframe using gensim

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.

## 1 Answer

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.

• This is a good suggestion. Unfortunately, when I tested on the entire set of documents this was about 2 seconds slower, which is a little confusing given your timing on a single document. I'm still pretty surprised at how long this kind of query takes given the library gensim has been around for a while. Jan 28 '21 at 15:36