import argparse
import distutils.util

import paddle.v2 as paddle
from network_conf import DSSM
import reader
from utils import TaskType, load_dic, logger, ModelType, ModelArch, display_args

parser = argparse.ArgumentParser(description="PaddlePaddle DSSM example")

parser.add_argument(
    "-i",
    "--train_data_path",
    type=str,
    required=False,
    help="The path of training data.")
parser.add_argument(
    "-t",
    "--test_data_path",
    type=str,
    required=False,
    help="The path of testing data.")
parser.add_argument(
    "-s",
    "--source_dic_path",
    type=str,
    required=False,
    help="The path of the source's word dictionary.")
parser.add_argument(
    "--target_dic_path",
    type=str,
    required=False,
    help=("The path of the target's word dictionary, "
          "if this parameter is not set, the `source_dic_path` will be used"))
parser.add_argument(
    "-b",
    "--batch_size",
    type=int,
    default=32,
    help="The size of mini-batch (default:32).")
parser.add_argument(
    "-p",
    "--num_passes",
    type=int,
    default=10,
    help="The number of passes to run(default:10).")
parser.add_argument(
    "-y",
    "--model_type",
    type=int,
    required=True,
    default=ModelType.CLASSIFICATION_MODE,
    help=("model type, %d for classification, %d for pairwise rank, "
          "%d for regression (default: classification).") %
    (ModelType.CLASSIFICATION_MODE, ModelType.RANK_MODE,
     ModelType.REGRESSION_MODE))
parser.add_argument(
    "-a",
    "--model_arch",
    type=int,
    required=True,
    default=ModelArch.CNN_MODE,
    help="The model architecture, %d for CNN, %d for FC, %d for RNN." %
    (ModelArch.CNN_MODE, ModelArch.FC_MODE, ModelArch.RNN_MODE))
parser.add_argument(
    "--share_network_between_source_target",
    type=distutils.util.strtobool,
    default=False,
    help="Whether to share network parameters between source and target.")
parser.add_argument(
    "--share_embed",
    type=distutils.util.strtobool,
    default=False,
    help="Whether to share word embedding between source and target.")
parser.add_argument(
    "--dnn_dims",
    type=str,
    default="256,128,64,32",
    help=("The dimentions of dnn layers, default is '256,128,64,32', "
          "which means create a 4-layer dnn. The dimention of each layer is "
          "'256, 128, 64 and 32'."))
parser.add_argument(
    "--num_workers",
    type=int,
    default=1,
    help="The number of worker threads, default 1.")
parser.add_argument(
    "--use_gpu",
    type=distutils.util.strtobool,
    default=False,
    help="Whether to use GPU devices (default: False)")
parser.add_argument(
    "-c",
    "--class_num",
    type=int,
    default=0,
    help="The number of categories for classification task.")
parser.add_argument(
    "--model_output_prefix",
    type=str,
    default="./",
    help="The prefix of the path to store the trained models (default: ./).")
parser.add_argument(
    "-g",
    "--num_batches_to_log",
    type=int,
    default=100,
    help=("The log period. Every num_batches_to_test batches, "
          "a training log will be printed. (default: 100)"))
parser.add_argument(
    "-e",
    "--num_batches_to_test",
    type=int,
    default=200,
    help=("The test period. Every num_batches_to_save_model batches, "
          "the specified test sample will be test (default: 200)."))
parser.add_argument(
    "-z",
    "--num_batches_to_save_model",
    type=int,
    default=400,
    help=("Every num_batches_to_save_model batches, "
          "a trained model will be saved (default: 400)."))

args = parser.parse_args()
args.model_type = ModelType(args.model_type)
args.model_arch = ModelArch(args.model_arch)
if args.model_type.is_classification():
    assert args.class_num > 1, ("The parameter class_num should be set in "
                                "classification task.")

layer_dims = [int(i) for i in args.dnn_dims.split(",")]
args.target_dic_path = args.source_dic_path if not \
        args.target_dic_path else args.target_dic_path


