seq2seq_example.py 文件源码

def test_main2():
    '''
    Integration test - training then prediction
    '''
    import tempfile
    tempdir = tempfile.mkdtemp()
    arglist = "--data-dir %s -e 2 --iter-num=1 -v -v --tensorboard-verbose=1 train 5000" % tempdir
    arglist = arglist.split(' ')
    tf.reset_default_graph()
    ts2s = CommandLine(arglist=arglist)
    wfn = ts2s.weights_output_fn
    assert os.path.exists(wfn)

    arglist = "-i %s predict 1 2 3 4 5 6 7 8 9 0" % wfn
    arglist = arglist.split(' ')
    tf.reset_default_graph()
    ts2s = CommandLine(arglist=arglist)
    assert len(ts2s.prediction_results[0][0])==10

    os.system("rm -rf %s" % tempdir)