add function to get top k most likely next words