7633129ba8d5f0e28bd6b6d6027b14352482ef31,fairseq/data/token_block_dataset.py,TokenBlockDataset,prefetch,#TokenBlockDataset#Any#,104

Before Change

    def prefetch(self, indices):
        indices.sort()
        total_size = 0
        for idx in indices:
            s, e = self.slice_indices[idx]
            total_size += e - s
        self.cache = np.empty(total_size, dtype=np.int32)
        start = 0
        for idx in indices:
            s, e = self.slice_indices[idx]

After Change

        return getattr(self.dataset, "supports_prefetch", False)

    def prefetch(self, indices):
        self.dataset.prefetch({
            ds_idx
            for index in indices
            for start_ds_idx, _, end_ds_idx in [self.block_to_dataset_index[index]]
            for ds_idx in range(start_ds_idx, end_ds_idx + 1)
        })

Instances