'Tests training and inference scripts.'

def test_train_infer(self):

    (sources_train, targets_train) = test_utils.create_temp_parallel_data(sources=[u'a    a    a    a', u'b    b    b    b', u'c    c    c    c', u'\u7b11    \u7b11    \u7b11    \u7b11'], targets=[u'b    b    b    b', u'a    a    a    a', u'c    c    c    c', u'\u6ce3    \u6ce3    \u6ce3    \u6ce3'])
    (sources_dev, targets_dev) = test_utils.create_temp_parallel_data(sources=[u'a    a', u'b    b', u'c    c    c', u'\u7b11    \u7b11    \u7b11'], targets=[u'b    b', u'a    a', u'c    c    c', u'\u6ce3    \u6ce3    \u6ce3'])
    vocab_source = test_utils.create_temporary_vocab_file([u'a', u'b', u'c', u'\u7b11'])
    vocab_target = test_utils.create_temporary_vocab_file([u'a', u'b', u'c', u'\u6ce3'])
    _clear_flags()
    tf.reset_default_graph()
    train_script = imp.load_source(u'seq2seq.test.train_bin', os.path.join(BIN_FOLDER, u'train.py'))
    tf.app.flags.FLAGS.output_dir = self.output_dir
    tf.app.flags.FLAGS.hooks = u'\n                        -    class:    PrintModelAnalysisHook\n                        -    class:    MetadataCaptureHook\n                        -    class:    TrainSampleHook\n                '
    tf.app.flags.FLAGS.metrics = u'\n                        -    class:    LogPerplexityMetricSpec\n                        -    class:    BleuMetricSpec\n                        -    class:    RougeMetricSpec\n                                params:\n                                        rouge_type:    rouge_1/f_score\n                '
    tf.app.flags.FLAGS.model = u'AttentionSeq2Seq'
    tf.app.flags.FLAGS.model_params = u'\n                attention.params:\n                        num_units:    10\n                vocab_source:    {}\n                vocab_target:    {}\n                '.format(vocab_source.name, vocab_target.name)
    tf.app.flags.FLAGS.batch_size = 2
    config_path = os.path.join(self.output_dir, u'train_config.yml')
    with gfile.GFile(config_path, u'w') as config_file:
        yaml.dump({u'input_pipeline_train': {u'class': u'ParallelTextInputPipeline', u'params': {u'source_files': [sources_train.name], u'target_files': [targets_train.name]}}, u'input_pipeline_dev': {u'class': u'ParallelTextInputPipeline', u'params': {u'source_files': [sources_dev.name], u'target_files': [targets_dev.name]}}, u'train_steps': 50, u'model_params': {u'embedding.dim': 10, u'decoder.params': {u'rnn_cell': {u'cell_class': u'GRUCell', u'cell_params': {u'num_units': 8}}}, u'encoder.params': {u'rnn_cell': {u'cell_class': u'GRUCell', u'cell_params': {u'num_units': 8}}}}}, config_file)
    tf.app.flags.FLAGS.config_paths = config_path
    tf.logging.set_verbosity(tf.logging.INFO)
    train_script.main([])
    expected_checkpoint = os.path.join(self.output_dir, u'model.ckpt-50.data-00000-of-00001')
    self.assertTrue(os.path.exists(expected_checkpoint))
    _clear_flags()
    tf.reset_default_graph()
    infer_script = imp.load_source(u'seq2seq.test.infer_bin', os.path.join(BIN_FOLDER, u'infer.py'))
    attention_dir = os.path.join(self.output_dir, u'att')
    tf.app.flags.FLAGS.model_dir = self.output_dir
    tf.app.flags.FLAGS.input_pipeline = u'\n                        class:    ParallelTextInputPipeline\n                        params:\n                                source_files:\n                                        -    {}\n                                target_files:\n                                        -    {}\n                '.format(sources_dev.name, targets_dev.name)
    tf.app.flags.FLAGS.batch_size = 2
    tf.app.flags.FLAGS.checkpoint_path = os.path.join(self.output_dir, u'model.ckpt-50')
    tf.app.flags.FLAGS.tasks = u'\n                -    class:    DecodeText\n                -    class:    DumpAttention\n                        params:\n                                output_dir:    {}\n                '.format(attention_dir)
    infer_script.main([])
    self.assertTrue(os.path.exists(os.path.join(attention_dir, u'attention_scores.npz')))
    self.assertTrue(os.path.exists(os.path.join(attention_dir, u'00002.png')))
    scores = np.load(os.path.join(attention_dir, u'attention_scores.npz'))
    self.assertIn(u'arr_0', scores)
    self.assertEqual(scores[u'arr_0'].shape[1], 3)
    self.assertIn(u'arr_1', scores)
    self.assertEqual(scores[u'arr_1'].shape[1], 3)
    self.assertIn(u'arr_2', scores)
    self.assertEqual(scores[u'arr_2'].shape[1], 4)
    self.assertIn(u'arr_3', scores)
    self.assertEqual(scores[u'arr_3'].shape[1], 4)
    _clear_flags()
    tf.reset_default_graph()
    infer_script = imp.load_source(u'seq2seq.test.infer_bin', os.path.join(BIN_FOLDER, u'infer.py'))
    tf.app.flags.FLAGS.model_dir = self.output_dir
    tf.app.flags.FLAGS.input_pipeline = u'\n                        class:    ParallelTextInputPipeline\n                        params:\n                                source_files:\n                                        -    {}\n                                target_files:\n                                        -    {}\n                '.format(sources_dev.name, targets_dev.name)
    tf.app.flags.FLAGS.batch_size = 2
    tf.app.flags.FLAGS.checkpoint_path = os.path.join(self.output_dir, u'model.ckpt-50')
    tf.app.flags.FLAGS.model_params = u'\n                        inference.beam_search.beam_width:    5\n                '
    tf.app.flags.FLAGS.tasks = u'\n                -    class:    DecodeText\n                        params:\n                                postproc_fn:    seq2seq.data.postproc.decode_sentencepiece\n                -    class:    DumpBeams\n                        params:\n                                file:    {}\n                '.format(os.path.join(self.output_dir, u'beams.npz'))
    infer_script.main([])
    self.assertTrue(os.path.exists(os.path.join(self.output_dir, u'beams.npz')))