def train(train_data_path=None,
          test_data_path=None,
          source_dic_path=None,
          target_dic_path=None,
          model_type=ModelType.create_classification(),
          model_arch=ModelArch.create_cnn(),
          batch_size=32,
          num_passes=10,
          share_semantic_generator=False,
          share_embed=False,
          class_num=None,
          num_workers=1,
          use_gpu=False):
    """
    Train the DSSM.
    """
    default_train_path = "./data/rank/train.txt"
    default_test_path = "./data/rank/test.txt"
    default_dic_path = "./data/vocab.txt"
    if not model_type.is_rank():
        default_train_path = "./data/classification/train.txt"
        default_test_path = "./data/classification/test.txt"

    use_default_data = not train_data_path

    if use_default_data:
        train_data_path = default_train_path
        test_data_path = default_test_path
        source_dic_path = default_dic_path
        target_dic_path = default_dic_path

    dataset = reader.Dataset(
        train_path=train_data_path,
        test_path=test_data_path,
        source_dic_path=source_dic_path,
        target_dic_path=target_dic_path,
        model_type=model_type, )

    train_reader = paddle.batch(
        paddle.reader.shuffle(
            dataset.train, buf_size=1000),
        batch_size=batch_size)

    test_reader = paddle.batch(
        paddle.reader.shuffle(
            dataset.test, buf_size=1000),
        batch_size=batch_size)

    paddle.init(use_gpu=use_gpu, trainer_count=num_workers)

    cost, prediction, label = DSSM(
        dnn_dims=layer_dims,
        vocab_sizes=[
            len(load_dic(path)) for path in [source_dic_path, target_dic_path]
        ],
        model_type=model_type,
        model_arch=model_arch,
        share_semantic_generator=share_semantic_generator,
        class_num=class_num,
        share_embed=share_embed)()

    parameters = paddle.parameters.create(cost)

    adam_optimizer = paddle.optimizer.Adam(
        learning_rate=2e-4,
        regularization=paddle.optimizer.L2Regularization(rate=1e-3),
        model_average=paddle.optimizer.ModelAverage(average_window=0.5))

    trainer = paddle.trainer.SGD(
        cost=cost,
        extra_layers=paddle.evaluator.auc(input=prediction, label=label)
        if not model_type.is_rank() else None,
        parameters=parameters,
        update_equation=adam_optimizer)

    feeding = {}
    if model_type.is_classification() or model_type.is_regression():
        feeding = {"source_input": 0, "target_input": 1, "label_input": 2}
    else:
        feeding = {
            "source_input": 0,
            "left_target_input": 1,
            "right_target_input": 2,
            "label_input": 3
        }

    def _event_handler(event):
        """
        Define batch handler
        """
        if isinstance(event, paddle.event.EndIteration):
            # output train log
            if event.batch_id % args.num_batches_to_log == 0:
                logger.info("Pass %d, Batch %d, Cost %f, %s" % (
                    event.pass_id, event.batch_id, event.cost, event.metrics))

            # test model
            if event.batch_id > 0 and \
                    event.batch_id % args.num_batches_to_test == 0:
                if test_reader is not None:
                    if model_type.is_classification():
                        result = trainer.test(
                            reader=test_reader, feeding=feeding)
                        logger.info("Test at Pass %d, %s" % (event.pass_id,
                                                             result.metrics))
                    else:
                        result = None
            # save model
            if event.batch_id > 0 and \
                    event.batch_id % args.num_batches_to_save_model == 0:
                model_desc = "{type}_{arch}".format(
                    type=str(args.model_type), arch=str(args.model_arch))
                with open("%sdssm_%s_pass_%05d.tar" %
                          (args.model_output_prefix, model_desc,
                           event.pass_id), "w") as f:
                    trainer.save_parameter_to_tar(f)

    trainer.train(
        reader=train_reader,
        event_handler=_event_handler,
        feeding=feeding,
        num_passes=num_passes)

    logger.info("Training has finished.")


if __name__ == "__main__":
    display_args(args)
    train(
        train_data_path=args.train_data_path,
        test_data_path=args.test_data_path,
        source_dic_path=args.source_dic_path,
        target_dic_path=args.target_dic_path,
        model_type=ModelType(args.model_type),
        model_arch=ModelArch(args.model_arch),
        batch_size=args.batch_size,
        num_passes=args.num_passes,
        share_semantic_generator=args.share_network_between_source_target,
        share_embed=args.share_embed,
        class_num=args.class_num,
        num_workers=args.num_workers,
        use_gpu=args.use_gpu)

db49c7f6c0b6ada71cd197eac61d0287607be4f6
<mask token>

<mask token>
parser.add_argument('-i', '--train_data_path', type=str, required=False,
    help='The path of training data.')
parser.add_argument('-t', '--test_data_path', type=str, required=False,
    help='The path of testing data.')
parser.add_argument('-s', '--source_dic_path', type=str, required=False,
    help="The path of the source's word dictionary.")
parser.add_argument('--target_dic_path', type=str, required=False, help=
    "The path of the target's word dictionary, if this parameter is not set, the `source_dic_path` will be used"
    )
parser.add_argument('-b', '--batch_size', type=int, default=32, help=
    'The size of mini-batch (default:32).')
parser.add_argument('-p', '--num_passes', type=int, default=10, help=
    'The number of passes to run(default:10).')
parser.add_argument('-y', '--model_type', type=int, required=True, default=
    ModelType.CLASSIFICATION_MODE, help=
    'model type, %d for classification, %d for pairwise rank, %d for regression (default: classification).'
     % (ModelType.CLASSIFICATION_MODE, ModelType.RANK_MODE, ModelType.
    REGRESSION_MODE))
parser.add_argument('-a', '--model_arch', type=int, required=True, default=
    ModelArch.CNN_MODE, help=
    'The model architecture, %d for CNN, %d for FC, %d for RNN.' % (
    ModelArch.CNN_MODE, ModelArch.FC_MODE, ModelArch.RNN_MODE))
