Szekspir

import torch
import sys
from torch import nn,optim
#Szekspir na znakach
'''
Zadanie

ngram model jezyka na znakach neuronowy
wziac kawal tekstu (obojetnie), wytrenowac model i zrobic generator
'''

history_length = 32
embedding_size = 10
nb_of_char_codes = 128
hidden_size = 100

history_encoded = [ord("\n")] * history_length

device = torch.device('cpu')


def char_source():
    for line in sys.stdin:
        for char in line:
            if ord(char) < nb_of_char_codes:
                yield ord(char)

class NGramLanguageModel(nn.Module):
    def __init__(self, nb_of_char_codes ,history_length, embedding_size,hidden_size):
        super(NGramLanguageModel, self).__init__()

        self.embeddings = nn.Embedding(nb_of_char_codes,embedding_size).to(device)
        self.model = nn.Sequential(
            nn.Linear(history_length * embedding_size, hidden_size),
            nn.Linear(hidden_size, nb_of_char_codes),
            nn.LogSoftmax()
        ).to(device)

    def forward(self,inputs):
        embedded_inputs = self.embeddings(inputs)
        return self.model(embedded_inputs.view(-1))

    def generate(self,to_be_contiuned,n):
        t = ((" ") * history_length + to_be_contiuned)[-history_length:]
        history = [ord(c) for c in t] #zamienia tekst na historie (ciag znakow ciag liczb)

        with torch.no_grad():
            for _ in range(n):
                x = torch.tensor(history,dtype=torch.long)
                y = torch.exp(model(x))

                best = (sorted(range(nb_of_char_codes), key=lambda i: -y[i]))[0:4] #cztery najlepsze znaki

                yb = torch.tensor([
                    y[ix] if ix in best else 0.0
                    for ix in range(nb_of_char_codes)]) #zmodyfikowane y, ktore dla czterych najlepszych ma przypisane prawd
                c = torch.multinomial(yb,1)[0].item()
                t += chr(c)

                history.pop(0)
                history.append(c)
                #Losujemy z rozkładu
                '''
                c = torch.multinomial(y,1)[0].item()

                t += chr(c)
                history.pop(0)
                history.append(c)
                '''
        return t

model = NGramLanguageModel(nb_of_char_codes,history_length,embedding_size,hidden_size)

counter = 0
step  = 1000
losses = []

criterion = nn.NLLLoss()
optimizer = optim.Adam(model.parameters())

for c in char_source():
    x = torch.tensor(history_encoded, dtype=torch.long, device=device)
    model.zero_grad()
    y = model(x)

    loss = criterion(y.view(1,-1), torch.tensor([c],dtype=torch.long,device=device))

    losses.append(loss.item())
    if len(losses) > step:
        losses.pop(0)

    if counter % step == 0:
        avg_loss = sum(losses) / len(losses)
        print(counter)
        print(avg_loss)
        print(y)
        print(model.generate("Machine translation is",200))

    loss.backward()
    optimizer.step()

    history_encoded.pop(0)
    history_encoded.append(c)
    counter += 1