1
\$\begingroup\$

I writing a function in Python to substitute word tokens in a sentence with IDs.

The sentence is a list of tokens (list_of_tokens). The IDs are provided in a dictionary mapping tokens with an IDs (tokens2IDs_dict).

However, the functions I've written is not very efficient because of the for-loop. I am wondering if there is any more pythonic and efficient way to perform this mapping.

def tokens2IDs(list_of_tokens, tokens2IDs_dict):
    new_sentence = []
    for token in list_of_tokens:
        if len(token) == 0 or token == " ":
               continue
        new_sentence.append(tokens2IDs_dict.get(token, tokens2IDs_dict.get('UNK')))
    return new_sentence
\$\endgroup\$
  • \$\begingroup\$ If you want to go through a word you don't have a choice to actually go through each character meaning however you want to approach the problem you'll face some kind of iterator (loop, list comprehension etc) \$\endgroup\$ – IEatBagels May 4 '18 at 19:39
5
\$\begingroup\$

There is nothing wrong in general with for loops. If you have to iterate over some elements there is no magic that can avoid this. So why is your function "not very efficient"? And why is this because of the for-loop?

However you asked for a pythonic solution. In Python explicit for loops are avoided by using list comprehension. Internally Python again loops over all elements, however slightly more efficient. Better readability is the other big gain.

So let's measure (I did some renaming for readability). I compare

  • your original explicit for loop
  • a 1:1 translation to a list comprehension
  • further improved by extraction of the default lookup
  • also apply Jean-François Fabres strip()

.

import timeit

def encode(tokens, dictionary):
    ids = []
    for token in tokens:
        if len(token) == 0 or token == " ":
            continue
        ids.append(dictionary.get(token, dictionary.get('UNK')))
    return ids

def encode1(tokens, dictionary):
    ids = [dictionary.get(t, dictionary.get("UNK")) for t in tokens if t and t != " "]
    return ids

def encode2(tokens, dictionary):
    unk = dictionary.get("UNK")
    ids = [dictionary.get(t, unk) for t in tokens if t and t != " "]
    return ids

def encode3(tokens, dictionary):
    unk = dictionary.get("UNK")
    ids = [dictionary.get(t, unk) for t in tokens if t.strip()]
    return ids

def test(tokens, dictionary, result, encoders):
    for encode in encoders:
        assert result == encode(tokens, dictionary)

def time(tokens, dictionary, encoders):
    print(tokens)
    print(encoders[0](tokens, dictionary))
    for encode in encoders:
        print(timeit.timeit(lambda: encode(tokens, dictionary)))

encoders = [encode, encode1, encode2, encode3]

tokens = ["asdf", "qwer", "yxcv", " ", ""]
dictionary = {"UNK": 0, "asdf": 1, "qwer": 2}
test(tokens, dictionary, [1, 2, 0], encoders)

string = "asdwerasfeqtdfgrtzdfhetzufgjhtzui"
tokens = list(string)
dictionary = {k: v for v, k in enumerate(["UNK"] + list(dict.fromkeys(tokens)))}
time(tokens, dictionary, encoders)

tokens = [" "] * len(string)
time(tokens, dictionary, encoders)

tokens = [""] * len(string)
time(tokens, dictionary, encoders)

results on my machine

['a', 's', 'd', 'w', 'e', 'r', 'a', 's', 'f', 'e', 'q', 't', 'd', 'f', 'g', 'r', 't', 'z', 'd', 'f', 'h', 'e', 't', 'z', 'u', 'f', 'g', 'j', 'h', 't', 'z', 'u', 'i']
[7, 9, 6, 11, 10, 12, 7, 9, 8, 10, 5, 14, 6, 8, 15, 12, 14, 1, 6, 8, 2, 10, 14, 1, 3, 8, 15, 4, 2, 14, 1, 3, 13]
7.497583819553256
5.152629452757537
3.7371353656053543
4.100516518577933
[' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ']
[]
1.744785476475954
1.0442649777978659
1.1422674069181085
1.721993056125939
['', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '', '']
[]
1.2569031789898872
0.6375910267233849
0.7521892953664064
1.699515514075756

which shows:

  • list comprehension is indeed significantly faster than the explicit for-loop
  • avoiding the evaluation of the default also has significant impact
  • strip() is a little slower than your explicit test (while more generic)
  • for the explicit tests the order of evaluation is important

Altogether this is doubling the speed, but there is no magic to find magnitudes.

\$\endgroup\$
4
\$\begingroup\$

The code can be packed into a list comprehension:

def tokens2IDs(list_of_tokens, tokens2IDs_dict):
    dictget = tokens2IDs_dict.get
    default = dictget('UNK')
    return [dictget(token, default) for token in list_of_tokens if token.strip()]

The comprehension filters out the empty/blank tokens thanks to testing of the "truthy" value of the stripped token string.

So the comprehension should be slightly faster than the classical loop, and I pre-computed the default value so it's only fetched once.

Another trick is to store the dictionary get function reference in local variable: one less lookup to perform.

Note that depending on the data (a lot of blank tokens vs a lot of empty tokens) it may be better to change the following:

if token.strip()   # faster if there are a lot of blank tokens

by

if token and token.strip()  # faster if there are a lot of empty tokens

the latter short-circuits on an empty string, so strip isn't called at all.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.