test_words_str = ['And', 'the', 'name', 'of']
for _ in range(100):
    # Example: use last 3 words from the text
    test_word_indices = [word2idx.get(w, unk_idx) for w in test_words_str]
    test_heights = [word_heights[w] for w in test_word_indices]
    predictions = predict_next(model, test_word_indices, test_heights, idx2word, unk_idx)
    #print(f"Top predictions for next word:")
    test_words_str.append(predictions[2][0])
    #for word, prob in predictions:
    #    print(f"  {word}: {prob:.3f}")

print(f"\nInput words: {test_words_str}")