text: # encoding:utf-8
  import sys
  import time
  import os
  os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
  import pickle
  import argparse
  import logging
  from dataset import Dataset
  from model import Model
  from utils import *
  import pprint
  from tensorboardX import SummaryWriter
  
  def parse_args():
      """
      Parses command line arguments.
      """
      parser = argparse.ArgumentParser('GraphCM')
      parser.add_argument('--train', action='store_true',
                          help='train the model')
      parser.add_argument('--valid', action='store_true',
                          help='perform click prediction task on valid set')
      parser.add_argument('--test', action='store_true',
                          help='perform click prediction task on test set')
      parser.add_argument('--rank', action='store_true',
                          help='perform relevance estimation task on human labeled test set')
      parser.add_argument('--num_iter', type=int, default=1,
                          help='the number of duplicated evaluation for valid/test/rank')
  
      train_settings = parser.add_argument_group('train settings')
      train_settings.add_argument('--optim', default='adadelta',
                                  help='optimizer type')
      train_settings.add_argument('--learning_rate', type=float, default=0.01,
                                  help='learning rate')
      train_settings.add_argument('--weight_decay', type=float, default=1e-5,
                                  help='weight decay')
      train_settings.add_argument('--momentum', type=float, default=0.99,
                                  help='momentum')
      train_settings.add_argument('--dropout_rate', type=float, default=0.5,
                                  help='dropout rate')
      train_settings.add_argument('--batch_size', type=int, default=64,
                                  help='train batch size')
      train_settings.add_argument('--num_steps', type=int, default=20000,
                                  help='number of training steps')
      train_settings.add_argument('--reg_relevance', type=float, default=1.0,
                                  help='regularization for relevance training')
  
      model_settings = parser.add_argument_group('model settings')
      model_settings.add_argument('--algo', default='GraphCM',
                                  help='choose the algorithm to use')
      model_settings.add_argument('--combine', default='mul',
                                  help='the combination type for examination and relevance')
      model_settings.add_argument('--embed_size', type=int, default=64,
                                  help='size of the query/doc embeddings')
      model_settings.add_argument('--pos_embed_size', type=int, default=4,
                                  help='size of the position embeddings')
      model_settings.add_argument('--click_embed_size', type=int, default=4,
                                  help='size of the click embeddings')
      model_settings.add_argument('--vtype_embed_size', type=int, default=8,
                                  help='size of the vtype embeddings')
      model_settings.add_argument('--hidden_size', type=int, default=128,
                                  help='size of RNN hidden units')
      model_settings.add_argument('--max_d_num', type=int, default=10,
                                  help='max number of docs in a session')
      model_settings.add_argument('--use_pretrain_embed', action='store_true',
                                  help='whether use pretrained embeddings')
      model_settings.add_argument('--use_gnn', action='store_true',
                                  help='whether use gnn layer')
      model_settings.add_argument('--gnn_neigh_sample', type=int, default=5,
                                  help='the number of neighbors to be sampled for GAT')
      model_settings.add_argument('--gnn_att_heads', type=int, default=2,
                                  help='the number of multi-head attention for GAT')
      model_settings.add_argument('--gnn_dropout', type=int, default=0,
                                  help='the dropout for the gat layer')
      model_settings.add_argument('--gnn_leaky_slope', type=float, default=0.2,
                                  help='leaky slope of leakyrelu for gat layer')
      model_settings.add_argument('--gnn_concat', type=bool, default=False,
                                  help='whether perform concatenation in gat layer')
      model_settings.add_argument('--inter_neigh_sample', type=int, default=0,
                                  help='the number of neighbor to be sampled for interaction')
      model_settings.add_argument('--inter_leaky_slope', type=float, default=0.2,
                                  help='leaky slope of leakyrelu for interaction')
  
      path_settings = parser.add_argument_group('path settings')
      path_settings.add_argument('--dataset', default='TianGong-ST',
                                  help='name of the dataset to be used')
      path_settings.add_argument('--model_dir', default='./outputs/models/',
                                  help='the dir to store models')
      path_settings.add_argument('--result_dir', default='./outputs/results/',
                                  help='the dir to output the results')
      path_settings.add_argument('--summary_dir', default='./outputs/summary/',
                                  help='the dir to write tensorboard summary')
      path_settings.add_argument('--log_dir', default='./outputs/log/',
                                  help='path of the log file. If not set, logs are printed to console')
  
      path_settings.add_argument('--eval_freq', type=int, default=100,
                                  help='the frequency of evaluating on the valid set when training')
      path_settings.add_argument('--check_point', type=int, default=100,
                                  help='the frequency of saving model')
      path_settings.add_argument('--patience', type=int, default=5,
                                  help='lr decay when more than the patience times of evaluation where loss/ppl do not decrease')
      path_settings.add_argument('--lr_decay', type=float, default=0.5,
                                  help='lr decay')
      path_settings.add_argument('--load_model', type=int, default=-1,
                                  help='load model global step')
      path_settings.add_argument('--data_parallel', type=bool, default=False,
                                  help='data_parallel')
      path_settings.add_argument('--gpu_num', type=int, default=1,
                                  help='gpu_num')
  
      return parser.parse_args()
  
  def train(args, dataset):
      """
      trains the model
      """
      logger = logging.getLogger("GraphCM")
      logger.info('Initialize the model...')
      model = Model(args, dataset.query_size, dataset.doc_size, dataset.vtype_size, dataset)
      logger.info('model.global_step: {}'.format(model.global_step))
      if args.load_model > -1:
          logger.info('Reloading the model...')
          model.load_model(model_dir=args.model_dir, model_prefix=args.algo, global_step=args.load_model)
      logger.info('Training the model...')
      model.train(dataset)
      logger.info('Done with model training!')
  
  def valid(args, dataset):
      """
      compute perplexity and log-likelihood for valid file
      """
      logger = logging.getLogger("GraphCM")
      logger.info('Initialize the model...')
      model = Model(args, dataset.query_size, dataset.doc_size, dataset.vtype_size, dataset)
      logger.info('model.global_step: {}'.format(model.global_step))
      assert args.load_model > -1, 'args.load_model is required to specify the model file to be loaded!'
      logger.info('Reloading the model...')
      model.load_model(model_dir=args.model_dir, model_prefix=args.algo, global_step=args.load_model)
      logger.info('Evaluating the model on valid set...')
      summary_writer = SummaryWriter(args.summary_dir)
      sum_click_loss, sum_perplexity = 0.0, 0.0
      for i in range(args.num_iter):
          valid_batches = dataset.gen_mini_batches('valid', dataset.validset_size, shuffle=False)
          valid_click_loss, valid_rel_loss, perplexity = model.evaluate(valid_batches, dataset)
          sum_click_loss += valid_click_loss
          sum_perplexity += perplexity
          summary_writer.add_scalar('final_valid/avg_click_loss', sum_click_loss / (i + 1), i)
          summary_writer.add_scalar('final_valid/avg_perplexity', sum_perplexity / (i + 1), i)
          summary_writer.add_scalar('final_valid/click_loss', valid_click_loss, i)
          summary_writer.add_scalar('final_valid/perplexity', perplexity, i)
  
  def test(args, dataset):
      """
      compute perplexity and log-likelihood for test file
      """
      logger = logging.getLogger("GraphCM")
      logger.info('Initialize the model...')
      model = Model(args, dataset.query_size, dataset.doc_size, dataset.vtype_size, dataset)
      logger.info('model.global_step: {}'.format(model.global_step))
      assert args.load_model > -1, 'args.load_model is required to specify the model file to be loaded!'
      logger.info('Reloading the model...')
      model.load_model(model_dir=args.model_dir, model_prefix=args.algo, global_step=args.load_model)
      logger.info('Predict on test files...')
      summary_writer = SummaryWriter(args.summary_dir)
      sum_click_loss, sum_perplexity = 0.0, 0.0
      for i in range(args.num_iter):
          test_batches = dataset.gen_mini_batches('test', dataset.testset_size, shuffle=False)
          test_click_loss, test_rel_loss, perplexity = model.evaluate(test_batches, dataset)
          sum_click_loss += test_click_loss
          sum_perplexity += perplexity
          summary_writer.add_scalar('final_test/avg_click_loss', sum_click_loss / (i + 1), i)
          summary_writer.add_scalar('final_test/avg_perplexity', sum_perplexity / (i + 1), i)
          summary_writer.add_scalar('final_test/click_loss', test_click_loss, i)
          summary_writer.add_scalar('final_test/perplexity', perplexity, i)
  
  def rank(args, dataset):
      """
      ranking performance on test files
      """
      logger = logging.getLogger("GraphCM")
      logger.info('Initialize the model...')
      model = Model(args, dataset.query_size, dataset.doc_size, dataset.vtype_size, dataset)
      logger.info('model.global_step: {}'.format(model.global_step))
      assert args.load_model > -1, 'args.load_model is required to specify the model file to be loaded!'
      logger.info('Reloading the model...')
      model.load_model(model_dir=args.model_dir, model_prefix=args.algo, global_step=args.load_model)
      logger.info('Start computing NDCG@k for ranking performance')
      trunc_levels = [1, 3, 5, 10]
      sum_ndcgs = {trunc_level: 0.0 for trunc_level in trunc_levels}
      summary_writer = SummaryWriter(args.summary_dir)
      for i in range(args.num_iter):
          label_batches = dataset.gen_mini_batches('label', dataset.labelset_size, shuffle=False)
          ndcgs = model.ranking(label_batches, dataset)
          for trunc_level in trunc_levels:
              sum_ndcgs[trunc_level] += ndcgs[trunc_level]
              summary_writer.add_scalar('final_NDCG/avg_{}'.format(trunc_level), sum_ndcgs[trunc_level] / (i + 1), i)
              summary_writer.add_scalar('final_NDCG/{}'.format(trunc_level), ndcgs[trunc_level], i)
          
  def run():
      """
      Prepares and runs the whole system.
      """
      # get arguments
      args = parse_args()
      assert args.batch_size % args.gpu_num == 0
      assert args.hidden_size % 2 == 0
  
      # create a logger
      logger = logging.getLogger("GraphCM")
      logger.setLevel(logging.INFO)
      formatter = logging.Formatter('%(asctime)s - %(message)s')
      check_path(args.model_dir)
      check_path(args.result_dir)
      check_path(args.summary_dir)
      if args.log_dir:
          check_path(args.log_dir)
          file_handler = logging.FileHandler(args.log_dir + time.strftime('%Y-%m-%d-%H:%M:%S',time.localtime(time.time())) + '.txt')
          file_handler.setLevel(logging.INFO)
          file_handler.setFormatter(formatter)
          logger.addHandler(file_handler)
      else:
          console_handler = logging.StreamHandler()
          console_handler.setLevel(logging.INFO)
          console_handler.setFormatter(formatter)
          logger.addHandler(console_handler)
  
      logger.info('Running with args : {}'.format(args))
  
      logger.info('Checking the directories...')
      for dir_path in [args.model_dir, args.result_dir, args.summary_dir]:
          if not os.path.exists(dir_path):
              os.makedirs(dir_path)
  
      logger.info('Loading train/valid/test/label/graph data...')
      dataset = Dataset(args)
      
      if args.train:
          train(args, dataset)
      if args.valid:
          valid(args, dataset)
      if args.test:
          test(args, dataset)
      if args.rank:
          rank(args, dataset)
      logger.info('run done.')
  
  if __name__ == '__main__':
      run()
  