277 lines
6.7 KiB
Python
277 lines
6.7 KiB
Python
#!/usr/bin/env python
|
|
# coding: utf-8
|
|
|
|
# In[1]:
|
|
|
|
|
|
get_ipython().system('git clone --single-branch git://gonito.net/challenging-america-word-gap-prediction -b master')
|
|
|
|
|
|
# In[2]:
|
|
|
|
|
|
from torch import device as dev
|
|
|
|
device = dev("cuda")
|
|
|
|
|
|
# In[3]:
|
|
|
|
|
|
import lzma
|
|
def read_xz_file(fname):
|
|
with lzma.open(fname, mode='rt', encoding='utf-8') as f:
|
|
return [line.strip() for line in f.readlines()]
|
|
|
|
|
|
def read_file(fname):
|
|
with open(fname, mode='rt', encoding='utf-8') as f:
|
|
return [line.strip() for line in f.readlines()]
|
|
|
|
|
|
def get_contexts(input_text):
|
|
all_fields = input_text.replace(r'\n', ' ').split('\t')
|
|
return {'left': all_fields[6], 'right': all_fields[7]}
|
|
|
|
|
|
def compose_sentences(raw_input, labels):
|
|
result = []
|
|
for input, label in zip(raw_input, labels):
|
|
context = get_contexts(input)
|
|
result.append(f'{context["left"]} {input} {context["right"]}')
|
|
return result
|
|
|
|
|
|
# In[4]:
|
|
|
|
|
|
train_input_raw = read_xz_file('challenging-america-word-gap-prediction/train/in.tsv.xz')
|
|
train_labels = read_file('challenging-america-word-gap-prediction/train/expected.tsv')
|
|
train_sentences = compose_sentences(train_input_raw, train_labels)
|
|
|
|
|
|
# In[5]:
|
|
|
|
|
|
from torchtext.data import get_tokenizer
|
|
from torchtext.vocab import build_vocab_from_iterator
|
|
|
|
from torch import save as save_model
|
|
|
|
def tokenize_dataset(lines, tokenizer):
|
|
for line in lines:
|
|
yield tokenizer(line)
|
|
|
|
vocabulary_max_size = 16384
|
|
unknown_token = '<0>'
|
|
tokenizer = get_tokenizer('basic_english')
|
|
vocabulary = build_vocab_from_iterator(
|
|
tokenize_dataset(train_sentences, tokenizer),
|
|
specials=[unknown_token],
|
|
max_tokens=vocabulary_max_size
|
|
)
|
|
vocabulary.set_default_index(vocabulary[unknown_token])
|
|
save_model(vocabulary, 'vocabulary.pth')
|
|
|
|
|
|
# In[6]:
|
|
|
|
|
|
from torch import LongTensor
|
|
|
|
|
|
class TrigramDataset:
|
|
def __init__(self, lines, vocab, tokenizer, unknown_token):
|
|
self.unknown_token = unknown_token
|
|
self.vocab = vocab
|
|
self.tokenizer = tokenizer
|
|
self.lines = lines
|
|
|
|
def __getitem__(self, idx):
|
|
x = []
|
|
y = []
|
|
sentence = [self.vocab[token] for token in self.tokenizer(self.lines[idx])]
|
|
for pos, _ in enumerate(sentence):
|
|
prev = sentence[pos-1] if pos > 0 else self.vocab[self.unknown_token]
|
|
current = sentence[pos]
|
|
next = sentence[pos+1] if pos < len(sentence) - 1 else self.vocab[self.unknown_token]
|
|
x.append([prev, next])
|
|
y.append([current])
|
|
return LongTensor(x), LongTensor(y)
|
|
|
|
def __len__(self):
|
|
return len(self.lines)
|
|
|
|
|
|
# In[7]:
|
|
|
|
|
|
train_dataset = TrigramDataset(train_sentences, vocabulary, tokenizer, unknown_token)
|
|
|
|
|
|
# In[8]:
|
|
|
|
|
|
from torch import nn
|
|
|
|
|
|
class LanguageModel(nn.Module):
|
|
grams_count = 3
|
|
def __init__(self, vocabulary_size, embedding_size, hidden_size):
|
|
super(LanguageModel, self).__init__()
|
|
self.embedding_size = embedding_size
|
|
self.embedding = nn.Embedding(vocabulary_size, embedding_size)
|
|
self.layers = nn.Sequential(
|
|
nn.Linear((self.grams_count - 1) * embedding_size, hidden_size),
|
|
nn.ReLU(),
|
|
nn.Linear(hidden_size, vocabulary_size),
|
|
nn.Softmax(dim=1)
|
|
)
|
|
|
|
def forward(self, x):
|
|
x = self.embedding(x).view((-1, (self.grams_count - 1) * self.embedding_size))
|
|
return self.layers(x)
|
|
|
|
|
|
# In[12]:
|
|
|
|
|
|
from torch.optim import Adam
|
|
from torch import log
|
|
from tqdm import tqdm
|
|
|
|
|
|
def train(model, dataset, output_file, epochs):
|
|
optimizer = Adam(model.parameters(), lr=0.00007)
|
|
criterion = nn.NLLLoss()
|
|
model.to(device)
|
|
model.train()
|
|
for epoch in range(epochs):
|
|
for i, (x, y) in(bar := tqdm(enumerate(dataset), total=len(dataset))):
|
|
x = x.to(device)
|
|
y = y.to(device)
|
|
optimizer.zero_grad()
|
|
ypredicted = model(x)
|
|
loss = criterion(log(ypredicted), y[:,0])
|
|
if not i % 100:
|
|
bar.set_description(f'Epoch: {epoch}, Loss: {loss}, Batch: {i}')
|
|
loss.backward()
|
|
try:
|
|
nn.utils.clip_grad_norm_(model.parameters(), 5, error_if_nonfinite=True)
|
|
optimizer.step()
|
|
except RuntimeError:
|
|
print("Grad overflow")
|
|
|
|
save_model(model.state_dict(), output_file)
|
|
|
|
|
|
# In[13]:
|
|
|
|
|
|
embedding_size = 256
|
|
model = LanguageModel(len(vocabulary), embedding_size, 128)
|
|
|
|
|
|
# In[ ]:
|
|
|
|
|
|
train(model, train_dataset, 'test_model', 5)
|
|
|
|
|
|
# In[28]:
|
|
|
|
|
|
embedding_size = 256
|
|
model_512 = LanguageModel(len(vocabulary), embedding_size, 512)
|
|
train(model_512, train_dataset, 'test_model_512', 5)
|
|
|
|
|
|
# In[ ]:
|
|
|
|
|
|
embedding_size = 256
|
|
model = LanguageModel(len(vocabulary), embedding_size, 128)
|
|
|
|
|
|
# In[18]:
|
|
|
|
|
|
dev_input_raw = read_xz_file('challenging-america-word-gap-prediction/dev-0/in.tsv.xz')
|
|
dev_contexts = [get_contexts(t) for t in dev_input_raw]
|
|
test_input_raw = read_xz_file('challenging-america-word-gap-prediction/test-A/in.tsv.xz')
|
|
test_sentences = [get_contexts(t) for t in test_input_raw]
|
|
|
|
|
|
# In[15]:
|
|
|
|
|
|
from torch import load as load_model
|
|
|
|
vocabulary = load_model('vocabulary.pth')
|
|
tokenizer = get_tokenizer('basic_english')
|
|
|
|
#model = load_model('test_model')
|
|
model.eval()
|
|
|
|
|
|
# In[16]:
|
|
|
|
|
|
from torch import LongTensor, topk, log
|
|
from tqdm import tqdm
|
|
|
|
|
|
def predict_words(dataset, tokenizer, vocab, model):
|
|
preds = []
|
|
for entry in tqdm(dataset):
|
|
tokenized_left = tokenizer(entry['left'])
|
|
tokenized_right = tokenizer(entry['right'])
|
|
# [last word from left context, y, first word from right context]
|
|
src = LongTensor([vocab[tokenized_left[-1]], vocab[tokenized_right[0]]]).to(device)
|
|
output = model(src)
|
|
top = topk(output[0], 50)
|
|
probs, tokens = top.values.tolist(), vocab.lookup_tokens(top.indices.tolist())
|
|
current_output = ''
|
|
accumulated_probability = 0
|
|
for prob, token in zip(probs, tokens):
|
|
accumulated_probability += prob
|
|
current_output += f'{token.strip()}:{prob} '
|
|
current_output += f':{1 - accumulated_probability}'
|
|
preds.append(current_output)
|
|
return preds
|
|
|
|
|
|
# In[24]:
|
|
|
|
|
|
preds = predict_words(dev_contexts, tokenizer, vocabulary, model)
|
|
|
|
|
|
# In[25]:
|
|
|
|
|
|
with open('challenging-america-word-gap-prediction/dev-0/out-hidden_size=128.tsv', 'w') as f:
|
|
f.writelines(line + '\n' for line in preds)
|
|
|
|
|
|
# In[26]:
|
|
|
|
|
|
test_preds = predict_words(test_sentences, tokenizer, vocabulary, model)
|
|
with open('challenging-america-word-gap-prediction/test-A/out-hidden_size=128.tsv', 'w') as f:
|
|
f.writelines(line + '\n' for line in test_preds)
|
|
|
|
|
|
# In[29]:
|
|
|
|
|
|
#model_512 = load_model('test_model_512')
|
|
model_512.eval()
|
|
preds_512 = predict_words(dev_contexts, tokenizer, vocabulary, model_512)
|
|
with open('challenging-america-word-gap-prediction/dev-0/out-hidden_size=512.tsv', 'w') as f:
|
|
f.writelines(line + '\n' for line in preds)
|
|
test_preds_512 = predict_words(test_sentences, tokenizer, vocabulary, model_512)
|
|
with open('challenging-america-word-gap-prediction/test-A/out-hidden_size=512.tsv', 'w') as f:
|
|
f.writelines(line + '\n' for line in test_preds)
|
|
|