Calculating text similarity, filtering the results, and reshaping matrix into a list-of-tuples generates MemoryError

Question

I have a constantly growing corpus of currently ~36,000 documents (growing daily) and I want calculate the similarity between each pair of documents. After calculating the similarity scores, I want to filter the results to only include scores above 0.9, and capture the (row label, column label, metric) in a list of tuples that can be written to a CSV file.

The process that I currently have works as intended, but as my corpus of documents has continued to grow, I'm now receiving a MemoryError when trying to process my actual dataset.

Error generated with my actual dataset:

2020-08-10 10:37:08,933 There are 35845 records loaded
2020-08-10 10:37:19,570 Completed pre-processing all documents
2020-08-10 10:38:05,458 Completed calculating similarity
Traceback (most recent call last):
  File "/home/curtis/project/text_similarity.py", line 97, in <module>
    scores = get_scores(pairwise_similarity, doc_keys, threshold=0.9)
  File "/home/curtis/project/text_similarity.py", line 61, in get_scores
    arr[np.tril_indices(arr.shape[0], -1)] = np.nan
  File "/home/curtis/miniconda3/envs/scraper-dev/lib/python3.7/site-packages/numpy/lib/twodim_base.py", line 868, in tril_indices
    return nonzero(tri(n, m, k=k, dtype=bool))
  File "<__array_function__ internals>", line 6, in nonzero
  File "/home/curtis/miniconda3/envs/project-dev/lib/python3.7/site-packages/numpy/core/fromnumeric.py", line 1896, in nonzero
    return _wrapfunc(a, 'nonzero')
  File "/home/curtis/miniconda3/envs/project-dev/lib/python3.7/site-packages/numpy/core/fromnumeric.py", line 61, in _wrapfunc
    return bound(*args, **kwds)
MemoryError: Unable to allocate 9.57 GiB for an array with shape (642414090, 2) and data type int64

Existing process text_similarity.py with sample data:

import csv
import logging
import os
import string

import numpy as np
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.datasets import fetch_20newsgroups


class clean_document:
    """Class to pre-process documents"""

    def __init__(self, input_string, stopwords):
        self.input_string = input_string
        self.string_lower = self.lower_string()
        self.string_no_emoji = self.drop_emoji()
        self.string_no_punct = self.remove_punction()
        self.tokens = self.tokenizer()
        self.tokens_no_stopwords = self.remove_stopwords(stopwords)

    def lower_string(self):
        string_lower = self.input_string.lower()
        return string_lower

    def drop_emoji(self):
        """Thanks to https://stackoverflow.com/a/49986645"""
        no_emoji = self.string_lower.encode("ascii", "ignore").decode("ascii")
        return no_emoji

    def remove_punction(self):
        """Thanks to https://stackoverflow.com/a/266162"""
        no_punct = self.string_no_emoji.translate(str.maketrans("", "", string.punctuation))
        return no_punct

    def tokenizer(self):
        tokens = word_tokenize(self.string_no_punct)
        return tokens

    def remove_stopwords(self, stopwords):
        no_stopwords = [line for line in self.tokens if line not in stopwords]
        return no_stopwords


def calc_pairwise_similarity(corpus):
    """Thanks to https://stackoverflow.com/a/8897648"""
    vect = TfidfVectorizer(min_df=1)
    tfidf = vect.fit_transform(corpus)
    pairwise_similarity = tfidf * tfidf.T

    return pairwise_similarity


def get_scores(pairwise_similarity, doc_keys, threshold=0.9):
    """Extract scores into a list-of-tuples"""
    arr = pairwise_similarity.toarray()
    arr[arr <= threshold] = np.nan
    np.fill_diagonal(arr, np.nan)
    arr[np.tril_indices(arr.shape[0], -1)] = np.nan
    idx = (~np.isnan(arr)).nonzero()
    vals = arr[idx].tolist()
    keys = list(zip(idx[0].tolist(), idx[1].tolist()))
    output = [(x[0][0], x[0][1], x[1]) for x in list(zip(keys, vals))]
    final = [(doc_keys[line[0]], doc_keys[line[1]], line[2]) for line in output]

    return final


# MAIN PROGRAM

# set up basic logging
logging.basicConfig(format="%(asctime)s %(message)s", level=logging.INFO)

# load the dataset
newsgroups_train = fetch_20newsgroups(subset='train')
documents = {i: line for i, line in enumerate(newsgroups_train['data'])}
logging.info(f"There are {len(documents)} records loaded")

# define the stop words to use
stop_words = set(stopwords.words("english"))

# process the original strings and create a cleaned corpus
corpus = []
for line in documents.values():
    x = clean_document(line, stop_words)
    corpus.append(x.string_no_punct)
logging.info("Completed pre-processing all documents")

# calculate pairwise similaritry
pairwise_similarity = calc_pairwise_similarity(corpus)
logging.info("Completed calculating similarity")

# get similiarity metrics
doc_keys = list(documents.keys())
scores = get_scores(pairwise_similarity, doc_keys, threshold=0.9)
logging.info("Extracted similarity scores")

# write scores to CSV file
with open("scores.csv", "w") as f:
    writer = csv.writer(f)
    writer.writerows(scores)
logging.info("Successfully wrote metrics to file")

Answer 1

A quick fix to the MemoryError issue is to avoid expanding the sparse similarity matrix into a dense matrix for extracting indices in the get_scores function. Instead, the indices can be extracted by operating directly on the sparse matrix.

def get_scores(pairwise_similarity, doc_keys, threshold=0.9):
    sim_coo = pairwise_similarity.tocoo(copy=False)
    doc_keys = np.array(doc_keys)
    mask = (sim_coo.data > threshold) & (sim_coo.row > sim_coo.col)
    row_doc_keys = doc_keys[sim_coo.row[mask]]
    col_doc_keys = doc_keys[sim_coo.col[mask]]
    sim_values = sim_coo.data[mask]
    return zip(row_doc_keys, col_doc_keys, sim_values)

Given the size of your data and memory limit, this fix would no longer work if your data size increases by 40%+ since the similarity matrix itself would already run out of memory. To further reduce memory usage, the similarity values need to be updated incrementally rather than computed all at once.

I'll leave other code improvements to other reviews.