File size: 632 Bytes
9dd3461 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
from torch.utils.data.datapipes.datapipe import DataChunk
from torch.utils.data.datapipes.dataframe import dataframe_wrapper as df_wrapper
__all__ = ["DataChunkDF", ]
class DataChunkDF(DataChunk):
"""
DataChunkDF iterating over individual items inside of DataFrame containers,
to access DataFrames user `raw_iterator`
"""
def __iter__(self):
for df in self.items:
for record in df_wrapper.iterate(df):
yield record
def __len__(self):
total_len = 0
for df in self.items:
total_len += df_wrapper.get_len(df)
return total_len
|