parser.add_argument('--share_network_between_source_target', type=distutils
    .util.strtobool, default=False, help=
    'Whether to share network parameters between source and target.')
parser.add_argument('--share_embed', type=distutils.util.strtobool, default
    =False, help='Whether to share word embedding between source and target.')
parser.add_argument('--dnn_dims', type=str, default='256,128,64,32', help=
    "The dimentions of dnn layers, default is '256,128,64,32', which means create a 4-layer dnn. The dimention of each layer is '256, 128, 64 and 32'."
    )
parser.add_argument('--num_workers', type=int, default=1, help=
    'The number of worker threads, default 1.')
parser.add_argument('--use_gpu', type=distutils.util.strtobool, default=
    False, help='Whether to use GPU devices (default: False)')
parser.add_argument('-c', '--class_num', type=int, default=0, help=
    'The number of categories for classification task.')
parser.add_argument('--model_output_prefix', type=str, default='./', help=
    'The prefix of the path to store the trained models (default: ./).')
parser.add_argument('-g', '--num_batches_to_log', type=int, default=100,
    help=
    'The log period. Every num_batches_to_test batches, a training log will be printed. (default: 100)'
    )
parser.add_argument('-e', '--num_batches_to_test', type=int, default=200,
    help=
    'The test period. Every num_batches_to_save_model batches, the specified test sample will be test (default: 200).'
    )
parser.add_argument('-z', '--num_batches_to_save_model', type=int, default=
    400, help=
    'Every num_batches_to_save_model batches, a trained model will be saved (default: 400).'
    )
<mask token>
if args.model_type.is_classification():
    assert args.class_num > 1, 'The parameter class_num should be set in classification task.'
<mask token>


def train(train_data_path=None, test_data_path=None, source_dic_path=None,
    target_dic_path=None, model_type=ModelType.create_classification(),
    model_arch=ModelArch.create_cnn(), batch_size=32, num_passes=10,
    share_semantic_generator=False, share_embed=False, class_num=None,
    num_workers=1, use_gpu=False):
    """
    Train the DSSM.
    """
    default_train_path = './data/rank/train.txt'
    default_test_path = './data/rank/test.txt'
    default_dic_path = './data/vocab.txt'
    if not model_type.is_rank():
        default_train_path = './data/classification/train.txt'
        default_test_path = './data/classification/test.txt'
    use_default_data = not train_data_path
    if use_default_data:
        train_data_path = default_train_path
        test_data_path = default_test_path
        source_dic_path = default_dic_path
        target_dic_path = default_dic_path
    dataset = reader.Dataset(train_path=train_data_path, test_path=
        test_data_path, source_dic_path=source_dic_path, target_dic_path=
        target_dic_path, model_type=model_type)
    train_reader = paddle.batch(paddle.reader.shuffle(dataset.train,
        buf_size=1000), batch_size=batch_size)
    test_reader = paddle.batch(paddle.reader.shuffle(dataset.test, buf_size
        =1000), batch_size=batch_size)
    paddle.init(use_gpu=use_gpu, trainer_count=num_workers)
    cost, prediction, label = DSSM(dnn_dims=layer_dims, vocab_sizes=[len(
        load_dic(path)) for path in [source_dic_path, target_dic_path]],
        model_type=model_type, model_arch=model_arch,
        share_semantic_generator=share_semantic_generator, class_num=
        class_num, share_embed=share_embed)()
    parameters = paddle.parameters.create(cost)
    adam_optimizer = paddle.optimizer.Adam(learning_rate=0.0002,
        regularization=paddle.optimizer.L2Regularization(rate=0.001),
        model_average=paddle.optimizer.ModelAverage(average_window=0.5))
    trainer = paddle.trainer.SGD(cost=cost, extra_layers=paddle.evaluator.
        auc(input=prediction, label=label) if not model_type.is_rank() else
        None, parameters=parameters, update_equation=adam_optimizer)
    feeding = {}
    if model_type.is_classification() or model_type.is_regression():
        feeding = {'source_input': 0, 'target_input': 1, 'label_input': 2}
    else:
        feeding = {'source_input': 0, 'left_target_input': 1,
            'right_target_input': 2, 'label_input': 3}

    def _event_handler(event):
        """
        Define batch handler
        """
        if isinstance(event, paddle.event.EndIteration):
            if event.batch_id % args.num_batches_to_log == 0:
                logger.info('Pass %d, Batch %d, Cost %f, %s' % (event.
                    pass_id, event.batch_id, event.cost, event.metrics))
            if (event.batch_id > 0 and event.batch_id % args.
                num_batches_to_test == 0):
                if test_reader is not None:
                    if model_type.is_classification():
                        result = trainer.test(reader=test_reader, feeding=
                            feeding)
                        logger.info('Test at Pass %d, %s' % (event.pass_id,
                            result.metrics))
                    else:
                        result = None
            if (event.batch_id > 0 and event.batch_id % args.
                num_batches_to_save_model == 0):
                model_desc = '{type}_{arch}'.format(type=str(args.
                    model_type), arch=str(args.model_arch))
                with open('%sdssm_%s_pass_%05d.tar' % (args.
                    model_output_prefix, model_desc, event.pass_id), 'w') as f:
                    trainer.save_parameter_to_tar(f)
    trainer.train(reader=train_reader, event_handler=_event_handler,
        feeding=feeding, num_passes=num_passes)
    logger.info('Training has finished.')


