Changes
1 changed files (+0/-157)
-
old.py (deleted)
-
@@ -1,157 +0,0 @@#!/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 # https://closeheat.com/blog/pytorch-lstm-text-generation-tutorial # https://github.com/ChunML/NLP/tree/master/text_generation/ # https://towardsdatascience.com/simple-text-generation-d1c93f43f340 # https://karpathy.github.io/2019/04/25/recipe/ # 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] # Remove empty strings self.words = [word for word in self.words if word != ''] print(self.words[0:1000]) 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) - 16 def __getitem__(self, index): return (torch.tensor(self.words_indexes[index:index+16]), torch.tensor(self.words_indexes[index+1:index+16+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 = 512 self.embedding_dim = 512 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(2): model.train() state_h, state_c = model.init_state(16) cnt = 0 for batch, (x, y) in enumerate(dataloader): cnt += 1 if cnt > 100: 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('i am'))
-