12345678910111213141516171819202122 |
- 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
|