increase speed & lower memory consumption of text normalization function

Question

I'm trying to preprocess a large text document. I've written a text normalization function which takes a disproprtionate amount of time and memory. How can I format the function to lower these two?

The time result below was for this example.

t = "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. Curabitur pretium tincidunt lacus. Nulla gravida orci a odio. Nullam varius, turpis et commodo pharetra, est eros bibendum elit, nec luctus magna felis sollicitudin mauris. Integer in mauris eu nibh euismod gravida. Duis ac tellus et risus vulputate vehicula. Donec lobortis risus a elit. Etiam tempor. Ut ullamcorper, ligula eu tempor congue, eros est euismod turpis, id tincidunt sapien risus a quam. Maecenas fermentum consequat mi. Donec fermentum. Pellentesque malesuada nulla a mi. Duis sapien sem, aliquet nec, commodo eget, consequat quis, neque. Aliquam faucibus, elit ut dictum aliquet, felis nisl adipiscing sapien, sed malesuada diam lacus eget erat. Cras mollis scelerisque nunc. Nullam arcu. Aliquam consequat. Curabitur augue lorem, dapibus quis, laoreet et, pretium ac, nisi. Aenean magna nisl, mollis quis, molestie eu, feugiat in, orci. In hac habitasse platea dictumst."

This is the script.

import time
import string
import nltk
import numpy as np

# text preprocessing module, use boolean flags to customize function
def normalize_text(text, lowercase=True, punctuation=True):

    # Step 1: Tokenize
    output = nltk.word_tokenize(text)

    # Step 2: Convert to lowercase (optional)
    if lowercase:
        output = [word.lower() for word in output]

    # Step 3: Remove punctuation:
    if punctuation:
        output = [str(token).translate(str.maketrans('', '', string.punctuation)) for token in output]
        output = [token for token in output if token != '']

    return(output)

Removing all filters and the corresponding if-statements sped up processing by a mere 0.7%.

def normalize_text2(text):

    # Step 1: Tokenize
    output = nltk.word_tokenize(text)

    # Step 2: Convert to lowercase (optional)
    output = [word.lower() for word in output]

    # Step 3: Remove punctuation:
    output = [str(token).translate(str.maketrans('', '', string.punctuation)) for token in output]
    output = [token for token in output if token != '']

    return(output)

Here is the bagged comparison.

times1 = []

for i in range(1000):
    start = time.time()
    tokens = normalize_text(t)
    end = time.time()
    times1.append(end - start)

time1 = np.mean(times1)
print(time1)    

times2 = []

for i in range(1000):
    start = time.time()
    tokens = normalize_text2(t)
    end = time.time()
    times2.append(end - start)

time2 = np.mean(times2)
print(time2) 

print(time2/time1)

Here are the results:

0.0021646411418914796
0.0021491129398345946
0.9928264312470212

Any advice on how to improve further? For example, how could I reduce the number of different list comprehensions, so that the same sequence of text does not need to crunched anew this many times?

Answer 1

You can save a little bit of time by not re-running str.maketrans for each token, since it's always going to produce the same result:

import nltk
from statistics import mean
import string
import time
from typing import List


def normalize_text3(text: str) -> List[str]:
    output: List[str] = []
    punctuation_filter = str.maketrans('', '', string.punctuation)
    for token in nltk.word_tokenize(text):
        token = token.translate(punctuation_filter)
        if not token:
            continue
        output.append(token.lower())
    return output

tested with:

for func in [normalize_text, normalize_text2, normalize_text3]:
    times = []
    for _ in range(1000):
        start = time.time()
        tokens = normalize_text(t)
        end = time.time()
        times.append(end - start)
    print(f"{func.__name__.rjust(15)}: {mean(times)}")

gets me:

dog runs
 normalize_text: 0.003226396322250366
normalize_text2: 0.0032752704620361327
normalize_text3: 0.0030987038612365725

If you want to lower memory consumption, you might consider having this function return a generator rather than a list...

Answer 2

A slight change to the answer:

def normalize_text4(text):
    output: List[str] = []
    punctuation_filter = str.maketrans('', '', string.punctuation)
    for token in nltk.word_tokenize(text, preserve_line=True):
        token = token.translate(punctuation_filter)
        if token:
            output.append(token.lower())
    return output

Usage with preserve_line=True is a bit faster, and with the punctuation removed, the result is the same as with the default preserve_line=False. As most of the time is spent in word_tokenize, this is the first you want to optimize, though I haven't looked farther than this.

Here is the measurement (same as above):

times = []
base = None
for fct in (normalize_text, normalize_text2, normalize_text3, normalize_text4):
    for i in range(1000):
        start = time.time()
        tokens = fct(t)
        end = time.time()
        times.append(end - start)

    avg = np.mean(times)
    if not base:
        base = avg
    print(f'{fct.__name__:15}: {avg * 1000:4.3} ms, {avg / base * 100:6.4} %')

and the results (on my Windows 10 notebook):

normalize_text : 4.88 ms,  100.0 %
normalize_text2: 4.86 ms,  99.44 %
normalize_text3: 4.64 ms,  94.93 %
normalize_text4: 3.85 ms,  78.88 %

The results vary, with the percentage somewhere between 74 and 82%, but this is a typical outcome.

EDIT:
Something I noticed afterwards, and that I don't have an explanation for: if you run normalize_text4 before any of the other scripts (that use preserve_line=False) instead of after them, it is quite a bit faster:

normalize_text4: 1.81 ms,  41.07 %
normalize_text : 4.42 ms,  100.0 %
normalize_text4: 3.57 ms,  80.76 %

(I changed the script to have normalize_text as base like before)
I would guess that some caching is happening that is counter-productive in this (constructed) case.