124 lines
3.0 KiB
Python
124 lines
3.0 KiB
Python
|
|
|
|
from torch import nn
|
|
import torch
|
|
|
|
|
|
from torch.utils.data import IterableDataset
|
|
import itertools
|
|
import lzma
|
|
import regex as re
|
|
import pickle
|
|
|
|
def look_ahead_iterator(gen):
|
|
prev = None
|
|
current = None
|
|
next = None
|
|
for next in gen:
|
|
if prev is not None and current is not None:
|
|
yield (prev, current, next)
|
|
prev = current
|
|
current = next
|
|
|
|
def get_words_from_line(line):
|
|
line = line.rstrip()
|
|
yield '<s>'
|
|
for m in re.finditer(r'[\p{L}0-9\*]+|\p{P}+', line):
|
|
yield m.group(0).lower()
|
|
yield '</s>'
|
|
|
|
|
|
def get_word_lines_from_file(file_name):
|
|
counter=0
|
|
with lzma.open(file_name, 'r') as fh:
|
|
for line in fh:
|
|
counter+=1
|
|
if counter == 100000:
|
|
break
|
|
line = line.decode("utf-8")
|
|
yield get_words_from_line(line)
|
|
|
|
|
|
|
|
class Trigrams(IterableDataset):
|
|
def load_vocab(self):
|
|
with open("vocab.pickle", 'rb') as handle:
|
|
vocab = pickle.load( handle)
|
|
return vocab
|
|
|
|
def __init__(self, text_file, vocabulary_size):
|
|
self.vocab = self.load_vocab()
|
|
self.vocab.set_default_index(self.vocab['<unk>'])
|
|
self.vocabulary_size = vocabulary_size
|
|
self.text_file = text_file
|
|
|
|
def __iter__(self):
|
|
return look_ahead_iterator(
|
|
(self.vocab[t] for t in itertools.chain.from_iterable(get_word_lines_from_file(self.text_file))))
|
|
|
|
vocab_size = 20000
|
|
|
|
train_dataset = Trigrams('train/in.tsv.xz', vocab_size)
|
|
|
|
|
|
|
|
#=== trenowanie
|
|
from torch import nn
|
|
import torch
|
|
from torch.utils.data import DataLoader
|
|
embed_size = 100
|
|
|
|
class SimpleTrigramNeuralLanguageModel(nn.Module):
|
|
def __init__(self, vocabulary_size, embedding_size):
|
|
super(SimpleTrigramNeuralLanguageModel, self).__init__()
|
|
self.embedings = nn.Embedding(vocabulary_size, embedding_size)
|
|
self.linear = nn.Linear(embedding_size*2, vocabulary_size)
|
|
self.softmax = nn.Softmax()
|
|
|
|
# self.model = nn.Sequential(
|
|
# nn.Embedding(vocabulary_size, embedding_size),
|
|
# nn.Linear(embedding_size, vocabulary_size),
|
|
# nn.Softmax()
|
|
# )
|
|
|
|
def forward(self, x):
|
|
emb_1 = self.embedings(x[0])
|
|
emb_2 = self.embedings(x[1])
|
|
|
|
concated = self.linear(torch.cat((emb_1, emb_2), dim=1))
|
|
y = self.softmax(concated)
|
|
|
|
return y
|
|
|
|
model = SimpleTrigramNeuralLanguageModel(vocab_size, embed_size)
|
|
|
|
vocab = train_dataset.vocab
|
|
|
|
vocab.set_default_index(vocab['<unk>'])
|
|
ixs = torch.tensor(vocab.forward(['pies']))
|
|
# out[0][vocab['jest']]
|
|
|
|
|
|
device = 'cpu'
|
|
model = SimpleTrigramNeuralLanguageModel(vocab_size, embed_size).to(device)
|
|
data = DataLoader(train_dataset, batch_size=5000)
|
|
optimizer = torch.optim.Adam(model.parameters())
|
|
criterion = torch.nn.NLLLoss()
|
|
|
|
model.train()
|
|
step = 0
|
|
for x, y, z in data:
|
|
x = x.to(device)
|
|
y = y.to(device)
|
|
z = z.to(device)
|
|
optimizer.zero_grad()
|
|
ypredicted = model([x, z])
|
|
loss = criterion(torch.log(ypredicted), y)
|
|
if step % 100 == 0:
|
|
print(step, loss)
|
|
torch.save(model.state_dict(), f'model1_{step}.bin')
|
|
step += 1
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
torch.save(model.state_dict(), 'model_tri1.bin') |