MCPcopy Create free account
hub / github.com/zai-org/CodeGeeX / prefetch

Method prefetch

codegeex/megatron/data/indexed_dataset.py:243–265  ·  view source on GitHub ↗
(self, indices)

Source from the content-addressed store, hash-verified

241 return True
242
243 def prefetch(self, indices):
244 if all(i in self.cache_index for i in indices):
245 return
246 if not self.data_file:
247 self.read_data(self.path)
248 indices = sorted(set(indices))
249 total_size = 0
250 for i in indices:
251 total_size += self.data_offsets[i + 1] - self.data_offsets[i]
252 self.cache = np.empty(total_size, dtype=self.dtype)
253 ptx = 0
254 self.cache_index.clear()
255 for i in indices:
256 self.cache_index[i] = ptx
257 size = self.data_offsets[i + 1] - self.data_offsets[i]
258 a = self.cache[ptx : ptx + size]
259 self.data_file.seek(self.data_offsets[i] * self.element_size)
260 self.data_file.readinto(a)
261 ptx += size
262 if self.data_file:
263 # close and delete data file after prefetch so we can pickle
264 self.data_file.close()
265 self.data_file = None
266
267 # @lru_cache(maxsize=8)
268 def __getitem__(self, idx):

Callers

nothing calls this directly

Calls 2

read_dataMethod · 0.80
clearMethod · 0.45

Tested by

no test coverage detected