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

Before Change

            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]
            self.dataset.read_into(s, self.cache[start:start + e - s])
            self.cache_index[idx] = (start, start + e - s)
            start += e - s

    @property
    def supports_prefetch(self):
        return True

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