def raw_data_to_tfrecord(item):
  """Process data into TFRecord compatible format to prepare for writing."""
  image = tf.io.serialize_tensor(tf.convert_to_tensor(item['image'])).numpy()
  context_features = {
      'index': get_int64_feature([item['index']]),
      'token': get_int64_feature(item['token']),
      'txt_mask': get_int64_feature(item['txt_mask']),
      'target_token': get_int64_feature(item['target_token']),
      'target_txt_mask': get_int64_feature(item['target_txt_mask']),
      'image': get_bytes_feature([image]),
  }
  ex = tf.train.Example(features=tf.train.Features(feature=context_features))

  return ex
Process data into TFRecord compatible format to prepare for writing.
language/gscan/xattn_model/dataset/preprocess.py