if __name__ == '__main__':
    display_args(args)
    train(train_data_path=args.train_data_path, test_data_path=args.
        test_data_path, source_dic_path=args.source_dic_path,
        target_dic_path=args.target_dic_path, model_type=ModelType(args.
        model_type), model_arch=ModelArch(args.model_arch), batch_size=args
        .batch_size, num_passes=args.num_passes, share_semantic_generator=
        args.share_network_between_source_target, share_embed=args.
        share_embed, class_num=args.class_num, num_workers=args.num_workers,
        use_gpu=args.use_gpu)

<mask token>
parser = argparse.ArgumentParser(description='PaddlePaddle DSSM example')
parser.add_argument('-i', '--train_data_path', type=str, required=False,
    help='The path of training data.')
parser.add_argument('-t', '--test_data_path', type=str, required=False,
    help='The path of testing data.')
parser.add_argument('-s', '--source_dic_path', type=str, required=False,
    help="The path of the source's word dictionary.")
parser.add_argument('--target_dic_path', type=str, required=False, help=
    "The path of the target's word dictionary, if this parameter is not set, the `source_dic_path` will be used"
    )
parser.add_argument('-b', '--batch_size', type=int, default=32, help=
    'The size of mini-batch (default:32).')
parser.add_argument('-p', '--num_passes', type=int, default=10, help=
    'The number of passes to run(default:10).')
parser.add_argument('-y', '--model_type', type=int, required=True, default=
    ModelType.CLASSIFICATION_MODE, help=
    'model type, %d for classification, %d for pairwise rank, %d for regression (default: classification).'
     % (ModelType.CLASSIFICATION_MODE, ModelType.RANK_MODE, ModelType.
    REGRESSION_MODE))
parser.add_argument('-a', '--model_arch', type=int, required=True, default=
    ModelArch.CNN_MODE, help=
    'The model architecture, %d for CNN, %d for FC, %d for RNN.' % (
    ModelArch.CNN_MODE, ModelArch.FC_MODE, ModelArch.RNN_MODE))
parser.add_argument('--share_network_between_source_target', type=distutils
    .util.strtobool, default=False, help=
    'Whether to share network parameters between source and target.')
parser.add_argument('--share_embed', type=distutils.util.strtobool, default
    =False, help='Whether to share word embedding between source and target.')
parser.add_argument('--dnn_dims', type=str, default='256,128,64,32', help=
    "The dimentions of dnn layers, default is '256,128,64,32', which means create a 4-layer dnn. The dimention of each layer is '256, 128, 64 and 32'."
    )
parser.add_argument('--num_workers', type=int, default=1, help=
    'The number of worker threads, default 1.')
parser.add_argument('--use_gpu', type=distutils.util.strtobool, default=
    False, help='Whether to use GPU devices (default: False)')
parser.add_argument('-c', '--class_num', type=int, default=0, help=
    'The number of categories for classification task.')
parser.add_argument('--model_output_prefix', type=str, default='./', help=
    'The prefix of the path to store the trained models (default: ./).')
parser.add_argument('-g', '--num_batches_to_log', type=int, default=100,
    help=
    'The log period. Every num_batches_to_test batches, a training log will be printed. (default: 100)'
    )
parser.add_argument('-e', '--num_batches_to_test', type=int, default=200,
    help=
    'The test period. Every num_batches_to_save_model batches, the specified test sample will be test (default: 200).'
    )
parser.add_argument('-z', '--num_batches_to_save_model', type=int, default=
    400, help=
    'Every num_batches_to_save_model batches, a trained model will be saved (default: 400).'
    )
args = parser.parse_args()
args.model_type = ModelType(args.model_type)
args.model_arch = ModelArch(args.model_arch)
if args.model_type.is_classification():
    assert args.class_num > 1, 'The parameter class_num should be set in classification task.'
layer_dims = [int(i) for i in args.dnn_dims.split(',')]
args.target_dic_path = (args.source_dic_path if not args.target_dic_path else
    args.target_dic_path)


