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
29. 29
30. 30
31. 31
32. 32
33. 33
34. 34
35. 35
36. 36
37. 37
38. 38
39. 39
40. 40
41. 41
42. 42
43. 43
44. 44
45. 45
46. 46
47. 47
48. 48
49. 49
50. 50
51. 51
52. 52
53. 53
54. 54
55. 55
56. 56
57. 57
58. 58
59. 59
60. 60
61. 61
62. 62
63. 63
64. 64
65. 65
66. 66
67. 67
68. 68
69. 69
70. 70
71. 71
72. 72
73. 73
74. 74
75. 75
76. 76
77. 77
78. 78
79. 79
80. 80
81. 81
82. 82
83. 83
84. 84
85. 85
86. 86
87. 87
88. 88
89. 89
90. 90
91. 91
92. 92
93. 93
94. 94
95. 95
96. 96
97. 97
98. 98
99. 99
100. 100
101. 101
102. 102
103. 103
104. 104
105. 105
106. 106
107. 107
108. 108
109. 109
110. 110
111. 111
112. 112
113. 113
114. 114
115. 115
116. 116
117. 117

#!/usr/bin/python3

import re
from collections import Counter

import psycopg2
import torch
from torch import nn, optim
from torch.utils.data import DataLoader


# Fetch messages from database since it's way faster than using the API
conn = psycopg2.connect(dbname="mastodon_production")
cur = conn.cursor()
cur.execute('SELECT * FROM statuses')
statuses = cur.fetchall()

# Get the messages as plain text
# TODO: Remove punctuation and other junk
text = [re.sub(r'<[^>]*>', '', status[2]) for status in statuses]  # Use regex to remove HTML stuff
#print(text[0:100])


class Model(nn.Module):
    def __init__(self, dataset):
        super(Model, self).__init__()

        self.lstm_size = 128
        self.embedding_dim = 128
        self.num_layers = 3

        n_vocab = len(dataset.uniq_words)
        self.embedding = nn.Embedding(
            num_embeddings=n_vocab,
            embedding_dim=self.embedding_dim
        )
        self.lstm = nn.LSTM(
            input_size=self.lstm_size,
            hidden_size=self.lstm_size,
            num_layers=self.num_layers,
            dropout=0.2
        )
        self.fc = nn.Linear(self.lstm_size, n_vocab)
    
    def forward(self, x, prev_state):
        embed = self.embedding(x)
        output, state = self.lstm(embed, prev_state)
        logits = self.fc(output)

        return logits, state
    
    def init_state(self, sequence_length):
        return (torch.zeros(self.num_layers, sequence_length, self.lstm_size),
                torch.zeros(self.num_layers, sequence_length, self.lstm_size))


class Dataset(torch.utils.data.Dataset):
    def __init__(self):
        self.words = [word for message in text for word in message.split()]
        self.word_counts = Counter(self.words)
        self.uniq_words = sorted(self.word_counts, key=self.word_counts.get)

        self.index_to_word = {index: word for index, word in enumerate(self.uniq_words)}
        self.word_to_index = {word: index for index, word in enumerate(self.uniq_words)}

        self.words_indexes = [self.word_to_index[w] for w in self.words]
    
    def __len__(self):
        return len(self.words_indexes) - 4
    
    def __getitem__(self, index):
        return (torch.tensor(self.words_indexes[index:index+4]),
                torch.tensor(self.words_indexes[index+1:index+4+1]))


dataset = Dataset()
model = Model(dataset)


dataloader = DataLoader(dataset, batch_size=1024)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(1):
    state_h, state_c = model.init_state(4)
    for batch, (x, y) in enumerate(dataloader):
        optimizer.zero_grad()

        y_pred, (state_h, state_c) = model(x, (state_h, state_c))
        loss = criterion(y_pred.transpose(1, 2), y)

        state_h = state_h.detach()
        state_c = state_c.detach()
        
        loss.backward()
        optimizer.step()
        
        print({ 'epoch': epoch, 'batch': batch, 'loss': loss.item() })


def predict(text, next_words=100):
    model.eval()

    words = text.split(' ')
    state_h, state_c = model.init_state(len(words))
    
    for i in range(0, next_words):
        x = torch.tensor([[dataset.word_to_index[w] for w in words[i:]]])
        y_pred, (state_h, state_c) = model(x, (state_h, state_c))
        
        last_word_logits = y_pred[0][-1]
        p = torch.nn.functional.softmax(last_word_logits, dim=0).detach().numpy()
        word_index = np.random.choice(len(last_word_logits), p=p)
        words.append(dataset.index_to_word[word_index])
    
    return words

predict('This is a test')