Helper function to get training data.

def get_train_data(args):
    """Helper function to get training data."""
    counter = dict()
    with io.open(args.vocab, 'r', encoding='utf-8') as f:
        for line in f:
            token, count = line.split('\t')
            counter[token] = int(count)
    vocab = nlp.Vocab(counter, unknown_token=None, padding_token=None,
                      bos_token=None, eos_token=None, min_freq=1)

    npz = np.load(args.cooccurrences)
    row, col, counts = npz['row'], npz['col'], npz['data']

    rank_dtype = 'int32'
    if row.max() >= np.iinfo(np.int32).max:
        rank_dtype = 'int64'
        # MXNet has no support for uint32, so we must fall back to int64
        logging.info('More words than could be counted using int32. '
                     'Using int64 to represent word indices.')
    row = mx.nd.array(row, dtype=rank_dtype)
    col = mx.nd.array(col, dtype=rank_dtype)
    # row is always used as 'source' and col as 'context' word. Therefore
    # duplicate the entries.

    assert row.shape == col.shape
    row = mx.nd.concatenate([row, col])
    col = mx.nd.concatenate([col, row[:len(row) // 2]])

    counts = mx.nd.array(counts, dtype='float32')
    counts = mx.nd.concatenate([counts, counts])

    return vocab, row, col, counts