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

Method __getitem__

codegeex/megatron/data/indexed_dataset.py:503–522  ·  view source on GitHub ↗
(self, idx)

Source from the content-addressed store, hash-verified

501
502 # @lru_cache(maxsize=8)
503 def __getitem__(self, idx):
504 if isinstance(idx, int):
505 ptr, size = self._index[idx]
506 np_array = np.frombuffer(
507 self._bin_buffer, dtype=self._index.dtype, count=size, offset=ptr
508 )
509 return np_array
510 elif isinstance(idx, slice):
511 start, stop, step = idx.indices(len(self))
512 if step != 1:
513 raise ValueError("Slices into indexed_dataset must be contiguous")
514 ptr = self._index._pointers[start]
515 sizes = self._index._sizes[idx]
516 offsets = list(accumulate(sizes))
517 total_size = sum(sizes)
518 np_array = np.frombuffer(
519 self._bin_buffer, dtype=self._index.dtype, count=total_size, offset=ptr
520 )
521 sents = np.split(np_array, offsets[:-1])
522 return sents
523
524 def get(self, idx, offset=0, length=None):
525 """Retrieves a single item from the dataset with the option to only

Callers

nothing calls this directly

Calls

no outgoing calls

Tested by

no test coverage detected