//  [python]
def compute_pred_dict(candidates_dict, dev_features, raw_results,
                      candidate_beam):
  """Computes official answer key from raw logits.

  This function joins three pieces needed for eval script for each example,
  based on the unique_id:

  1. Examples, which come from the original JSON definition of the dataset;
    each has a unique `example_id`.
  2. Features, which are the windowed sequences of wordpieces given to the
    neural network graph (and may be smaller than a single passage);
    each has a `unique_id`.
  3. Raw results, which are the predictions coming from the execution of the
    neural network graph; each has a `unique_id`.

  Because of the way `unique_ids` are assigned by `CreateTFExampleFn`,
  all `unique_ids` associated with an `example_id` should numerically be sorted
  after that `example_id`. The `for` loop over `datum`s below takes advantage
  of this in order to merge these three things together.

  Finally, with all of these things available together, this function delegates
  to `compute_predictions(...)` to post-process the prediction for each example
  and turn it into the JSON prediction format expected by the eval script.

  Args:
    candidates_dict: A dictionary containing the annotations from jsonl file.
    dev_features: Features loaded from tf_record file.
    raw_results: Output from running tensorflow graph.
    candidate_beam: see FLAGS.candidate_beam.

  Returns:
    A dictionary containing predictions.
  """
  logging.info("Post-processing predictions started.")

  raw_results_by_id = [(int(res["unique_id"] + 1), res) for res in raw_results]
  # Cast example id to int32 for each example, similarly to the raw results.
  sess = tf.Session()
  all_candidates = candidates_dict.items()
  example_ids = tf.cast(np.array([int(k) for k, _ in all_candidates
                                     ]), tf.int32).eval(session=sess)
  examples_by_id = list(zip(example_ids, all_candidates))
  if not examples_by_id:
    raise ValueError("No examples candidates found.")
  feature_ids = []
  features = []
  for f in dev_features:
    feature_ids.append(f.features.feature["unique_ids"].int64_list.value[0] + 1)
    features.append(f.features.feature)
  feature_ids = tf.cast(np.array(feature_ids), tf.int32).eval(session=sess)
  features_by_id = list(zip(feature_ids, features))

  # Join examples with features and raw results.
  eval_examples = []
  merged = sorted(
      examples_by_id + raw_results_by_id + features_by_id,
      key=lambda pair: pair[0])

  # Error counters
  num_failed_matches = 0

  ex_count = 0
  feature_count = 0
  result_count = 0
  # `feature_unique_id` is an example ID or an example ID with something
  # appended on the end of it such that it sorts after the appropriate
  # example ID (see `convert_examples_to_features`).
  for feature_unique_id, datum in merged:
    # if from `examples_by_id`
    if isinstance(datum, tuple):
      ex_count += 1
      eval_examples.append(
          EvalExample(example_id=datum[0], candidates=datum[1]))
    # if from `features_by_id`
    elif "wp_start_offset" in datum:
      feature_count += 1
      # Join with the example that we just appended above, by
      # adding to the `EvalExample`'s `features` dict.
      if not eval_examples:
        logging.warning("Expected to already have example for this example id. "
                        "Dataset / predictions mismatch?")
        num_failed_matches += 1
        continue
      eval_examples[-1].features[feature_unique_id] = datum
    # if from `raw_results_by_id`
    else:
      result_count += 1
      # Join with the example that we just appended above, by
      # adding to the `EvalExample`'s `results` dict.
      if not eval_examples:
        logging.warning("Expected to already have example for this example id. "
                        "Dataset / predictions mismatch?")
        num_failed_matches += 1
        continue
      eval_examples[-1].results[feature_unique_id] = datum

  logging.info("  Num candidate examples found: %d", ex_count)
  logging.info("  Num candidate features found: %d", feature_count)
  logging.info("  Num results found: %d", result_count)
  logging.info("  len(merged): %d", len(merged))
  if num_failed_matches > 0:
    logging.warning("  Num failed matches: %d", num_failed_matches)

  # Construct prediction objects.
  tydi_pred_dict = {}
  for i, eval_example in enumerate(eval_examples):
    if i % 10000 == 0:
      logging.info("  Computing predictions for input example: %d/%d", i,
                   len(eval_examples))

    summary = compute_predictions(eval_example, candidate_beam)
    if summary is not None:
      tydi_pred_dict[eval_example.example_id] = summary.predicted_label

  return tydi_pred_dict