# Buggy code
def load_embedding_from_path(args):
    """Load a TokenEmbedding."""
    if args.embedding_path.endswith(".bin"):
        with utils.print_time("load fastText model."):
            model = nlp.model.train.FasttextEmbeddingModel.load_fasttext_format(
                args.embedding_path
            )
        idx_to_token = sorted(model._token_to_idx, key=model._token_to_idx.get)

        embedding = nlp.embedding.TokenEmbedding(
            unknown_token=None, unknown_lookup=model, allow_extend=True
        )

        # Analogy task is open-vocabulary, so must keep all known words.
        # But if not evaluating analogy, no need to precompute now as all
        # words for closed vocabulary task can be obtained via the unknown
        # lookup
        if not args.analogy_datasets:
            idx_to_token = []
        elif args.analogy_datasets and args.analogy_max_vocab_size:
            idx_to_token = idx_to_token[: args.analogy_max_vocab_size]

        embedding["<unk>"] = mx.nd.zeros(model.weight.shape[1])
        if idx_to_token:
            with utils.print_time(
                "compute vectors for {} known words.".format(len(idx_to_token))
            ):
                embedding[idx_to_token] = model[idx_to_token]
    else:
        embedding = nlp.embedding.TokenEmbedding.from_file(args.embedding_path)

    return embedding


# Traceback (KeyError)
Traceback (most recent call last):
File "evaluate_pretrained.py", line 216, in <module>
vocab.set_embedding(token_embedding_)
File "/Users/siyuanl/private/gluon-nlp/src/gluonnlp/vocab/vocab.py", line 412, in set_embedding
new_idx_to_vec[1:, col_start:col_end] = embs[self._idx_to_token[1:]]
File "/Users/siyuanl/private/gluon-nlp/src/gluonnlp/embedding/token_embedding.py", line 637, in __getitem__
indices = [self._token_to_idx[token] for token in tokens]
File "/Users/siyuanl/private/gluon-nlp/src/gluonnlp/embedding/token_embedding.py", line 637, in <listcomp>
indices = [self._token_to_idx[token] for token in tokens]
KeyError: '<pad>'

# Fixed code
def load_embedding_from_path(args):
    """Load a TokenEmbedding."""
    if args.embedding_path.endswith(".bin"):
        with utils.print_time("load fastText model."):
            model = nlp.model.train.FasttextEmbeddingModel.load_fasttext_format(
                args.embedding_path
            )
        idx_to_token = sorted(model._token_to_idx, key=model._token_to_idx.get)

        # Analogy task is open-vocabulary, so must keep all known words.
        # But if not evaluating analogy, no need to precompute now as all
        # words for closed vocabulary task can be obtained via the unknown
        # lookup
        if not args.analogy_datasets:
            # TODO(leezu): use shape (0, model.weight.shape[1]) once np shape
            # is supported by TokenEmbedding
            idx_to_token = ["<unk>"]
            idx_to_vec = mx.nd.zeros((1, model.weight.shape[1]))
        else:
            if args.analogy_max_vocab_size:
                idx_to_token = idx_to_token[: args.analogy_max_vocab_size]
            with utils.print_time(
                "compute vectors for {} known words.".format(len(idx_to_token))
            ):
                idx_to_vec = model[idx_to_token]

        embedding = nlp.embedding.TokenEmbedding(
            unknown_token=None,
            idx_to_token=idx_to_token,
            idx_to_vec=idx_to_vec,
            unknown_lookup=model,
        )
    else:
        embedding = nlp.embedding.TokenEmbedding.from_file(args.embedding_path)

    return embedding