99 lines
2.7 KiB
Python
99 lines
2.7 KiB
Python
|
#!/usr/bin/python3
|
||
|
|
|
||
|
|
# https://pytorch.org/tutorials/beginner/nlp/word_embeddings_tutorial.html
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from torch import nn, optim
|
||
|
|
|
||
|
|
history_length = 32
|
||
|
|
history_encoded = [ord('\n')] * history_length
|
||
|
|
nb_of_char_codes = 128
|
||
|
|
embedding_size = 30
|
||
|
|
step = 1000
|
||
|
|
|
||
|
|
device = torch.device('cpu')
|
||
|
|
|
||
|
|
f = open('shakespeare.txt')
|
||
|
|
def char_source():
|
||
|
|
for line in f:
|
||
|
|
for c in line:
|
||
|
|
c_code = ord(c)
|
||
|
|
if c_code < nb_of_char_codes:
|
||
|
|
yield(c_code)
|
||
|
|
|
||
|
|
|
||
|
|
class EncoderRNN(nn.Module):
|
||
|
|
def __init__(self, input_size, hidden_size, embedding_size):
|
||
|
|
super(EncoderRNN, self).__init__()
|
||
|
|
self.input_size = input_size
|
||
|
|
self.hidden_size = hidden_size
|
||
|
|
|
||
|
|
self.embedding = nn.Embedding(input_size, embedding_size)
|
||
|
|
self.gru = nn.GRU(embedding_size, nb_of_char_codes)
|
||
|
|
self.softmax = nn.LogSoftmax(dim=1)
|
||
|
|
|
||
|
|
def forward(self, input, hidden):
|
||
|
|
embedded = self.embedding(input)
|
||
|
|
output = embedded
|
||
|
|
output, hidden = self.gru(output, hidden)
|
||
|
|
output = self.softmax(output)
|
||
|
|
return output, hidden
|
||
|
|
|
||
|
|
def initHidden(self):
|
||
|
|
return torch.zeros(1, self.hidden_size, self.input_size, device=device)
|
||
|
|
|
||
|
|
def generate(self, n, encoder_hidden):
|
||
|
|
t = (" " * 200 + "To be or not to be")[-history_length:]
|
||
|
|
history = [ord(c) for c in t]
|
||
|
|
|
||
|
|
with torch.no_grad():
|
||
|
|
for _ in range(n):
|
||
|
|
x = torch.tensor(history, dtype=torch.long, device=device)
|
||
|
|
x = x.unsqueeze(0)
|
||
|
|
y = model(x,encoder_hidden)[0][:,-1,:][0]
|
||
|
|
y = torch.exp(y)
|
||
|
|
best = (sorted(range(nb_of_char_codes), key=lambda i: -y[i]))[0:2]
|
||
|
|
yb = torch.tensor([(y[ix] if ix in best else 0.0) for ix in range(nb_of_char_codes)])
|
||
|
|
c = torch.multinomial(yb, 1)[0].item()
|
||
|
|
|
||
|
|
t += chr(c)
|
||
|
|
|
||
|
|
history.pop(0)
|
||
|
|
history.append(c)
|
||
|
|
|
||
|
|
print(t)
|
||
|
|
|
||
|
|
|
||
|
|
model = EncoderRNN(nb_of_char_codes, history_length, embedding_size).to(device)
|
||
|
|
criterion = nn.NLLLoss().to(device)
|
||
|
|
optimizer = optim.Adam(model.parameters())
|
||
|
|
|
||
|
|
|
||
|
|
counter = 0
|
||
|
|
losses = []
|
||
|
|
|
||
|
|
for c in char_source():
|
||
|
|
x = torch.tensor(history_encoded, dtype=torch.long, device=device)
|
||
|
|
model.zero_grad()
|
||
|
|
x = x.unsqueeze(0)
|
||
|
|
encoder_hidden = model.initHidden()
|
||
|
|
y = model(x,encoder_hidden)[0][:,-1,:]
|
||
|
|
|
||
|
|
loss = criterion(y, torch.tensor([c]).to(device))
|
||
|
|
losses += [loss.item()]
|
||
|
|
if len(losses) > step:
|
||
|
|
losses.pop(0)
|
||
|
|
|
||
|
|
counter += 1
|
||
|
|
|
||
|
|
if counter % step == 0:
|
||
|
|
avg_loss = sum(losses)/len(losses)
|
||
|
|
print(f"{counter}: {loss} {avg_loss}")
|
||
|
|
model.generate(200, encoder_hidden)
|
||
|
|
|
||
|
|
loss.backward()
|
||
|
|
optimizer.step()
|
||
|
|
|
||
|
|
history_encoded.pop(0)
|
||
|
|
history_encoded.append(c)
|