def train(train_data_path=None, test_data_path=None, source_dic_path=None,
    target_dic_path=None, model_type=ModelType.create_classification(),
    model_arch=ModelArch.create_cnn(), batch_size=32, num_passes=10,
    share_semantic_generator=False, share_embed=False, class_num=None,
    num_workers=1, use_gpu=False):
    """
    Train the DSSM.
    """
    default_train_path = './data/rank/train.txt'
    default_test_path = './data/rank/test.txt'
    default_dic_path = './data/vocab.txt'
    if not model_type.is_rank():
        default_train_path = './data/classification/train.txt'
        default_test_path = './data/classification/test.txt'
    use_default_data = not train_data_path
    if use_default_data:
        train_data_path = default_train_path
        test_data_path = default_test_path
        source_dic_path = default_dic_path
        target_dic_path = default_dic_path
    dataset = reader.Dataset(train_path=train_data_path, test_path=
        test_data_path, source_dic_path=source_dic_path, target_dic_path=
        target_dic_path, model_type=model_type)
    train_reader = paddle.batch(paddle.reader.shuffle(dataset.train,
        buf_size=1000), batch_size=batch_size)
    test_reader = paddle.batch(paddle.reader.shuffle(dataset.test, buf_size
        =1000), batch_size=batch_size)
    paddle.init(use_gpu=use_gpu, trainer_count=num_workers)
    cost, prediction, label = DSSM(dnn_dims=layer_dims, vocab_sizes=[len(
        load_dic(path)) for path in [source_dic_path, target_dic_path]],
        model_type=model_type, model_arch=model_arch,
        share_semantic_generator=share_semantic_generator, class_num=
        class_num, share_embed=share_embed)()
    parameters = paddle.parameters.create(cost)
    adam_optimizer = paddle.optimizer.Adam(learning_rate=0.0002,
        regularization=paddle.optimizer.L2Regularization(rate=0.001),
        model_average=paddle.optimizer.ModelAverage(average_window=0.5))
    trainer = paddle.trainer.SGD(cost=cost, extra_layers=paddle.evaluator.
        auc(input=prediction, label=label) if not model_type.is_rank() else
        None, parameters=parameters, update_equation=adam_optimizer)
    feeding = {}
    if model_type.is_classification() or model_type.is_regression():
        feeding = {'source_input': 0, 'target_input': 1, 'label_input': 2}
    else:
        feeding = {'source_input': 0, 'left_target_input': 1,
            'right_target_input': 2, 'label_input': 3}

    def _event_handler(event):
        """
        Define batch handler
        """
        if isinstance(event, paddle.event.EndIteration):
            if event.batch_id % args.num_batches_to_log == 0:
                logger.info('Pass %d, Batch %d, Cost %f, %s' % (event.
                    pass_id, event.batch_id, event.cost, event.metrics))
            if (event.batch_id > 0 and event.batch_id % args.
                num_batches_to_test == 0):
                if test_reader is not None:
                    if model_type.is_classification():
                        result = trainer.test(reader=test_reader, feeding=
                            feeding)
                        logger.info('Test at Pass %d, %s' % (event.pass_id,
                            result.metrics))
                    else:
                        result = None
            if (event.batch_id > 0 and event.batch_id % args.
                num_batches_to_save_model == 0):
                model_desc = '{type}_{arch}'.format(type=str(args.
                    model_type), arch=str(args.model_arch))
                with open('%sdssm_%s_pass_%05d.tar' % (args.
                    model_output_prefix, model_desc, event.pass_id), 'w') as f:
                    trainer.save_parameter_to_tar(f)
    trainer.train(reader=train_reader, event_handler=_event_handler,
        feeding=feeding, num_passes=num_passes)
    logger.info('Training has finished.')


if __name__ == '__main__':
    display_args(args)
    train(train_data_path=args.train_data_path, test_data_path=args.
        test_data_path, source_dic_path=args.source_dic_path,
        target_dic_path=args.target_dic_path, model_type=ModelType(args.
        model_type), model_arch=ModelArch(args.model_arch), batch_size=args
        .batch_size, num_passes=args.num_passes, share_semantic_generator=
        args.share_network_between_source_target, share_embed=args.
        share_embed, class_num=args.class_num, num_workers=args.num_workers,
        use_gpu=args.use_gpu)

import argparse
import distutils.util
import paddle.v2 as paddle
from network_conf import DSSM
import reader
from utils import TaskType, load_dic, logger, ModelType, ModelArch, display_args
parser = argparse.ArgumentParser(description='PaddlePaddle DSSM example')
parser.add_argument('-i', '--train_data_path', type=str, required=False,
    help='The path of training data.')
parser.add_argument('-t', '--test_data_path', type=str, required=False,
    help='The path of testing data.')
parser.add_argument('-s', '--source_dic_path', type=str, required=False,
    help="The path of the source's word dictionary.")
parser.add_argument('--target_dic_path', type=str, required=False, help=
    "The path of the target's word dictionary, if this parameter is not set, the `source_dic_path` will be used"
    )
parser.add_argument('-b', '--batch_size', type=int, default=32, help=
    'The size of mini-batch (default:32).')
parser.add_argument('-p', '--num_passes', type=int, default=10, help=
    'The number of passes to run(default:10).')
parser.add_argument('-y', '--model_type', type=int, required=True, default=
    ModelType.CLASSIFICATION_MODE, help=
    'model type, %d for classification, %d for pairwise rank, %d for regression (default: classification).'
     % (ModelType.CLASSIFICATION_MODE, ModelType.RANK_MODE, ModelType.
    REGRESSION_MODE))
