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
118. 118
119. 119
120. 120
121. 121
122. 122
123. 123
124. 124
125. 125
126. 126
127. 127
128. 128
129. 129
130. 130
131. 131
132. 132
133. 133
134. 134
135. 135
136. 136
137. 137
138. 138
139. 139
140. 140
141. 141
142. 142
143. 143
144. 144
145. 145
146. 146
147. 147

#!/usr/bin/python3

import re
from collections import Counter

import psycopg2
import numpy as np
import torch
from torch import nn
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()


# Extract all words from statuses
# Use regex to remove HTML stuff
text = [re.sub(r'<[^>]*>', '', status[2]) for status in statuses]
# print(text[0:100])


class Dataset(torch.utils.data.Dataset):
    def __init__(self):
        # Flatten text into words
        self.words = [word for message in text for word in message.split()]
        # Remove URLs and special characters and convert to lowercase
        self.words = [re.sub(r'[^a-z0-9]', '', word.lower()) for word in self.words if word.find('://') == -1]
        
        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()
dataloader = DataLoader(dataset, batch_size=256)
print(len(dataloader))


device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")


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))


model = Model(dataset).to(device)
print(model)

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)


for epoch in range(1):
    model.train()
    
    state_h, state_c = model.init_state(4)

    cnt = 0
    for batch, (x, y) in enumerate(dataloader):
        cnt += 1
        if cnt > 50:
            break

        optimizer.zero_grad()

        # Compute prediction error
        y_pred, (state_h, state_c) = model(x, (state_h, state_c))
        loss = loss_fn(y_pred.transpose(1, 2), y)

        state_h = state_h.detach()
        state_c = state_c.detach()

        # Backpropogation
        optimizer.zero_grad()
        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


print(predict('this is a test'))