def model_webqsp(self, features, params,
                   top_k):
    """Compute logits for more evaluation.

    Args:
      features: features
      params: hyper-parameters
      top_k: top k to retrieve at each hop

    Returns:
      predictions and losses
    """

    tensors = dict()
    tensors['intermediate_logits'] = list()
    tensors['intermediate_labels'] = list()
    tensors['intermediate_objs'] = list()
    tensors['intermediate_answerable'] = list()

    (question_token_ids, segment_ids, question_mask, question_entity_id,
     question_entity_sketch, constraint_entity_id, constraint_entity_sketch,
     answer_ids) = features

    # Compute question embeddings and question entity embeddings.
    question_embs = self.compute_bert_cls_embs(question_token_ids, segment_ids,
                                               question_mask)
    # batch_size, bert_hidden_size
    question_entity_embs = tf.nn.embedding_lookup(self.entity_embeddings_mat,
                                                  question_entity_id)
    # batch_size, hidden_size

    # Initialize query embeddings before iteration.
    query_entity_embs = question_entity_embs
    # batch_size, hidden_size
    query_entity_sketch = question_entity_sketch
    # batch_size, depth, width

    # Questions in WebQuestionsSP are either one hop questions or
    # two hop questions where the intermediate entities is CVT
    # entities. In this experiment, we set the max hop to 2.
    for hid in range(2):
      # Compute query embeddings of relation by projecting the original
      # question embeddings into another space.
      with tf.variable_scope('question_emb_ffn_%d' % hid):
        query_relation_embs = tf.keras.layers.Dense(
            units=params['relation_emb_size'])(
                question_embs)
        # batch_size, hidden_size

      # We concatenate the subject, relation, and object embeddings to form
      # a query. Note that we set relation embeddings as 0, because
      # we observe that this could makes the initial retrieval process
      # more stable at training time. This retrieval will only return a fact id.
      # We will recompute the similarity score with non-zero query relation
      # embeddings which will eventually be used to compute logits.
      # Another possibility is to set a small coeffient, \alpha, with small
      # values in the beginning, and slightly increase it as training going.
      query_embs = tf.concat([
          query_entity_embs,
          tf.zeros_like(query_relation_embs),
          tf.zeros_like(query_entity_embs)
      ],
                             axis=1)
      # batch_size, hiddent_size * 3

      # Retrieve the topk facts and gather their subjects and objects.
      _, topk_fact_ids = util.topk_search_fn(
          query_embs, self.kb_embeddings_mat, top_k)
      # batch_size, topk
      topk_subj_ids = tf.gather(self.all_fact_subjids, topk_fact_ids)
      # batch_size, topk
      topk_obj_ids = tf.gather(self.all_fact_objids, topk_fact_ids)
      # batch_size, topk

      # We check if the retrieved triple is eligible. To do so, we check
      # if the subject of the triple passes the cm-sketch with non-zero
      # probability. The probability of entities in the sketch is computed
      # from the previous iterations (or directly initialized) as the
      # softmax of logits. To recover the logits from softmax, we take the
      # log values and further mask out those that are not eligible.
      is_topk_fact_eligible, topk_fact_prev_probs = \
          module.check_topk_fact_eligible(
              topk_subj_ids, query_entity_sketch,
              self.all_entity_sketches, params)
      # batch_size, topk
      topk_fact_prev_logits = tf.math.log(topk_fact_prev_probs + 1e-6)
      # batch_size, topk
      topk_fact_prev_logits += (1.0 - is_topk_fact_eligible) * VERY_NEG
      # batch_size, topk
      with tf.variable_scope('topk_fact_eligibility_%d' % hid):
        # We learn another bias term here. This helps to adjust how
        # significant signals from previous iterations contribute to
        # later predictions.
        topk_fact_prev_logit_bias = tf.get_variable(
            name='topk_fact_logit_bias',
            dtype=tf.float32,
            shape=[1, 1],
            initializer=tf.random_normal_initializer(),
            trainable=True)
        topk_fact_prev_logits += topk_fact_prev_logit_bias

      # Now, we take the full fact embedding and compute the similarity
      # score between query and facts (with query relation embedding). We
      # further added the logits from previous iterations (after passing
      # the sketch) if we would like to use cm sketch.
      query_embs = tf.concat(
          [tf.zeros_like(query_entity_embs), query_relation_embs], axis=1)
      # batch_size, hidden_size * 2
      topk_fact_logit = self.compute_topk_fact_logits(topk_fact_ids, query_embs,
                                                      params)
      # batch_size, topk
      if params['use_cm_sketch']:
        topk_fact_logit += topk_fact_prev_logits
      # batch_size, topk

      # We filter the logits of CVT and regular entities. In the WebQuestionsSP
      # dataset, questions are either 1 hop questions or 2 hop questions
      # that travel through a CVT node. Thus, logits of regular entities
      # are considered to be answer candidates. Logits of CVT entities
      # will be passed to the next hop.
      # To distinguish CVT nodes, we could add another bit to entity embeddings
      # which we refer to as "hard" type. For simplicity, we use another
      # vector to store such information (just not appended to the end of
      # the embedding table).
      is_topk_obj_cvt = tf.gather(self.all_entity_is_cvt, topk_obj_ids)
      # batch_size, topk
      topk_ent_fact_logit = topk_fact_logit + is_topk_obj_cvt * VERY_NEG
      # batch_size, topk
      topk_cvt_fact_logit = topk_fact_logit + (1.0 - is_topk_obj_cvt) * VERY_NEG
      # batch_size, topk
      tensors['intermediate_logits'].append(topk_ent_fact_logit)
      tensors['intermediate_objs'].append(topk_obj_ids)

      # Then we compute the new query embedding and cm-sketch for the next
      # iteration.
      query_entity_embs, query_entity_sketch = \
          self.compute_query_embs_and_sketch(
              topk_obj_ids, topk_cvt_fact_logit, params)

      # Finally, we check if any of the retrieved facts actually contain
      # the correct answer. We treated them as labels and store them
      # for future use. We also compute how many percent of questions are
      # retrieved but not correctly answered. This is the upper bound of
      # the performance of our model.
      _, topk_fact_labels = util.compute_x_in_set(topk_obj_ids, answer_ids)
      topk_fact_labels = tf.cast(topk_fact_labels, dtype=tf.float32)
      # batch_size, topk
      topk_fact_labels *= is_topk_fact_eligible
      # batch_size, topk
      _, topk_objs_in_answers = util.compute_x_in_set(topk_obj_ids, answer_ids)
      # batch_size, topk
      topk_objs_in_answers = tf.logical_and(topk_objs_in_answers,
                                            is_topk_fact_eligible > 0.0)
      # batch_size, topk
      topk_objs_in_answers = tf.reduce_any(topk_objs_in_answers, axis=1)
      # batch_size
      tensors['intermediate_labels'].append(topk_fact_labels)
      tensors['intermediate_answerable'].append(topk_objs_in_answers)

    # After a few iterations, we concatenate the logits, labels and predictions
    # to unify the retrieval and prediction results of all iterations. The
    # concatenated results will be used for final prediction.
    concat_topk_obj_ids = tf.concat(tensors['intermediate_objs'], axis=1)
    # batch_size, topk * 2
    concat_topk_fact_logit = tf.concat(tensors['intermediate_logits'], axis=1)
    # batch_size, topk * 2
    concat_topk_fact_labels = tf.concat(tensors['intermediate_labels'], axis=1)
    # batch_size, topk * 2

    # We observe that there are ties between top predicted facts. They share
    # similar prediction scores but only a few of them satisfy the constraint.
    # Thus, we implement a constraint module to discriminate those entities.
    # We first compute an average of entity embeddings of tied facts.
    # Constraint entity embeddings are directly loaded from embedding table.
    concat_topk_fact_best_logit = tf.reduce_max(
        concat_topk_fact_logit, axis=1, keepdims=True)
    # batch_size, 1
    filtered_concat_topk_fact_best_logit = tf.cast(
        tf.equal(concat_topk_fact_logit, concat_topk_fact_best_logit),
        dtype=tf.float32)
    # batch_size, topk * 2
    concat_topk_subj_query_embs, concat_topk_subj_sketches = \
        self.compute_query_embs_and_sketch(
            concat_topk_obj_ids, filtered_concat_topk_fact_best_logit, params)
    # batch_size, topk * 2, hidden_size
    constraint_entity_embs = util.embedding_lookup_with_padding(
        self.entity_embeddings_mat, constraint_entity_id)
    # batch_size, hidden_size
    with tf.variable_scope('question_emb_ffn_constraint'):
      # Project question embeddings to get query relation embeddings.
      constraint_relation_embs = tf.keras.layers.Dense(
          units=params['relation_emb_size'])(
              question_embs)
      # batch_size, hidden_size
    constraint_query_embs = tf.concat([
        concat_topk_subj_query_embs,
        tf.zeros_like(constraint_relation_embs), constraint_entity_embs
    ],
                                      axis=1)
    # batch_size, hiddent_size * 3
    constraint_topk_fact_logits, constraint_topk_fact_ids = \
        util.topk_search_fn(constraint_query_embs, self.kb_embeddings_mat,
                            top_k)
    # batch_size, topk

    # Similar as previous retrieval steps, we check if retrieved facts for
    # constraints are eligible. We mask out logits of ineligible facts.
    constraint_topk_subj_ids = tf.gather(self.all_fact_subjids,
                                         constraint_topk_fact_ids)
    # batch_size, topk
    is_constraint_topk_subj_eligible, _ = module.check_topk_fact_eligible(
        constraint_topk_subj_ids, concat_topk_subj_sketches,
        self.all_entity_sketches, params)
    # batch_size, topk
    constraint_topk_obj_ids = tf.gather(self.all_fact_objids,
                                        constraint_topk_fact_ids)
    # batch_size, topk
    is_constraint_topk_obj_eligible, _ = module.check_topk_fact_eligible(
        constraint_topk_obj_ids, constraint_entity_sketch,
        self.all_entity_sketches, params)
    # batch_size, topk
    is_constraint_topk_eligible = tf.minimum(is_constraint_topk_subj_eligible,
                                             is_constraint_topk_obj_eligible)
    # batch_size, topk
    constraint_topk_fact_logits += (1.0 -
                                    is_constraint_topk_eligible) * VERY_NEG
    # batch_size, topk
    constraint_topk_fact_logits = tf.nn.relu(constraint_topk_fact_logits)
    # batch_size, topk

    # We need to add the logits from constraints to the logits at the
    # end of reasoning. However, the order of subject and object entities
    # does not match. We first find the mapping, map the logits, and add
    # it to the original logits.
    constraint_topk_fact_logits_mapped_to_concat = self.map_constraint_logits(
        concat_topk_obj_ids, constraint_topk_subj_ids,
        constraint_topk_fact_logits)
    concat_topk_fact_logit_with_constraint = concat_topk_fact_logit + \
        constraint_topk_fact_logits_mapped_to_concat
    # batch_size, topk * 2

    # Finally compute loss, cache a few tensors.
    answerable_at_topk = tf.metrics.mean(
        tf.cast(
            tf.logical_or(tensors['intermediate_answerable'][0],
                          tensors['intermediate_answerable'][1]),
            dtype=tf.float32))
    loss = self.compute_tf_loss(concat_topk_fact_logit, concat_topk_fact_labels)
    tensors['logits'] = concat_topk_fact_logit_with_constraint
    tensors['labels'] = concat_topk_fact_labels
    tensors['answerable_at_topk'] = answerable_at_topk

    return loss, tensors
Compute logits for more evaluation.

    Args:
      features: features
      params: hyper-parameters
      top_k: top k to retrieve at each hop

    Returns:
      predictions and losses
    
language/emql/model.py