Generate examples as dicts.

def _build_pcollection(self, pipeline, folder, split):
    """Generate examples as dicts."""
    beam = tfds.core.lazy_imports.apache_beam

    split_type = self.builder_config.split_type
    filename = os.path.join(folder, "{}.tar.gz".format(split_type))


    def _extract_data(inputs):
      """Extracts files from the tar archives."""
      filename, split = inputs
      with tf.io.gfile.GFile(filename, "rb") as f:
        with tarfile.open(fileobj=f, mode="r") as tar:
          for tarinfo in tar:
            split_name = tarinfo.name.split("_")
            if len(split_name) > 2 and split_name[2] == split:
              buf = six.BytesIO()
              shutil.copyfileobj(tar.extractfile(tarinfo), buf)
              yield [tarinfo.name, buf.getvalue()]

    def _process_example(inputs):
      filename, data_string = inputs
      buf = six.BytesIO(data_string)
      buf.seek(0)
      data = np.load(buf)
      # Extract the images and convert to uint8. The reshape is required, see
      # https://github.com/deepmind/abstract-reasoning-matrices.
      all_images = np.uint8(data["image"].reshape(16, 160, 160, 1))
      return {
          "relation_structure_encoded": data["relation_structure_encoded"],
          "target": data["target"],
          "meta_target": data["meta_target"],
          "context": all_images[:8],
          "answers": all_images[8:],
          "filename": filename,
      }

    # Beam might fuse together the _extract_data and _process_example which
    # defeats the purpose of parallel processing. As a result, we reshard by
    # doing a GroupByKey on random keys, and then flattening again.
    def _add_random_keys(inputs):
      key = str(random.randrange(10**10))
      return key, inputs

    def _remove_keys(inputs):
      _, rows = inputs
      for row in rows:
        yield row

    return (pipeline
            | beam.Create([(filename, split)])
            | beam.FlatMap(_extract_data)
            | beam.Map(_add_random_keys)
            | beam.GroupByKey()
            | beam.FlatMap(_remove_keys)
            | beam.Map(_process_example))