def main(_):
if FLAGS.model == 'dmn':
word2vec = load_glove(FLAGS.glove_size)
words = WordTable(word2vec, FLAGS.glove_size)
from models.old.dmn import DMN
elif FLAGS.model == 'dmn+':
words = WordTable()
from models.new.dmn_plus import DMN
elif FLAGS.model == 'mc_dmn+':
word2vec = load_glove(FLAGS.glove_size)
words = WordTable(word2vec, FLAGS.glove_size)
from models.new.mc_dmn_plus import DMN
elif FLAGS.model == 'dmn_embed':
words = WordTable()
from models.old.dmn_embedding import DMN
else:
print('Unknown model type: %s' % FLAGS.model)
return
# Read data
train = read_babi(FLAGS.data_dir, FLAGS.task, 'train', FLAGS.batch_size,
words)
test = read_babi(FLAGS.data_dir, FLAGS.task, 'test', FLAGS.batch_size,
words)
val = train.split_dataset(FLAGS.val_ratio)
FLAGS.max_sent_size, FLAGS.max_ques_size, FLAGS.max_fact_count = get_max_sizes(
train, test, val)
print('Word count: %d, Max sentence len : %d' %
(words.vocab_size, FLAGS.max_sent_size))
# Modify save dir
FLAGS.save_dir += '/task_%d/' % FLAGS.task
if not os.path.exists(FLAGS.save_dir):
os.makedirs(FLAGS.save_dir) #, exist_ok=True)
with tf.Session() as sess:
model = DMN(FLAGS, words)
sess.run(tf.initialize_all_variables())
if FLAGS.test:
model.load(sess)
model.eval(sess, test, name='Test')
else:
if FLAGS.load: model.load(sess)
model.train(sess, train, val)
def main(_):
if FLAGS.model == 'dmn':
word2vec = load_glove(FLAGS.glove_size)
words = WordTable(word2vec, FLAGS.glove_size)
from models.old.dmn import DMN
elif FLAGS.model == 'dmn+':
words = WordTable()
from models.new.dmn_plus import DMN
elif FLAGS.model == 'dmn_embed':
words = WordTable()
from models.old.dmn_embedding import DMN
else:
print('Unknown model type: %s' % FLAGS.model)
return
# Read data
train = read_babi(FLAGS.data_dir, FLAGS.task, 'train', FLAGS.batch_size, words)
test = read_babi(FLAGS.data_dir, FLAGS.task, 'test', FLAGS.batch_size, words)
val = train.split_dataset(FLAGS.val_ratio)
FLAGS.max_sent_size, FLAGS.max_ques_size, FLAGS.max_fact_count = get_max_sizes(train, test, val)
print('Word count: %d, Max sentence len : %d' % (words.vocab_size, FLAGS.max_sent_size))
# Modify save dir
FLAGS.save_dir += '/task_%d/' % FLAGS.task
if not os.path.exists(FLAGS.save_dir):
os.makedirs(FLAGS.save_dir, exist_ok=True)
with tf.Session() as sess:
model = DMN(FLAGS, words)
sess.run(tf.initialize_all_variables())
if FLAGS.test:
model.load(sess)
model.eval(sess, test, name='Test')
else:
if FLAGS.load: model.load(sess)
model.train(sess, train, val)