forked from davidgolub/QuestionGeneration
-
Notifications
You must be signed in to change notification settings - Fork 0
/
iob_trainer_test.py
62 lines (46 loc) · 1.82 KB
/
iob_trainer_test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
from data_loaders.iob_loader import IOBLoader
from models.iob.iob_model import IOBModel
from helpers import constants, utils
base_directory = 'datasets/iob_test'
data_loader = IOBLoader(base_directory, tokenizer_type=constants.TOKENIZER_SPACE)
data_loader.mix_indices()
config = {
'input_max_length': data_loader.input_max_length,
'vocab_size': data_loader.vocab.size(),
'embeddings_size': 25,
'hidden_size': 30,
'out_size': 5,
'num_classes': data_loader.label_vocab.size(),
'batch_size': 3,
'learning_rate': 1e-2,
'save_path': 'iob/logs'}
config_path = 'iob/logs/config.json'
params_path = 'iob/logs/model_params.ckpt'
model = IOBModel(config, embeddings=None)
model.save(config_path, params_path)
model.restore(params_path)
batch = data_loader.get_batch(constants.DATASET_TRAIN, config['batch_size'])
for i in range(0, 100):
while batch is not None:
loss, predictions = model.forward(batch)
batch = data_loader.get_batch(constants.DATASET_TRAIN, config['batch_size'])
print(predictions)
print(loss)
if i % 3 == 0:
data_loader.reset_indices()
total_predictions = []
while True:
batch = data_loader.get_batch(constants.DATASET_TEST, config['batch_size'])
if batch is None:
break
predictions = model.predict(batch)
texts = data_loader.label_vocab.tokens_list(predictions)
for i in range(0, len(texts)):
cur_input_length = batch['input_lengths'][i]
cur_text = texts[i]
text_str = " ".join(cur_text[0:cur_input_length])
total_predictions.append(text_str)
utils.save_lines(total_predictions, \
'%s/predictions_test_%s.txt' % (config['save_path'], i))
data_loader.mix_indices()
batch = data_loader.get_batch(constants.DATASET_TRAIN, config['batch_size'])