def get_predict_fn():
  """Setup and initialize function to get answers."""

  if FLAGS.model_type == "onehop":
    create_model_fn = model_fns.create_onehop_model
  elif FLAGS.model_type == "twohop":
    create_model_fn = model_fns.create_twohop_model
  elif FLAGS.model_type == "twohop-cascaded":
    create_model_fn = model_fns.create_twohopcascade_model
  elif FLAGS.model_type == "threehop":
    create_model_fn = functools.partial(
        model_fns.create_twohop_model, num_hops=3)
  elif FLAGS.model_type == "threehop-cascaded":
    create_model_fn = functools.partial(
        model_fns.create_twohopcascade_model, num_hops=3)
  elif FLAGS.model_type == "wikimovie":
    create_model_fn = model_fns.create_wikimovie_model
  elif FLAGS.model_type == "wikimovie-2hop":
    create_model_fn = functools.partial(
        model_fns.create_wikimovie_model, num_hops=2)
  elif FLAGS.model_type == "wikimovie-3hop":
    create_model_fn = functools.partial(
        model_fns.create_wikimovie_model, num_hops=3)
  elif FLAGS.model_type == "hotpotqa":
    create_model_fn = functools.partial(
        model_fns.create_hotpotqa_model, num_hops=FLAGS.num_hops)

  bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)

  # Load mention and entity files.
  mention2text = json.load(tf.gfile.Open(os.path.join(FLAGS.train_data_dir,
                                                      "mention2text.json")))
  tf.logging.info("Loading metadata about entities and mentions...")
  entity2id, entity2name = json.load(tf.gfile.Open(os.path.join(
      FLAGS.train_data_dir, "entities.json")))
  entityid2name = {str(i): entity2name[e] for e, i in entity2id.items()}

  # Load raw paragraphs.
  tf.logging.info("Loading raw passages...")
  all_passages = {}
  with tf.gfile.Open(FLAGS.raw_passages) as f:
    for ii, line in enumerate(f):
      if ii % 100000 == 0:
        tf.logging.info("Loaded %d", ii)
      item = json.loads(line.strip())
      all_passages[item["kb_id"].lower()] = (
          item["title"], item["context"], item["mentions"])

  qa_config = QAConfig(
      qry_layers_to_use=FLAGS.qry_layers_to_use,
      qry_aggregation_fn=FLAGS.qry_aggregation_fn,
      dropout=FLAGS.question_dropout,
      qry_num_layers=FLAGS.question_num_layers,
      projection_dim=FLAGS.projection_dim,
      load_only_bert=FLAGS.load_only_bert,
      num_entities=len(entity2id),
      max_entity_len=FLAGS.max_entity_len,
      ensure_answer_sparse=FLAGS.ensure_answer_sparse,
      ensure_answer_dense=FLAGS.ensure_answer_dense,
      train_with_sparse=FLAGS.train_with_sparse,
      predict_with_sparse=FLAGS.predict_with_sparse,
      fix_sparse_to_one=FLAGS.fix_sparse_to_one,
      supervision="entity",
      l2_normalize_db=FLAGS.l2_normalize_db,
      entity_score_aggregation_fn=FLAGS.entity_score_aggregation_fn,
      entity_score_threshold=FLAGS.entity_score_threshold,
      softmax_temperature=FLAGS.softmax_temperature,
      sparse_reduce_fn=FLAGS.sparse_reduce_fn,
      intermediate_loss=False,
      light=FLAGS.light,
      sparse_strategy=FLAGS.sparse_strategy,
      train_batch_size=1,
      predict_batch_size=1)

  mips_config = MipsConfig(
      ckpt_path=os.path.join(FLAGS.train_data_dir, "mention_feats"),
      ckpt_var_name="db_emb",
      num_mentions=len(mention2text),
      emb_size=FLAGS.projection_dim * 2,
      num_neighbors=FLAGS.num_mips_neighbors)

  tf.gfile.MakeDirs(FLAGS.output_dir)

  tokenizer = tokenization.FullTokenizer(
      vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case)

  tpu_cluster_resolver = None
  is_per_host = contrib_tpu.InputPipelineConfig.PER_HOST_V2
  run_config = contrib_tpu.RunConfig(
      cluster=tpu_cluster_resolver,
      master=FLAGS.master,
      model_dir=FLAGS.output_dir,
      save_checkpoints_steps=1000,
      tpu_config=contrib_tpu.TPUConfig(
          iterations_per_loop=1000,
          num_shards=FLAGS.num_tpu_cores,
          per_host_input_for_training=is_per_host),
      session_config=tf.ConfigProto(log_device_placement=False))

  if tf.gfile.Exists(os.path.join(FLAGS.init_checkpoint, "best_model.index")):
    bert_ckpt = os.path.join(FLAGS.init_checkpoint, "best_model")
  else:
    bert_ckpt = tf.train.latest_checkpoint(FLAGS.init_checkpoint)
  tf.logging.info("Initializing model with %s.data-00000-of-00001", bert_ckpt)
  model_fn = model_fn_builder(
      bert_config=bert_config,
      qa_config=qa_config,
      mips_config=mips_config,
      init_checkpoint=bert_ckpt,
      e2m_checkpoint=os.path.join(FLAGS.train_data_dir, "ent2ment.npz"),
      m2e_checkpoint=os.path.join(FLAGS.train_data_dir, "coref.npz"),
      entity_id_checkpoint=os.path.join(FLAGS.train_data_dir, "entity_ids"),
      entity_mask_checkpoint=os.path.join(FLAGS.train_data_dir, "entity_mask"),
      use_tpu=FLAGS.use_tpu,
      use_one_hot_embeddings=FLAGS.use_tpu,
      create_model_fn=create_model_fn,
      summary_obj=None)

  # If TPU is not available, this will fall back to normal Estimator on CPU
  # or GPU.
  estimator = contrib_tpu.TPUEstimator(
      use_tpu=FLAGS.use_tpu,
      model_fn=model_fn,
      config=run_config,
      train_batch_size=1,
      predict_batch_size=1)

  # Answer predictor.
  if tf.gfile.Exists(os.path.join(FLAGS.hotpot_init_checkpoint,
                                  "best_model.index")):
    bert_ckpt = os.path.join(FLAGS.hotpot_init_checkpoint, "best_model")
  else:
    bert_ckpt = tf.train.latest_checkpoint(FLAGS.hotpot_init_checkpoint)
  tf.logging.info("Initializing answer model with %s.data-00000-of-00001",
                  bert_ckpt)
  answer_predictor = BERTPredictor(tokenizer, bert_ckpt)
  predictor = Predictor(tokenizer, estimator, entity2id, entityid2name,
                        all_passages, answer_predictor, FLAGS.train_data_dir,
                        FLAGS.num_entities_linked)
  return predictor
Setup and initialize function to get answers.
language/labs/drkit/hotpotqa/demo.py