other.py

def transformed_name(key: Text) -> Text:
    return key + '_xf'
    
def gzip_reader_fn(filenames):
    return tf.data.TFRecordDataset(filenames, compression_type='GZIP')
  
def input_fn(file_pattern, tf_transform_output,
         feature_spec,
         # feature_keys, input_feature_keys, predict_feature_keys, or anything you like
         batch_size=256):

  apply_tf_transform_map_fn = get_apply_tft_map_fn()

  dataset = tf.data.experimental.make_batched_features_dataset(
      file_pattern=file_pattern,
      features=feature_spec,
      reader=gzip_reader_fn,
      shuffle=True,
      sloppy_ordering=True,
      batch_size=batch_size) \
      ... \
      .map(apply_tf_transform_map_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE) \
      .prefetch(tf.data.experimental.AUTOTUNE) \

  return dataset