text: from config.Config import Config
  from models.QuatRE import QuatRE
  import json
  import os
  
  
  from argparse import ArgumentParser
  parser = ArgumentParser("QuatRE")
  parser.add_argument("--dataset", default="WN18RR", help="Name of the dataset.")
  parser.add_argument("--learning_rate", default=0.1, type=float, help="Learning rate")
  parser.add_argument("--nbatches", default=100, type=int, help="Number of batches")
  parser.add_argument("--num_epochs", default=8000, type=int, help="Number of training epochs")
  parser.add_argument("--model_name", default='', help="")
  parser.add_argument('--neg_num', default=5, type=int, help='')
  parser.add_argument('--hidden_size', type=int, default=256, help='')
  parser.add_argument('--save_steps', type=int, default=10000, help='')
  parser.add_argument('--valid_steps', type=int, default=400, help='')
  parser.add_argument("--lmbda", default=0.5, type=float, help="")
  parser.add_argument("--lmbda2", default=0.01, type=float, help="")
  parser.add_argument("--mode", choices=["train", "predict"], default="train", type=str)
  parser.add_argument("--checkpoint_path", default=None, type=str)
  parser.add_argument("--test_file", default="", type=str)
  parser.add_argument("--optim", default='adagrad', help="")
  args = parser.parse_args()
  
  if args.model_name is None or len(args.model_name.strip()) == 0:
      args.model_name = "{}_lda-{}_nneg-{}_hs-{}_lr-{}_nepochs-{}".format(args.dataset,
                                                                          args.lmbda,
                                                                          args.neg_num,
                                                                          args.hidden_size,
                                                                          args.learning_rate,
                                                                          args.num_epochs)
  print(args)
  out_dir = os.path.abspath(os.path.join("../runs_QTransE/"))
  print("Writing to {}\n".format(out_dir))
  # Checkpoint directory
  checkpoint_dir = os.path.abspath(os.path.join(out_dir, "checkpoints"))
  result_dir = os.path.abspath(os.path.join(checkpoint_dir, args.model_name))
  if not os.path.exists(checkpoint_dir):
      os.makedirs(checkpoint_dir)
  if not os.path.exists(result_dir):
      os.makedirs(result_dir)
  
  con = Config()
  in_path = "./benchmarks/" + args.dataset + "/"
  con.set_in_path(in_path)
  test_file_path = ""
  if args.test_file != "":
      test_file_path = in_path + args.test_file
  con.set_test_file_path(test_file_path)
  con.set_work_threads(8)
  con.set_train_times(args.num_epochs)
  con.set_nbatches(args.nbatches)
  con.set_alpha(args.learning_rate)
  con.set_bern(1)
  con.set_dimension(args.hidden_size)
  con.set_lmbda(args.lmbda)
  con.set_lmbda_two(0.01)
  con.set_margin(1.0)
  con.set_ent_neg_rate(args.neg_num)
  con.set_opt_method(args.optim)
  con.set_save_steps(args.save_steps)
  con.set_valid_steps(args.valid_steps)
  con.set_early_stopping_patience(10)
  con.set_checkpoint_dir(checkpoint_dir)
  con.set_result_dir(result_dir)
  # knowledge graph completion ~ link prediction
  con.set_test_link(True)
  con.init()
  
  # training mode
  if args.mode == "train":
      con.set_train_model(QuatRE)
      con.training_model()
  # prediction mode
  else:
      con.set_test_model(QuatRE, args.checkpoint_path)
      con.test()
  