Suppose there's a 3-level nested list, for example,

[[['<s>', '<s>', 'eu', 'rejects', 'german'], ['<s>', 'eu', 'rejects', 'german', 'call']],\
 [['eu', 'rejects', 'german', 'call', 'to'], ['rejects', 'german', 'call', 'to', 'boycott']]]
  • level 1: sentence level
  • level 2: context window level
  • level 3: word/token level

Task: I'd like to map each token to its corresponding index with a pre-defined Python dictionary (e.g, word2index['xx']=4).

Output example:

[[[2522, 2522, 475, 1620, 397], [2522, 475, 1620, 397, 439]],\
 [[475, 1620, 397, 439, 4], [1620, 397, 439, 4, 1443]]]

The following is what I've done,

def word2index(input, model):
    """ Map words to index    
    input2index = []
    for sent in input:
        sent2index = []
        for window in sent:
            window2index = []
            for elem in window:
                    index = model.vocab[elem].index
                except KeyError:
                    unk = model.vocab['unk'].index
    return input2index

which involves lots of loops and creations of new list as intermediate containers. It'd be very easy to lose the track forgetting which container to append current data, especially when the nested list gets deeper.

Question: Would there be a more cleaner way to code such task?


1 Answer 1


You can use the dict.get(key, default) method to provide a default value instead of catching an Exception. This makes your computation a single expression, which is good for further inlining:

for elem in window:
    index = model.vocab.get(elem, unk).index

Which can be rewritten as:

window2index = [model.vocab.get(elem, unk).index for elem in window]

Which is faster in Cpython, and fairly straightforward.

Now, of course, you have a similar pattern for sent, and a similar pattern for input:

input2index = [EXPR-USING-sent for sent in input]

sent2index = [EXPR-USING-window for window in sent]

window2index = ... see above ...

So you can stack them all together as:

unk = model.vocab['unk']
input2index = [EXPR-USING-sent for sent in input]
return input2index

which becomes:

unk = model.vocab['unk']
input2index = [[EXPR-USING-window for window in sent] for sent in input]
return input2index

which becomes:

unk = model.vocab['unk']
input2index = [[[model.vocab.get(elem, unk).index
                for elem in window]
                for window in sent] 
                for sent in input]
return input2index
  • 2
    \$\begingroup\$ Great answer. However I pattern matched your last bit of code to flatten a 3d list to a 1d list, before tripping up when I saw the three [[[ at the start of the comprehension. It may be something you want to look into, :) \$\endgroup\$
    – Peilonrayz
    Nov 7, 2017 at 9:28
  • \$\begingroup\$ Nice answer. One little nitpick on the final form. I would consider returning the nested list comprehension instead of assigning it to a variable simply to return it. \$\endgroup\$
    – magu_
    Nov 7, 2017 at 16:11
  • \$\begingroup\$ Actually, I would too. But I know there's going to be some printing of data before it gets returned, just to check that it works :-> \$\endgroup\$
    – aghast
    Nov 7, 2017 at 21:02
  • \$\begingroup\$ Can I set more than one default values by adding a if condition? For instance, when dict[k] raises KeyError, assign a default value of <s> if k<1, else </s>. \$\endgroup\$
    – Logan
    Nov 16, 2017 at 3:57

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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