parser.add_argument('-a', '--model_arch', type=int, required=True, default=
    ModelArch.CNN_MODE, help=
    'The model architecture, %d for CNN, %d for FC, %d for RNN.' % (
    ModelArch.CNN_MODE, ModelArch.FC_MODE, ModelArch.RNN_MODE))
parser.add_argument('--share_network_between_source_target', type=distutils
    .util.strtobool, default=False, help=
    'Whether to share network parameters between source and target.')
parser.add_argument('--share_embed', type=distutils.util.strtobool, default
    =False, help='Whether to share word embedding between source and target.')
parser.add_argument('--dnn_dims', type=str, default='256,128,64,32', help=
    "The dimentions of dnn layers, default is '256,128,64,32', which means create a 4-layer dnn. The dimention of each layer is '256, 128, 64 and 32'."
    )
parser.add_argument('--num_workers', type=int, default=1, help=
    'The number of worker threads, default 1.')
parser.add_argument('--use_gpu', type=distutils.util.strtobool, default=
    False, help='Whether to use GPU devices (default: False)')
parser.add_argument('-c', '--class_num', type=int, default=0, help=
    'The number of categories for classification task.')
parser.add_argument('--model_output_prefix', type=str, default='./', help=
    'The prefix of the path to store the trained models (default: ./).')
parser.add_argument('-g', '--num_batches_to_log', type=int, default=100,
    help=
    'The log period. Every num_batches_to_test batches, a training log will be printed. (default: 100)'
    )
parser.add_argument('-e', '--num_batches_to_test', type=int, default=200,
    help=
    'The test period. Every num_batches_to_save_model batches, the specified test sample will be test (default: 200).'
    )
parser.add_argument('-z', '--num_batches_to_save_model', type=int, default=
    400, help=
    'Every num_batches_to_save_model batches, a trained model will be saved (default: 400).'
    )
args = parser.parse_args()
args.model_type = ModelType(args.model_type)
args.model_arch = ModelArch(args.model_arch)
if args.model_type.is_classification():
    assert args.class_num > 1, 'The parameter class_num should be set in classification task.'
layer_dims = [int(i) for i in args.dnn_dims.split(',')]
args.target_dic_path = (args.source_dic_path if not args.target_dic_path else
    args.target_dic_path)


def train(train_data_path=None, test_data_path=None, source_dic_path=None,
    target_dic_path=None, model_type=ModelType.create_classification(),
    model_arch=ModelArch.create_cnn(), batch_size=32, num_passes=10,
    share_semantic_generator=False, share_embed=False, class_num=None,
    num_workers=1, use_gpu=False):
    """
    Train the DSSM.
    """
    default_train_path = './data/rank/train.txt'
    default_test_path = './data/rank/test.txt'
    default_dic_path = './data/vocab.txt'
    if not model_type.is_rank():
        default_train_path = './data/classification/train.txt'
        default_test_path = './data/classification/test.txt'
    use_default_data = not train_data_path
    if use_default_data:
        train_data_path = default_train_path
        test_data_path = default_test_path
        source_dic_path = default_dic_path
        target_dic_path = default_dic_path
    dataset = reader.Dataset(train_path=train_data_path, test_path=
        test_data_path, source_dic_path=source_dic_path, target_dic_path=
        target_dic_path, model_type=model_type)
    train_reader = paddle.batch(paddle.reader.shuffle(dataset.train,
        buf_size=1000), batch_size=batch_size)
    test_reader = paddle.batch(paddle.reader.shuffle(dataset.test, buf_size
        =1000), batch_size=batch_size)
    paddle.init(use_gpu=use_gpu, trainer_count=num_workers)
    cost, prediction, label = DSSM(dnn_dims=layer_dims, vocab_sizes=[len(
        load_dic(path)) for path in [source_dic_path, target_dic_path]],
        model_type=model_type, model_arch=model_arch,
        share_semantic_generator=share_semantic_generator, class_num=
        class_num, share_embed=share_embed)()
    parameters = paddle.parameters.create(cost)
    adam_optimizer = paddle.optimizer.Adam(learning_rate=0.0002,
        regularization=paddle.optimizer.L2Regularization(rate=0.001),
        model_average=paddle.optimizer.ModelAverage(average_window=0.5))
    trainer = paddle.trainer.SGD(cost=cost, extra_layers=paddle.evaluator.
        auc(input=prediction, label=label) if not model_type.is_rank() else
        None, parameters=parameters, update_equation=adam_optimizer)
    feeding = {}
    if model_type.is_classification() or model_type.is_regression():
        feeding = {'source_input': 0, 'target_input': 1, 'label_input': 2}
    else:
        feeding = {'source_input': 0, 'left_target_input': 1,
            'right_target_input': 2, 'label_input': 3}

    def _event_handler(event):
        """
        Define batch handler
        """
        if isinstance(event, paddle.event.EndIteration):
            if event.batch_id % args.num_batches_to_log == 0:
                logger.info('Pass %d, Batch %d, Cost %f, %s' % (event.
                    pass_id, event.batch_id, event.cost, event.metrics))
            if (event.batch_id > 0 and event.batch_id % args.
                num_batches_to_test == 0):
                if test_reader is not None:
                    if model_type.is_classification():
                        result = trainer.test(reader=test_reader, feeding=
                            feeding)
                        logger.info('Test at Pass %d, %s' % (event.pass_id,
                            result.metrics))
                    else:
                        result = None
            if (event.batch_id > 0 and event.batch_id % args.
                num_batches_to_save_model == 0):
                model_desc = '{type}_{arch}'.format(type=str(args.
                    model_type), arch=str(args.model_arch))
                with open('%sdssm_%s_pass_%05d.tar' % (args.
                    model_output_prefix, model_desc, event.pass_id), 'w') as f:
                    trainer.save_parameter_to_tar(f)
    trainer.train(reader=train_reader, event_handler=_event_handler,
        feeding=feeding, num_passes=num_passes)
    logger.info('Training has finished.')


