def iter_rc_examples(context_file, qa_file, vocab):
  """Join examples by id and yield."""
  # image_id --> qid --> list(beams)
  contexts = collections.defaultdict(lambda: collections.defaultdict(list))
  with tf.io.gfile.GFile(context_file, "r") as f:
    for line in f:
      line = json.loads(line)
      for sample in line["token_ids"]:
        sample = [vocab.i2t(idx) for idx in sample]
        try:
          eos_idx = sample.index(vocab.SEP)
          sample = sample[:eos_idx]
        except ValueError:
          pass
        sample = " ".join(sample)
        question_id = line.get("question_id", -1)
        if question_id < 0:
          question_id = "default"
        contexts[line["image_id"]][question_id].append(sample)

  with tf.io.gfile.GFile(qa_file, "r") as f:
    rc_dataset = json.load(f)

  # JSON keys are automatically stored as strings, so we map between
  # ints and strings as necessary. We ignore answers.
  for image_id, qas in rc_dataset.items():
    image_id = int(image_id)
    for question_id, qa in qas.items():
      question_id = int(question_id)
      captions = contexts[image_id]
      if question_id in captions:
        conditional_contexts = captions[question_id]
      else:
        conditional_contexts = captions["default"]
      for i, context in enumerate(conditional_contexts):
        features = dict(
            beam_id=io_utils.int64_feature(i),
            image_id=io_utils.int64_feature(image_id),
            question_id=io_utils.int64_feature(question_id),
            context=io_utils.string_feature(context),
            question=io_utils.string_feature(qa["question"]))
        yield tf.train.Example(features=tf.train.Features(feature=features))
Join examples by id and yield.
language/capwap/datasets/rc_dataset.py