Generating a word bigram co-occurrence matrix

Question

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.)

Answer 1

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.