if __name__ == '__main__':
    display_args(args)
    train(train_data_path=args.train_data_path, test_data_path=args.
        test_data_path, source_dic_path=args.source_dic_path,
        target_dic_path=args.target_dic_path, model_type=ModelType(args.
        model_type), model_arch=ModelArch(args.model_arch), batch_size=args
        .batch_size, num_passes=args.num_passes, share_semantic_generator=
        args.share_network_between_source_target, share_embed=args.
        share_embed, class_num=args.class_num, num_workers=args.num_workers,
        use_gpu=args.use_gpu)

import argparse
import distutils.util

import paddle.v2 as paddle
from network_conf import DSSM
import reader
from utils import TaskType, load_dic, logger, ModelType, ModelArch, display_args

parser = argparse.ArgumentParser(description="PaddlePaddle DSSM example")

parser.add_argument(
    "-i",
    "--train_data_path",
    type=str,
    required=False,
    help="The path of training data.")
parser.add_argument(
    "-t",
    "--test_data_path",
    type=str,
    required=False,
    help="The path of testing data.")
parser.add_argument(
    "-s",
    "--source_dic_path",
    type=str,
    required=False,
    help="The path of the source's word dictionary.")
parser.add_argument(
    "--target_dic_path",
    type=str,
    required=False,
    help=("The path of the target's word dictionary, "
          "if this parameter is not set, the `source_dic_path` will be used"))
parser.add_argument(
    "-b",
    "--batch_size",
    type=int,
    default=32,
    help="The size of mini-batch (default:32).")
parser.add_argument(
    "-p",
    "--num_passes",
    type=int,
    default=10,
    help="The number of passes to run(default:10).")
parser.add_argument(
    "-y",
    "--model_type",
    type=int,
    required=True,
    default=ModelType.CLASSIFICATION_MODE,
    help=("model type, %d for classification, %d for pairwise rank, "
          "%d for regression (default: classification).") %
    (ModelType.CLASSIFICATION_MODE, ModelType.RANK_MODE,
     ModelType.REGRESSION_MODE))
parser.add_argument(
    "-a",
    "--model_arch",
    type=int,
    required=True,
    default=ModelArch.CNN_MODE,
    help="The model architecture, %d for CNN, %d for FC, %d for RNN." %
    (ModelArch.CNN_MODE, ModelArch.FC_MODE, ModelArch.RNN_MODE))
parser.add_argument(
    "--share_network_between_source_target",
    type=distutils.util.strtobool,
    default=False,
    help="Whether to share network parameters between source and target.")
parser.add_argument(
    "--share_embed",
    type=distutils.util.strtobool,
    default=False,
    help="Whether to share word embedding between source and target.")
parser.add_argument(
    "--dnn_dims",
    type=str,
    default="256,128,64,32",
    help=("The dimentions of dnn layers, default is '256,128,64,32', "
          "which means create a 4-layer dnn. The dimention of each layer is "
          "'256, 128, 64 and 32'."))
parser.add_argument(
    "--num_workers",
    type=int,
    default=1,
    help="The number of worker threads, default 1.")
parser.add_argument(
    "--use_gpu",
    type=distutils.util.strtobool,
    default=False,
    help="Whether to use GPU devices (default: False)")
parser.add_argument(
    "-c",
    "--class_num",
    type=int,
    default=0,
    help="The number of categories for classification task.")
parser.add_argument(
    "--model_output_prefix",
    type=str,
    default="./",
    help="The prefix of the path to store the trained models (default: ./).")
parser.add_argument(
    "-g",
    "--num_batches_to_log",
    type=int,
    default=100,
    help=("The log period. Every num_batches_to_test batches, "
          "a training log will be printed. (default: 100)"))
parser.add_argument(
    "-e",
    "--num_batches_to_test",
    type=int,
    default=200,
    help=("The test period. Every num_batches_to_save_model batches, "
          "the specified test sample will be test (default: 200)."))
parser.add_argument(
    "-z",
    "--num_batches_to_save_model",
    type=int,
    default=400,
    help=("Every num_batches_to_save_model batches, "
          "a trained model will be saved (default: 400)."))

