连接大型CSV的所有行

Question

So I have a large csv file (5 million rows) with multiple columns. Of particular interest to me is a column which contains text.

The input csv is of the following format:

system_id, member_name, message, is_post

0157e407,member1011, "I have had problems with my lungs for years now. It all started with an infection...", false

1915d457, member1055, "Looks like a lot of people take Paracetamol for managing pain and....",false

The column 'message' contains text and is of interest.

Now the task is to concatenate all the rows of this column into a one single large text, and then compute n-grams (n=1,2,3,4,5) on it. The output should be 5 different files corresponding to n-grams in the following format: For eg:

bigram.csv

n-gram, count

"word1 word2", 7

"word1 word3", 11

trigram.csv

n-gram, count

"word1 word2 word3", 22

"word 1 word2 word4", 24

Here is what I have tried so far:

from collections import OrderedDict
import csv
import re
import sys

import nltk


if __name__ == '__main__':
    if len(sys.argv) < 2:
        print "%d Arguments Given : Exiting..." % (len(sys.argv)-1)
        print "Usage: python %s <inp_file_path>" % sys.argv[0]
        exit(1)
    ifpath = sys.argv[1]
    with open(ifpath, 'r') as ifp:
        reader = csv.DictReader(ifp)
        all_msgs = []
        fieldnames = reader.fieldnames
        processed_rows = []
        for row in reader:
            msg = row['message']
            res = {'message': msg}
            txt = msg.decode('ascii', 'ignore')
            # some preprocessing
            txt = re.sub(r'[\.]{2,}', r". ", txt)
            txt = re.sub(r'([\.,;!?])([A-Z])', r'\1 \2', txt)
            sentences = nltk.tokenize.sent_tokenize(txt.strip())
            all_msgs.append(' '.join(sentences))
    text = ' '.join(all_msgs)

    tokens = nltk.word_tokenize(text)
    tokens = [token.lower() for token in tokens if len(token) > 1]
    bi_tokens = list(nltk.bigrams(tokens))
    tri_tokens = list(nltk.trigrams(tokens))
    bigrms = []
    for item in sorted(set(bi_tokens)):
        bb = OrderedDict()
        bb['bigrams'] = ' '.join(item)
        bb['count'] = bi_tokens.count(item)
        bigrms.append(bb)

    trigrms = []
    for item in sorted(set(tri_tokens)):
        tt = OrderedDict()
        tt['trigrams'] = ' '.join(item)
        tt['count'] = tri_tokens.count(item)
        trigrms.append(tt)

    with open('bigrams.csv', 'w') as ofp2:
        header = ['bigrams', 'count']
        dict_writer = csv.DictWriter(ofp2, header)
        dict_writer.writeheader()
        dict_writer.writerows(bigrms)

    with open('trigrams.csv', 'w') as ofp3:
        header = ['trigrams', 'count']
        dict_writer = csv.DictWriter(ofp3, header)
        dict_writer.writeheader()
        dict_writer.writerows(trigrms)

    tokens = nltk.word_tokenize(text)
    fourgrams = nltk.collocations.QuadgramCollocationFinder.from_words(tokens)
    quadgrams = []
    for fourgram, freq in fourgrams.ngram_fd.items():
        dd = OrderedDict()
        dd['quadgram'] = " ".join(fourgram)
        dd['count'] = freq
        quadgrams.append(dd)
    with open('quadgram.csv', 'w') as ofp4:
        header = ['quadgram', 'count']
        dict_writer = csv.DictWriter(ofp4, header)
        dict_writer.writeheader()
        dict_writer.writerows(quadgrams)

This has been running for past 2 days on a 4 core machine. How can I make this more efficient (using pandas and/or multiprocessing, perhaps) and speed it up as reasonably as possible?

Answer 1

I would make a few changes:

Find the bottleneck

What portion is taking so long?

• Reading the CSV
• Tokenizing
• Making the n-grams
• Counting the n-grams
• Writing to disk

So the first thing I would do is make a cleaner separation between the different steps, and ideally make it possible to restart half-way

Reading the text

I would extract this to a different method. And from what I read ( here for example) pandas reads csv-files a lot quicker than csv . If the reading of the csv only takes 1 minute of the 2 days, this might not be an issue, but I would do something like this:

def read_text(filename):  # you could add **kwarg to pass on to the read_csv
    df = pd.read_csv(filename) # add info on file encoding etc
    message = df['message'].str.replace(r'[\.]{2,}', r". ")  # http://pandas.pydata.org/pandas-docs/stable/generated/pandas.Series.str.replace.html
    message = message.str.replace(r'([\.,;!?])([A-Z])', r'\1 \2')

    messame = message.strip()
    sentences = message.apply(nltk.tokenize.sent_tokenize)
    return ' '.join(sentences.appy(' '.join))

You can even do this in chunks, and yield instead of return the sentences to make it a generator, potentially saving on memory

Is there a specific reason you join the sentences after the sent_tokenize because I find this in the documentation

The Treebank tokenizer uses regular expressions to tokenize text as in Penn Treebank. This is the method that is invoked by word_tokenize(). It assumes that the text has already been segmented into sentences, eg using sent_tokenize().

So you would call it like this:

text = read_text(csv_file)
with open(text_file, 'w') as file:
    file.write(text)
print('finished reading text from file') # or use logging

Tokenizing

stays rougly the same

tokens = nltk.word_tokenize(text)
print('finished tokenizing the text')

def save_tokens(filename, tokens):
    # save the list somewhere, either json or pickle, so you can pick up later if something goes wrong

making the n-grams, counting and writing them to disk

Your code contains a lot of boilerplate which does the same thing with just a different function or filename, so instead I abstract this away in a list of tuples containing the name, the function to get the bigrams to count them and the filename to save

ngrams = [
    ('bigrams', nltk.bigrams, collections.Counter, 'bigrams.csv'),
    ('trigrams', nltk.trigrams, collections.Counter, 'quadgrams.csv'),
    ('quadgrams', nltk.collocations.QuadgramCollocationFinder.from_words, parse_quadgrams, 'quadgrams.csv'),
]

If you want to count how much items are in a list, just use collections.Counter instead of making an (expensive) collection.OrderedDict on each item. If you want to do the counting yourself, better to use tuples than OrderedDict . You could also use pd.Series.value_counts()

def parse_quadgrams(quadgrams):
    return quadgrams.ngram_fd #from what I see in the code this dict already contains the counts

for name, ngram_method, parse_method, output_file in ngrams:
    grams = ngram_method(tokens)
    print('finished generating ', name)
    # You could write this intermediate result to a temporary file in case something goes wrong
    count_df = pd.Series(parse_method(grams)).reset_index().rename(columns={'index': name, 0: 'count')
    # if you need it sorted you can do this on the DataFrame
    print('finished counting ', name)
    count_df.to_csv(output_file)
    print('finished writing ', name, ' to file: ', output_file)