File size: 8,784 Bytes
d9115cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 | import os
import random
from abc import ABC, abstractmethod
from typing import Callable, Dict, List
from src.utils.decorators import retry
from src.utils.file_utils import open_pyarrow_file
# We suppress the tensorflow warnings. Needs to happend before the tf import
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
import tensorflow as tf
tf.config.set_visible_devices([], "GPU") # Disable all for tensorflow
# if GPU version of TF installed,
# it will automatically occupy the full GPU memory
class RawDataIterator(ABC):
"""the abstract class for raw data iterator (e.g., parquet, avro, etc.)
Parameters
----------
list_of_file_paths : List[str]
the list of file paths to read from
"""
def __init__(
self,
**kwargs,
):
self.list_of_file_paths = None
self.should_shuffle_rows = None
def update_list_of_file_paths(self, list_of_file_paths: List[str]):
self.list_of_file_paths = list_of_file_paths
@abstractmethod
def get_file_suffix(self) -> str:
raise NotImplementedError("Must be implemented in child classes")
@abstractmethod
def iterrows(self):
raise NotImplementedError("Must be implemented in child classes")
@abstractmethod
def shuffle(self, seed=42):
raise NotImplementedError("Must be implemented in child classes")
@abstractmethod
def iter_batches(self, batch_size: int):
raise NotImplementedError("Must be implemented in child classes")
@retry()
def _get_next_example(self, dataset_iterator):
try:
return next(dataset_iterator)
except StopIteration:
return None
class ParquetDataIterator(RawDataIterator):
"""Data iterator class for parquet files
Parameters
----------
list_of_file_paths : List[str]
the list of file paths to read from
"""
def __init__(self, buffer_size=1000, features_to_consider=[], *args, **kwargs):
super().__init__(*args, **kwargs)
self.buffer_size = buffer_size
self.features_to_consider = features_to_consider
def iterrows(self):
assert self.list_of_file_paths is not None, "list_of_file_paths is not set"
for batch in self.iter_batches(batch_size=self.buffer_size):
for row in batch.to_pylist():
yield row
def iter_batches(self, batch_size: int) -> Dict[str, tf.Tensor]: # type: ignore
assert self.list_of_file_paths is not None, "list_of_file_paths is not set"
for file_path in self.list_of_file_paths:
with open_pyarrow_file(file_path) as f:
parquet_file = pq.ParquetFile(f)
for batch in parquet_file.iter_batches(
columns=self.features_to_consider
if len(self.features_to_consider)
else None,
batch_size=batch_size,
):
yield batch
def shuffle(self, seed=42) -> RawDataIterator:
random.seed(seed)
random.shuffle(self.list_of_file_paths) # type: ignore
return self
def get_file_suffix(self) -> str:
return "parquet"
class TFRecordIterator(RawDataIterator):
"""Data iterator class for tfrecord files
Parameters
----------
use_ragged_tensor: bool
Whether to use ragged tensors.
batch_tf_processing_functions: List[Callable]
A list of tensorflow functions to apply to the batches.
should_drop_last_batch: bool
Whether to drop the last batch if it is not a multiple of the batch size.
"""
def __init__(
self,
use_ragged_tensor: bool = False,
batch_tf_processing_functions: List[Callable] = [],
should_drop_last_batch: bool = True,
**kwargs,
):
super().__init__(**kwargs)
self.feature_description = None
self.use_ragged_tensor = use_ragged_tensor
self.batch_tf_processing_functions = batch_tf_processing_functions
self.should_drop_last_batch = should_drop_last_batch
def initialize_feature_description(self, raw_dataset: tf.data.TFRecordDataset):
"""
If the feature description is not set, infer the feature description from the first record in the dataset.
"""
if self.feature_description is None:
sample_record: tf.Tensor = next(iter(raw_dataset)) # type: ignore
self.feature_description = self.infer_feature_type(
self.parse_tfrecord(sample_record).features.feature # type: ignore
)
def iterrows(self):
assert self.list_of_file_paths is not None, "list_of_file_paths is not set"
raw_dataset = tf.data.TFRecordDataset(
[self.list_of_file_paths], compression_type="GZIP"
)
if self.should_shuffle_rows:
# the buffer here is the number of records to shuffle
# the larger the buffer, the more memory it will use
# too large might cause OOM
raw_dataset = raw_dataset.shuffle(buffer_size=128)
self.initialize_feature_description(raw_dataset)
# We create an iterator and manually iterate to allow for retrying the
# "next" operation in case of a failure.
dataset_iterator = iter(raw_dataset)
curr_example = self._get_next_example(dataset_iterator)
while curr_example:
example = tf.io.parse_single_example(curr_example, self.feature_description)
yield example
curr_example = self._get_next_example(dataset_iterator)
def iter_batches(self, batch_size: int) -> Dict[str, tf.Tensor]: # type: ignore
assert self.list_of_file_paths is not None, "list_of_file_paths is not set"
raw_dataset = tf.data.TFRecordDataset(
self.list_of_file_paths, compression_type="GZIP"
)
if self.should_shuffle_rows:
# the buffer here is the number of records to shuffle
# the larger the buffer, the more memory it will use
# too large might cause OOM
raw_dataset = raw_dataset.shuffle(buffer_size=128)
self.initialize_feature_description(raw_dataset)
# to avoid the issues with tf record warnings, we drop the last instances
batch_function = (
raw_dataset.ragged_batch if self.use_ragged_tensor else raw_dataset.batch
)
batched_dataset = batch_function(
batch_size, drop_remainder=self.should_drop_last_batch
).prefetch(buffer_size=tf.data.AUTOTUNE)
for batch_tf_processing_function in self.batch_tf_processing_functions:
batched_dataset = batched_dataset.map(batch_tf_processing_function)
dataset_iterator = iter(batched_dataset)
curr_batch = self._get_next_example(dataset_iterator)
while curr_batch is not None:
example = tf.io.parse_example(curr_batch, self.feature_description)
yield example
curr_batch = self._get_next_example(dataset_iterator)
# dynamic inferring the feature description of tfrecord files
def infer_feature_type(self, example_proto: tf.Tensor) -> dict:
feature_description = {}
tf_feature_type = (
tf.io.RaggedFeature if self.use_ragged_tensor else tf.io.VarLenFeature
)
for key, value in example_proto.items(): # type: ignore
if isinstance(value, tf.train.Feature):
if value.HasField("bytes_list"):
feature_description[key] = tf_feature_type(tf.string)
elif value.HasField("float_list"):
feature_description[key] = tf_feature_type(tf.float32)
elif value.HasField("int64_list"):
feature_description[key] = tf_feature_type(tf.int64)
else:
raise ValueError("Unknown feature type")
return feature_description
# parsing the tfrecord files from bytes
def parse_tfrecord(self, record: tf.Tensor) -> tf.Tensor:
example = tf.train.Example()
example.ParseFromString(record.numpy()) # type: ignore
return example
def shuffle(self, seed=42) -> RawDataIterator:
# TODO(lneves): Unify the shuffle method for all iterators
# Currently this one shuffles only files, parquet shuffles rows.
random.seed(seed)
random.shuffle(self.list_of_file_paths) # type: ignore
return self
def get_file_suffix(self) -> str:
return "tfrecord.gz" |