Token Embedding Lookup

Problem

Look up token embeddings from an embedding table.

token_ids is a list of row indices and embedding is a list of vectors. Return the list of embedding vectors selected by the ids, preserving the order of token_ids.

Hint

An embedding table is just a list indexed by token id.

Test cases

case 1

in: [[0,2],[[1,2],[3,4],[5,6]]]

expected: [[1,2],[5,6]]

case 2

in: [[0],[[7,8]]]

expected: [[7,8]]

case 3

in: [[0,0,1],[[1],[2]]]

expected: [[1],[1],[2]]

case 4

in: [[],[[1,2]]]

expected: []

Starter codePython

def embed_tokens(token_ids, embedding):
    # Your code here
    pass

4 test cases · run your code against them in the editor