args = parser.parse_args()
args.model_type = ModelType(args.model_type)
args.model_arch = ModelArch(args.model_arch)
if args.model_type.is_classification():
    assert args.class_num > 1, ("The parameter class_num should be set in "
                                "classification task.")

layer_dims = [int(i) for i in args.dnn_dims.split(",")]
args.target_dic_path = args.source_dic_path if not \
        args.target_dic_path else args.target_dic_path


def train(train_data_path=None,
          test_data_path=None,
          source_dic_path=None,
          target_dic_path=None,
          model_type=ModelType.create_classification(),
          model_arch=ModelArch.create_cnn(),
          batch_size=32,
          num_passes=10,
          share_semantic_generator=False,
          share_embed=False,
          class_num=None,
          num_workers=1,
          use_gpu=False):
    """
    Train the DSSM.
    """
    default_train_path = "./data/rank/train.txt"
    default_test_path = "./data/rank/test.txt"
    default_dic_path = "./data/vocab.txt"
    if not model_type.is_rank():
        default_train_path = "./data/classification/train.txt"
        default_test_path = "./data/classification/test.txt"

    use_default_data = not train_data_path

    if use_default_data:
        train_data_path = default_train_path
        test_data_path = default_test_path
        source_dic_path = default_dic_path
        target_dic_path = default_dic_path

    dataset = reader.Dataset(
        train_path=train_data_path,
        test_path=test_data_path,
        source_dic_path=source_dic_path,
        target_dic_path=target_dic_path,
        model_type=model_type, )

    train_reader = paddle.batch(
        paddle.reader.shuffle(
            dataset.train, buf_size=1000),
        batch_size=batch_size)

    test_reader = paddle.batch(
        paddle.reader.shuffle(
            dataset.test, buf_size=1000),
        batch_size=batch_size)

    paddle.init(use_gpu=use_gpu, trainer_count=num_workers)

    cost, prediction, label = DSSM(
        dnn_dims=layer_dims,
        vocab_sizes=[
            len(load_dic(path)) for path in [source_dic_path, target_dic_path]
        ],
        model_type=model_type,
        model_arch=model_arch,
        share_semantic_generator=share_semantic_generator,
        class_num=class_num,
        share_embed=share_embed)()

    parameters = paddle.parameters.create(cost)

    adam_optimizer = paddle.optimizer.Adam(
        learning_rate=2e-4,
        regularization=paddle.optimizer.L2Regularization(rate=1e-3),
        model_average=paddle.optimizer.ModelAverage(average_window=0.5))

    trainer = paddle.trainer.SGD(
        cost=cost,
        extra_layers=paddle.evaluator.auc(input=prediction, label=label)
        if not model_type.is_rank() else None,
        parameters=parameters,
        update_equation=adam_optimizer)

    feeding = {}
    if model_type.is_classification() or model_type.is_regression():
        feeding = {"source_input": 0, "target_input": 1, "label_input": 2}
    else:
        feeding = {
            "source_input": 0,
            "left_target_input": 1,
            "right_target_input": 2,
            "label_input": 3
        }

    def _event_handler(event):
        """
        Define batch handler
        """
        if isinstance(event, paddle.event.EndIteration):
            # output train log
            if event.batch_id % args.num_batches_to_log == 0:
                logger.info("Pass %d, Batch %d, Cost %f, %s" % (
                    event.pass_id, event.batch_id, event.cost, event.metrics))

            # test model
            if event.batch_id > 0 and \
                    event.batch_id % args.num_batches_to_test == 0:
                if test_reader is not None:
                    if model_type.is_classification():
                        result = trainer.test(
                            reader=test_reader, feeding=feeding)
                        logger.info("Test at Pass %d, %s" % (event.pass_id,
                                                             result.metrics))
                    else:
                        result = None
            # save model
            if event.batch_id > 0 and \
                    event.batch_id % args.num_batches_to_save_model == 0:
                model_desc = "{type}_{arch}".format(
                    type=str(args.model_type), arch=str(args.model_arch))
                with open("%sdssm_%s_pass_%05d.tar" %
                          (args.model_output_prefix, model_desc,
                           event.pass_id), "w") as f:
                    trainer.save_parameter_to_tar(f)

    trainer.train(
        reader=train_reader,
        event_handler=_event_handler,
        feeding=feeding,
        num_passes=num_passes)

    logger.info("Training has finished.")


if __name__ == "__main__":
    display_args(args)
    train(
        train_data_path=args.train_data_path,
        test_data_path=args.test_data_path,
        source_dic_path=args.source_dic_path,
        target_dic_path=args.target_dic_path,
        model_type=ModelType(args.model_type),
        model_arch=ModelArch(args.model_arch),
        batch_size=args.batch_size,
        num_passes=args.num_passes,
        share_semantic_generator=args.share_network_between_source_target,
        share_embed=args.share_embed,
        class_num=args.class_num,
        num_workers=args.num_workers,
        use_gpu=args.use_gpu)
