Support reading compressed tfrecord files.

PiperOrigin-RevId: 436868704

official/common/dataset_fn.py

@@ -28,7 +28,8 @@
 # ==============================================================================
 """Utility library for picking an appropriate dataset function."""
 
-from typing import Any, Callable, Union, Type
+import functools
+from typing import Any, Callable, Type, Union
 
 import tensorflow as tf
 
@@ -38,5 +39,6 @@ PossibleDatasetType = Union[Type[tf.data.Dataset], Callable[[tf.Tensor], Any]]
 def pick_dataset_fn(file_type: str) -> PossibleDatasetType:
   if file_type == 'tfrecord':
     return tf.data.TFRecordDataset
-
+  if file_type == 'tfrecord_compressed':
+    return functools.partial(tf.data.TFRecordDataset, compression_type='GZIP')
   raise ValueError('Unrecognized file_type: {}'.format(file_type))