def retain_top_words(data_matrix, ntop_tokens):
    """
    Given a matrix data (list of np.array) retain the n most frequent tokens
    (ntop_tokens) of each row-vector.
    """
    res_matrix = []
    for row in data_matrix:
        if isinstance(row, list):
            res_matrix.append([el for el in row if el < ntop_tokens])
        else:
            res_matrix.append(row[row < ntop_tokens])
    return res_matrix