1. 1
2. 2
3. 3
4. 4
5. 5
6. 6
7. 7
8. 8
9. 9
10. 10
11. 11
12. 12
13. 13
14. 14
15. 15
16. 16
17. 17
18. 18
19. 19
20. 20
21. 21
22. 22
23. 23
24. 24
25. 25
26. 26
27. 27
28. 28

import numpy as np
import torch


def predict(device, dataset, model, text, next_words=100, top_k=5):
    model.eval()

    words = text.split(' ')

    state_h, state_c = model.zero_state(1)
    state_h = state_h.to(device)
    state_c = state_c.to(device)

    for w in words:
        ix = torch.tensor([[dataset.word_to_index[word]]]).to(device)
        output, (state_h, state_c) = model(ix, (state_h, state_c))
    
    for i in range(next_words):
        _, top_ix = torch.topk(output[0], k=top_k)
        choices = top_ix.tolist()
        choice = np.random.choice(choices[0])

        words.append(dataset.index_to_word[choice])

        ix = torch.tensor([[choice]]).to(device)
        output, (state_h, state_c) = net(ix, (state_h, state_c))

    return words