I am computing co-occurence matrix for a fixed windows size in python using scipy's lil_matrix for storing the counts and computing the counts by sliding the context window over each word and then counting in the window.

Now the code is taking too much time for relatively small corpus size also (100 MB Wikipedia dump). The code is :

def gen_coocur(window_size=5):
        Generates coocurrence matrix
        # vocab is precomputed. 
        coocur_matrix = lil_matrix((len(vocab)+1, len(vocab)+1), dtype=np.float64)
        for page in self.wiki_extract.get_page():
            # word_tokenize is tokenizer from nltk 
            doc_tokens = word_tokenize(page.decode('utf-8'))
            N = len(doc_tokens)
            for token in self.vocab:
                for i in xrange(0,window_size):
                    if (token in doc_tokens[0:i] or token in doc_tokens[i:(i+window_size+1)]) and token != doc_tokens[i]:
                        coocur_matrix[self.vocab[doc_tokens[i]],self.vocab[token]] +=1
                for i in xrange(window_size, (N-window_size)):
                    if token in doc_tokens[(i-window_size):(i+window_size+1)] and token != doc_tokens[i]:
                        coocur_matrix[self.vocab[doc_tokens[i]],self.vocab[token]] +=1
                for i in xrange(N-window_size, N):
                    if (token in doc_tokens[i:N] or token in doc_tokens[i-window_size:N]) and token != doc_tokens[i]:
                        coocur_matrix[self.vocab[doc_tokens[i]],self.vocab[token]] +=1            

vocab is a dictionary which maps words -> wordId. How can I optimize this code to run faster?


1 Answer 1


From easiest to hardest:

  1. Try running it in pypy or numba
  2. Find a faster implementation. Unfortunately I can not recommend one.
  3. Parallelize the loop over the documents. Not so hard since your vocabulary is precomputed. (Even if it weren't you could get away with it using the hashing trick.) Combine this with the first bullet.
  4. Rewrite the inner loop in Cython.
  5. Rewrite the whole thing in a faster language like C++ or Scala.

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.