trjhtr

I have written a method which is designed to calculate the word co-occurrence matrix in a corpus, such that element(i,j) is the number of times that word i follows word j in the corpus.

Here is my code with a small example:

import numpy as np 
import nltk
from nltk import bigrams 

def co_occurrence_matrix(corpus):
 vocab = set(corpus)
 vocab = list(vocab)

 # Key:Value = Word:Index
 vocab_to_index = word:i for i, word in enumerate(vocab) 

 # Create bigrams from all words in corpus
 bi_grams = list(bigrams(corpus))

 # Frequency distribution of bigrams ((word1, word2), num_occurrences)
 bigram_freq = nltk.FreqDist(bi_grams).most_common(len(bi_grams))

 # Initialise co-occurrence matrix
 # co_occurrence_matrix[current][previous]
 co_occurrence_matrix = np.zeros((len(vocab), len(vocab)))

 # Loop through the bigrams in the frequency distribution, noting the 
 # current and previous word, and the number of occurrences of the bigram.
 # Get the vocab index of the current and previous words.
 # Put the number of occurrences into the appropriate element of the array.
 for bigram in bigram_freq:
 current = bigram[0][1]
 previous = bigram[0][0]
 count = bigram[1]
 pos_current = vocab_to_index[current]
 pos_previous = vocab_to_index[previous]
 co_occurrence_matrix[pos_current][pos_previous] = count 

 co_occurrence_matrix = np.matrix(co_occurrence_matrix)

 return co_occurrence_matrix

test_sent = ['hello', 'i', 'am', 'hello', 'i', 'dont', 'want', 'to', 'i', 'dont']
m = co_occurrence_matrix(test_sent)

Output:

[[0. 2. 0. 0. 0. 0.]
 [0. 0. 0. 1. 0. 2.]
 [0. 1. 0. 0. 0. 0.]
 [0. 0. 0. 0. 1. 0.]
 [1. 0. 0. 0. 0. 0.]
 [0. 0. 1. 0. 0. 0.]]

Whilst the example shown works fine, when I scale this up to a much larger corpus, I get the following Killed:9 error. I assume this is because the matrix is very large.

I am looking to make this method more efficient so that I can use it for large corpuses! (A few million words.)

A 10⁶ × 10⁶ matrix would contain 10¹² entries. Optimistically assuming one byte per entry, that would already be 1 TB. I would expect that most of the matrix entries will be 0. Consider looking into sparse matrices.