LM_generate_most_likely.py

def generate_most_probable(lm, index2word):
    """
    Generates a string, taking the most probable word at each time step.
    :param lm - the language model: a function that gets a string and returns a distribution on the next word
    :param index2word - a mapping from the index of a word in the vocabulary to the word itself
    """    
    generated_sentence = '<s>'
    curr_token = None

    while curr_token != '</s>':
        curr_distribution = lm(generated_sentence)  # vector of probabilities
        sorted_by_probability = np.argsort(curr_distribution)
        curr_token = index2word[int(sorted_by_probability[-1])] # last token is the most probable
        generated_sentence += ' ' + curr_token        

    return generated_sentence
    

generated_str = generate_most_probable(stupid_lm, vocab)
print(generated_str)