structures.py 632 B

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