"""Load indexed datasets from paired binary and index files.
Supports the original mmap and file-pointer readers for local *.bin / *.idx
pairs, plus optional streaming readers for object storage (S3 and MSC).
All three calls below are equivalent for local data:
from hyper_parallel.data.indexed.io import IndexedDataset
ds = IndexedDataset("/path/to/shard_00_text_document")
print(len(ds), ds[0][:20])
ds = IndexedDataset("/path/to/shard_00_text_document.bin")
print(len(ds), ds[0][:20])
ds = IndexedDataset("/path/to/shard_00_text_document.idx")
print(len(ds), ds[0][:20])
For object-storage data, pass an :class:`ObjectStorageConfig`:
cfg = ObjectStorageConfig(path_to_idx_cache="/tmp/idx_cache")
ds = IndexedDataset("s3://bucket/path/shard_00_text_document", object_storage_config=cfg)
"""
from __future__ import annotations
import gc
import os
import shutil
import struct
from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import Enum
from types import TracebackType
from typing import Any, Dict, List, Optional, Tuple, Type, Union
import importlib
import traceback
import numpy
import torch
from hyper_parallel.data.dataset_logging import get_dataset_logger
logger = get_dataset_logger(__name__)
def safe_import(
module: str,
*,
msg: Optional[str] = None,
alt: Optional[Any] = None,
) -> Tuple[bool, Optional[Any]]:
"""
A function used to import modules that may not be available.
This function will attempt to import a module with the given name, but it
will not throw an ImportError if the module is not found. Instead, it will
return the configured alternate object or ``None``.
Args:
module (str): The name of the module to import.
msg (str or None): An optional error message to be displayed if this module is used
after a failed import.
alt (object): An optional module to be used in place of the given module if it
fails to import
Returns:
Tuple(bool, object): Whether the intended import succeeded and the imported module,
configured alternate, or ``None``.
"""
try:
return True, importlib.import_module(module)
except ImportError:
exception_text = traceback.format_exc()
logger.debug(f"Import of {module} failed with: {exception_text}")
if msg is None:
msg = f"{module} could not be imported"
if alt is None:
logger.debug(msg)
return False, None
return False, alt
HAS_BOTO3, boto3 = safe_import("boto3")
HAS_MSC, multi_storage_client = safe_import("multi_storage_client")
_S3_PREFIX = "s3://"
_MSC_PREFIX = "msc://"
@dataclass
class ObjectStorageConfig:
"""Configuration for reading ``.bin``/``.idx`` files from object storage.
Attributes:
path_to_idx_cache: Local directory where the ``.idx`` file is cached on
first use. Re-used across ranks via a per-host directory layout.
bin_chunk_nbytes: Size in bytes of each chunked range read against the
``.bin`` object. Defaults to 256 MiB. Larger values reduce request
count but increase per-rank memory footprint.
"""
path_to_idx_cache: str
bin_chunk_nbytes: int = 256 * 1024 * 1024
def _is_object_storage_path(path: str) -> bool:
"""Return ``True`` if ``path`` is an ``s3://`` or ``msc://`` URI."""
return path.startswith(_S3_PREFIX) or path.startswith(_MSC_PREFIX)
def _parse_s3_path(path: str) -> Tuple[str, str]:
"""Split an ``s3://bucket/key`` URI into ``(bucket, key)``."""
if not path.startswith(_S3_PREFIX):
raise ValueError(f"Not an S3 path: {path}")
parts = path[len(_S3_PREFIX) :].split("/")
bucket = parts[0]
key = "/".join(parts[1:]) if len(parts) > 1 else ""
return bucket, key
def _get_index_cache_path(idx_path: str, object_storage_config: ObjectStorageConfig) -> str:
"""Return the local cache path for ``idx_path`` under ``path_to_idx_cache``."""
if idx_path.startswith(_S3_PREFIX):
stripped = idx_path[len(_S3_PREFIX) :]
elif idx_path.startswith(_MSC_PREFIX):
stripped = idx_path[len(_MSC_PREFIX) :]
else:
raise ValueError(f"Not an object storage path: {idx_path}")
return os.path.join(object_storage_config.path_to_idx_cache, stripped)
def _cache_index_file(remote_path: str, local_path: str) -> None:
"""Download ``.idx`` from object storage to ``local_path``.
Rank 0 performs the download and other ranks wait on a ``torch.distributed``
barrier. If the local file already exists this is a no-op.
Raises:
ImportError: If the relevant client library (``boto3`` for ``s3://`` or
``multi_storage_client`` for ``msc://``) is not installed.
ValueError: If ``remote_path`` is neither an ``s3://`` nor an
``msc://`` URI.
"""
torch_dist_enabled = torch.distributed.is_initialized()
rank = torch.distributed.get_rank() if torch_dist_enabled else 0
if remote_path.startswith(_S3_PREFIX):
if not HAS_BOTO3:
raise ImportError("boto3 is required to read s3:// datasets. Install via `pip install boto3`.")
if not os.path.exists(local_path):
if not torch_dist_enabled or rank == 0:
os.makedirs(os.path.dirname(local_path), exist_ok=True)
bucket, key = _parse_s3_path(remote_path)
client = boto3.client("s3")
logger.info("Downloading %s -> %s", remote_path, local_path)
client.download_file(bucket, key, local_path)
client.close()
elif remote_path.startswith(_MSC_PREFIX):
if not HAS_MSC:
raise ImportError(
"multi_storage_client is required to read msc:// datasets. "
"Install via `pip install multi-storage-client`."
)
if not os.path.exists(local_path):
if not torch_dist_enabled or rank == 0:
os.makedirs(os.path.dirname(local_path), exist_ok=True)
multi_storage_client.download_file(remote_path, local_path)
else:
raise ValueError(f"Unsupported object storage path: {remote_path}")
if torch_dist_enabled:
torch.distributed.barrier()
if not os.path.exists(local_path):
raise RuntimeError(f"Index cache file not found after download: {local_path}")
_INDEX_HEADER = b"MMIDIDX\x00\x00"
class DType(Enum):
"""The NumPy data type Enum for reading the IndexedDataset indices"""
UINT8 = 1
INT8 = 2
INT16 = 3
INT32 = 4
INT64 = 5
FLOAT64 = 6
FLOAT32 = 7
UINT16 = 8
@classmethod
def code_from_dtype(cls, value: Type[numpy.number]) -> int:
"""Get the code from the dtype
Args:
value (Type[numpy.number]): The dtype
Returns:
int: The code
"""
return cls[value.__name__.upper()].value
@classmethod
def dtype_from_code(cls, value: int) -> Type[numpy.number]:
"""Get the dtype from the code
Args:
value (int): The code
Returns:
Type[numpy.number]: The dtype
"""
return getattr(numpy, cls(value).name.lower())
@classmethod
def size(cls, key: Union[int, Type[numpy.number]]) -> int:
"""Get the size of the dtype/code in bytes
Args:
key (Union[int, Type[numpy.number]]): The dtype or code
Raises:
ValueError: If the key is neither dtype nor integer code
Returns:
int: The size of the dtype/code in bytes
"""
if isinstance(key, int):
return cls.dtype_from_code(key)().itemsize
if numpy.number in key.__mro__:
return key().itemsize
raise ValueError("Invalid key passed to DType.size()")
@classmethod
def optimal_dtype(cls, cardinality: Optional[int]) -> Type[numpy.number]:
"""Get the dtype to use for an index of a certain cardinality
Args:
cardinality (Optional[int]): The number of elements to be indexed
Returns:
Type[numpy.number]: The dtype to use for the index
"""
if cardinality is not None and cardinality < 65500:
return numpy.uint16
return numpy.int32
class _IndexWriter:
"""Object class to write the index (.idx) file
Args:
idx_path (str): The path to the index file
dtype (Type[numpy.number]): The dtype of the index file
"""
def __init__(self, idx_path: str, dtype: Type[numpy.number]) -> None:
"""Initialize the _IndexWriter
Args:
idx_path (str): The path to the index file.
dtype (Type[numpy.number]): The dtype of the index file.
"""
self.idx_path = idx_path
self.dtype = dtype
self.idx_writer = None
def __enter__(self) -> "_IndexWriter":
"""Enter the context introduced by the 'with' keyword
Returns:
_IndexWriter: The instance
"""
self.idx_writer = open(self.idx_path, "wb")
self.idx_writer.write(_INDEX_HEADER)
self.idx_writer.write(struct.pack("<Q", 1))
self.idx_writer.write(struct.pack("<B", DType.code_from_dtype(self.dtype)))
return self
def __exit__(
self,
exc_type: Optional[Type[BaseException]],
exc_val: Optional[BaseException],
exc_tb: Optional[TracebackType],
) -> Optional[bool]:
"""Exit the context introduced by the 'with' keyword
Args:
exc_type (Optional[Type[BaseException]]): Exception type
exc_val (Optional[BaseException]): Exception value
exc_tb (Optional[TracebackType]): Exception traceback object
Returns:
Optional[bool]: Whether to silence the exception
"""
self.idx_writer.close()
def write(
self,
sequence_lengths: List[int],
sequence_modes: Optional[List[int]],
document_indices: List[int],
) -> None:
"""Write the index (.idx) file
Args:
sequence_lengths (List[int]): The length of each sequence
sequence_modes (Optional[List[int]]): The mode of each sequences
document_indices (List[int]): The seqyebce indices demarcating the end of each document
"""
sequence_pointers = self._sequence_pointers(sequence_lengths)
sequence_count = len(sequence_lengths)
self.idx_writer.write(struct.pack("<Q", sequence_count))
document_count = len(document_indices)
self.idx_writer.write(struct.pack("<Q", document_count))
self.idx_writer.write(numpy.array(sequence_lengths, dtype=numpy.int32).tobytes(order="C"))
self.idx_writer.write(numpy.array(sequence_pointers, dtype=numpy.int64).tobytes(order="C"))
self.idx_writer.write(numpy.array(document_indices, dtype=numpy.int64).tobytes(order="C"))
if sequence_modes is not None:
self.idx_writer.write(numpy.array(sequence_modes, dtype=numpy.int8).tobytes(order="C"))
def _sequence_pointers(self, sequence_lengths: List[int]) -> List[int]:
"""Build the sequence pointers per the sequence lengths and dtype size
Args:
sequence_lengths (List[int]): The length of each sequence
Returns:
List[int]: The pointer to the beginning of each sequence
"""
itemsize = DType.size(self.dtype)
curr_ptr = 0
list_ptr = []
for length in sequence_lengths:
list_ptr.append(curr_ptr)
curr_ptr += length * itemsize
return list_ptr
class _IndexReader:
"""Object class to read the index (.idx) file
Args:
idx_path (str): The path to the index file
multimodal (bool): Whether the dataset is multimodal
"""
def __init__(self, idx_path: str, multimodal: bool) -> None:
"""Initialize the _IndexReader
Args:
idx_path (str): The path to the index file.
multimodal (bool): Whether the dataset is multimodal.
"""
logger.info("Loading index file %s", idx_path)
with open(idx_path, "rb") as f:
header = f.read(9)
if header != _INDEX_HEADER:
raise ValueError(f"Bad header in {idx_path}")
version = struct.unpack("<Q", f.read(8))[0]
if version != 1:
raise ValueError(f"Unsupported index version {version} in {idx_path}")
code = struct.unpack("<B", f.read(1))[0]
self.dtype = DType.dtype_from_code(code)
self.dtype_size = DType.size(self.dtype)
self.sequence_count = struct.unpack("<Q", f.read(8))[0]
self.document_count = struct.unpack("<Q", f.read(8))[0]
payload_offset = f.tell()
self._mmap = numpy.memmap(idx_path, mode="r", order="C")
self._buffer = memoryview(self._mmap)
logger.info("Extracting sequence lengths")
self.sequence_lengths = numpy.frombuffer(
self._buffer, dtype=numpy.int32, count=self.sequence_count, offset=payload_offset
).copy()
logger.info("Extracting sequence pointers")
self.sequence_pointers = numpy.frombuffer(
self._buffer,
dtype=numpy.int64,
count=self.sequence_count,
offset=payload_offset + self.sequence_lengths.nbytes,
).copy()
logger.info("Extracting document indices")
self.document_indices = numpy.frombuffer(
self._buffer,
dtype=numpy.int64,
count=self.document_count,
offset=payload_offset + self.sequence_lengths.nbytes + self.sequence_pointers.nbytes,
).copy()
self.sequence_modes: Optional[numpy.ndarray] = None
if multimodal:
logger.info("Extracting sequence modes")
self.sequence_modes = numpy.frombuffer(
self._buffer,
dtype=numpy.int8,
count=self.sequence_count,
offset=payload_offset
+ self.sequence_lengths.nbytes
+ self.sequence_pointers.nbytes
+ self.document_indices.nbytes,
).copy()
sequence_length_count = self.sequence_lengths.shape[0]
if sequence_length_count != len(self):
raise ValueError("Sequence length count does not match the dataset length")
if sequence_length_count != self.sequence_count:
raise ValueError("Sequence length count does not match the index sequence count")
if sequence_length_count != self.document_indices[-1]:
raise ValueError("Sequence length count does not match the final document index")
logger.info("Sequences: %d | Documents: %d", len(self), self.document_indices.shape[0] - 1)
def close(self) -> None:
"""Release memory-mapped index resources."""
if getattr(self, "_mmap", None) is None:
return
self.sequence_lengths = None
self.sequence_pointers = None
self.document_indices = None
self.sequence_modes = None
if getattr(self, "_buffer", None) is not None:
self._buffer.release()
self._buffer = None
self._mmap = None
def __del__(self) -> None:
"""Clean up the object."""
try:
self.close()
except (AttributeError, BufferError, ValueError):
pass
def __len__(self) -> int:
"""Get the number of sequences in the dataset
Returns:
int: The number of sequences in the dataset
"""
return self.sequence_count
def __getitem__(self, idx: int) -> Tuple[numpy.int32, numpy.int64, Optional[numpy.int8]]:
"""Return the pointer, length, and mode at the index
Args:
idx (int): The index into the dataset
Returns:
Tuple[numpy.int32, numpy.int64, Optional[numpy.int8]]: The pointer, length and mode
at the index
"""
return (
self.sequence_pointers[idx],
self.sequence_lengths[idx],
self.sequence_modes[idx] if self.sequence_modes is not None else None,
)
class _BinReader(ABC):
"""Abstract class to read the data (.bin) file"""
@abstractmethod
def read(self, dtype: Type[numpy.number], count: int, offset: int) -> numpy.ndarray:
"""Read bytes into a numpy array.
Args:
dtype (Type[numpy.number]): Data-type of the returned array.
count (int): Number of items to read.
offset (int): Start reading from this offset (in bytes).
Returns:
numpy.ndarray: An array with `count` items and data-type `dtype` constructed from
reading bytes from the data file starting at `offset`.
"""
class _MMapBinReader(_BinReader):
"""A _BinReader that memory maps the data (.bin) file"""
def __init__(self, bin_path: str) -> None:
"""Initialize the _MMapBinReader
Args:
bin_path (str): The path to the data (.bin) file.
"""
self._file = open(bin_path, "rb")
self._mmap = numpy.memmap(self._file, mode="r", order="C")
self._buffer = memoryview(self._mmap.data)
def read(self, dtype: Type[numpy.number], count: int, offset: int) -> numpy.ndarray:
"""Read bytes into a numpy array.
Args:
dtype (Type[numpy.number]): Data-type of the returned array.
count (int): Number of items to read.
offset (int): Start reading from this offset (in bytes).
Returns:
numpy.ndarray: An array with `count` items and data-type `dtype` constructed from
reading bytes from the data file starting at `offset`.
"""
return numpy.frombuffer(self._buffer, dtype=dtype, count=count, offset=offset)
def close(self) -> None:
"""Close the backing file without invalidating arrays returned from mmap."""
file_obj = getattr(self, "_file", None)
if file_obj is not None:
file_obj.close()
self._file = None
def __del__(self) -> None:
"""Clean up the object."""
try:
self.close()
except (AttributeError, BufferError, ValueError):
pass
class _FileBinReader(_BinReader):
"""A _BinReader that reads from the data (.bin) file using a file pointer"""
def __init__(self, bin_path: str) -> None:
"""Initialize the _FileBinReader
Args:
bin_path (str): The path to the data (.bin) file.
"""
self._bin_path = bin_path
def read(self, dtype: Type[numpy.number], count: int, offset: int) -> numpy.ndarray:
"""Read bytes into a numpy array.
Args:
dtype (Type[numpy.number]): Data-type of the returned array.
count (int): Number of items to read.
offset (int): Start reading from this offset (in bytes).
Returns:
numpy.ndarray: An array with `count` items and data-type `dtype` constructed from
reading bytes from the data file starting at `offset`.
"""
out = numpy.empty(count, dtype=dtype)
with open(self._bin_path, "rb", buffering=0) as f:
f.seek(offset)
byte_view = out.view(numpy.uint8)
nbytes = f.readinto(byte_view)
if nbytes != out.nbytes:
raise ValueError(
f"Short read from {self._bin_path}: requested {out.nbytes} bytes at offset {offset}, "
f"read {nbytes or 0} bytes"
)
return out
class _S3BinReader(_BinReader):
"""Stream ``.bin`` data from S3 via chunked ranged ``GetObject`` calls.
A single in-memory chunk (sized by
:attr:`ObjectStorageConfig.bin_chunk_nbytes`) is cached so consecutive
reads within the same chunk avoid network round-trips. Random-access
reads outside the current chunk trigger a new ranged ``GetObject``.
"""
def __init__(self, bin_path: str, object_storage_config: ObjectStorageConfig) -> None:
"""Initialize the _S3BinReader
Args:
bin_path (str): The ``s3://bucket/key`` URI of the data (.bin) object.
object_storage_config (ObjectStorageConfig): Chunked-read configuration.
"""
if not HAS_BOTO3:
raise ImportError("boto3 is required to read s3:// datasets. Install via `pip install boto3`.")
if object_storage_config.bin_chunk_nbytes <= 0:
raise ValueError(f"bin_chunk_nbytes must be positive, got {object_storage_config.bin_chunk_nbytes}")
self._client = boto3.client("s3")
self._s3_bucket, self._s3_key = _parse_s3_path(bin_path)
self._cache_nbytes = object_storage_config.bin_chunk_nbytes
self._cache_bytes_start: int = 0
self._cache_bytes_end: int = 0
self._cache: Optional[bytes] = None
def _extract_from_cache(self, offset: int, size: int) -> bytes:
"""Return a requested byte range from the active S3 cache window."""
if self._cache is None:
raise RuntimeError("Cache is empty; cannot extract before first read")
start = offset - self._cache_bytes_start
end = start + size
if start < 0 or end > len(self._cache):
raise IndexError(
f"Cache window [{self._cache_bytes_start}, {self._cache_bytes_end}) "
f"does not contain requested range [{offset}, {offset + size})"
)
return self._cache[start:end]
def read(self, dtype: Type[numpy.number], count: int, offset: int) -> numpy.ndarray:
"""Read ``count`` elements of ``dtype`` starting at byte ``offset``."""
size = count * DType.size(dtype)
if self._cache is not None and offset >= self._cache_bytes_start and offset + size <= self._cache_bytes_end:
return numpy.frombuffer(self._extract_from_cache(offset, size), dtype=dtype)
bytes_start = (offset // self._cache_nbytes) * self._cache_nbytes
bytes_end = max(bytes_start + self._cache_nbytes, offset + size)
self._cache = self._client.get_object(
Bucket=self._s3_bucket,
Key=self._s3_key,
Range=f"bytes={bytes_start}-{bytes_end - 1}",
)["Body"].read()
self._cache_bytes_start = bytes_start
self._cache_bytes_end = bytes_start + len(self._cache)
return numpy.frombuffer(self._extract_from_cache(offset, size), dtype=dtype)
def __del__(self) -> None:
try:
self._client.close()
except (AttributeError, RuntimeError):
pass
class _MultiStorageClientBinReader(_BinReader):
"""Read ``.bin`` data via NVIDIA's :mod:`multi_storage_client`."""
def __init__(self, bin_path: str, object_storage_config: ObjectStorageConfig) -> None:
"""Initialize the _MultiStorageClientBinReader
Args:
bin_path (str): The ``msc://`` URI of the data (.bin) object.
object_storage_config (ObjectStorageConfig): Accepted for constructor
parity with the other object-storage bin readers; unused here.
"""
del object_storage_config
if not HAS_MSC:
raise ImportError(
"multi_storage_client is required to read msc:// datasets. "
"Install via `pip install multi-storage-client`."
)
self._client, self._bin_path = multi_storage_client.resolve_storage_client(bin_path)
def read(self, dtype: Type[numpy.number], count: int, offset: int) -> numpy.ndarray:
"""Read ``count`` elements of ``dtype`` starting at byte ``offset``."""
size = count * DType.size(dtype)
buffer = self._client.read(
path=self._bin_path,
byte_range=multi_storage_client.types.Range(offset=offset, size=size),
)
return numpy.frombuffer(buffer, dtype=dtype)
OBJECT_STORAGE_BIN_READERS: Dict[str, Type[_BinReader]] = {
"s3": _S3BinReader,
"msc": _MultiStorageClientBinReader,
}
class IndexedDataset(torch.utils.data.Dataset):
"""A fast, on-disk dataset backed by paired index and binary files."""
def __init__(
self,
path_prefix: str,
multimodal: bool = False,
mmap: bool = True,
object_storage_config: Optional[ObjectStorageConfig] = None,
) -> None:
"""Initialize the IndexedDataset
Args:
path_prefix (str): The index (.idx) and data (.bin) prefix. May be an S3 URI
(``s3://bucket/key``) when ``object_storage_config`` is provided.
multimodal (bool): Whether the dataset is multimodal. Defaults to False.
mmap (bool): Whether to mmap the .bin files. Defaults to True. Must be False
for object-storage paths.
object_storage_config (Optional[ObjectStorageConfig]): When provided and
``path_prefix`` is an S3/MSC URI, the .idx file is downloaded to
``object_storage_config.path_to_idx_cache`` and the .bin file is streamed
via chunked GETs.
"""
super().__init__()
normalized_prefix = _normalize_prefix(path_prefix)
if _is_object_storage_path(normalized_prefix) and object_storage_config is not None:
if mmap:
raise ValueError(
"mmap must be False for object-storage prefixes; "
"set mmap=False (or set `mmap_bin_files: false` in the recipe)."
)
idx_path = get_idx_path(normalized_prefix)
cache_idx_path = _get_index_cache_path(idx_path, object_storage_config)
_cache_index_file(idx_path, cache_idx_path)
self.initialize(normalized_prefix, multimodal, mmap, object_storage_config)
def initialize(
self,
path_prefix: str,
multimodal: bool,
mmap: bool,
object_storage_config: Optional[ObjectStorageConfig] = None,
) -> None:
"""Open the index and data readers for a dataset prefix.
Args:
path_prefix (str): The index (.idx) and data (.bin) prefix.
multimodal (bool): Whether the dataset is multimodal.
mmap (bool): Whether to mmap the .bin files.
object_storage_config (Optional[ObjectStorageConfig]): Object-storage
configuration used for ``s3://``/``msc://`` prefixes.
"""
idx_path = get_idx_path(path_prefix)
bin_path = get_bin_path(path_prefix)
if _is_object_storage_path(path_prefix) and object_storage_config is not None:
local_idx_path = _get_index_cache_path(idx_path, object_storage_config)
if not os.path.exists(local_idx_path):
raise RuntimeError(f"Cached .idx not found: {local_idx_path}")
access = "s3" if path_prefix.startswith(_S3_PREFIX) else "msc"
bin_reader: _BinReader = OBJECT_STORAGE_BIN_READERS[access](bin_path, object_storage_config)
index_reader = _IndexReader(local_idx_path, multimodal)
else:
if not os.path.exists(idx_path) or not os.path.exists(bin_path):
raise FileNotFoundError(f"Missing .idx or .bin at prefix {path_prefix}")
bin_reader = _MMapBinReader(bin_path) if mmap else _FileBinReader(bin_path)
index_reader = _IndexReader(idx_path, multimodal)
self.path_prefix = path_prefix
self.multimodal = multimodal
self.mmap = mmap
self.object_storage_config = object_storage_config
self.bin_reader = bin_reader
self.index = index_reader
def __len__(self) -> int:
"""Get the number of sequences in the dataset
Returns:
int: The number of sequences in the dataset
"""
return len(self.index)
def __getitem__(
self, idx: Union[int, numpy.integer, slice]
) -> Union[
numpy.ndarray,
Tuple[numpy.ndarray, Any],
List[numpy.ndarray],
Tuple[List[numpy.ndarray], numpy.ndarray],
]:
"""Read one sequence or a contiguous slice of sequences.
Args:
idx (Union[int, numpy.integer, slice]): The index or contiguous
(step=1) slice into the dataset.
Returns:
Union[numpy.ndarray, Tuple[numpy.ndarray, Any], List[numpy.ndarray],
Tuple[List[numpy.ndarray], numpy.ndarray]]: The sequence(s), each
paired with its mode when the dataset is multimodal.
"""
if isinstance(idx, (int, numpy.integer)):
ptr, length, mode = self.index[idx]
seq = self.bin_reader.read(self.index.dtype, length, ptr)
return (seq, mode) if mode is not None else seq
if isinstance(idx, slice):
start, _, step = idx.indices(len(self))
if step != 1:
raise ValueError("Slices into IndexedDataset must be contiguous (step=1)")
lengths = self.index.sequence_lengths[idx]
modes = self.index.sequence_modes[idx] if self.multimodal else None
offsets = numpy.cumsum(lengths, dtype=numpy.int64)
token_count = int(offsets[-1]) if offsets.size else 0
buffer = self.bin_reader.read(
self.index.dtype,
token_count,
int(self.index.sequence_pointers[start]),
)
sequences = numpy.split(buffer, offsets[:-1])
return (sequences, modes) if modes is not None else sequences
raise TypeError(f"Unexpected index type {type(idx)}")
def get(
self, idx: int, offset: int = 0, length: Optional[int] = None
) -> Union[numpy.ndarray, Tuple[numpy.ndarray, Any]]:
"""Read a sub-range of one sequence.
Args:
idx (int): The index into the dataset.
offset (int): Number of leading elements to skip. Defaults to 0.
length (Optional[int]): Number of elements to read; defaults to the
remainder of the sequence.
Returns:
Union[numpy.ndarray, Tuple[numpy.ndarray, Any]]: The sequence
sub-range, paired with its mode for multimodal datasets.
"""
ptr, seq_len, mode = self.index[idx]
length = seq_len - offset if length is None else length
ptr += offset * DType.size(self.index.dtype)
seq = self.bin_reader.read(self.index.dtype, length, ptr)
return (seq, mode) if mode is not None else seq
@property
def sequence_lengths(self) -> numpy.ndarray:
"""The length of each sequence in the dataset (numpy.int32 array)."""
return self.index.sequence_lengths
@property
def document_indices(self) -> numpy.ndarray:
"""The sequence indices demarcating documents (numpy.int64 array)."""
return self.index.document_indices
@staticmethod
def exists(path_prefix: str) -> bool:
"""Return whether the .idx and .bin files exist for a dataset prefix.
Args:
path_prefix (str): The index (.idx) and data (.bin) prefix.
Returns:
bool: Whether both files exist; object-storage prefixes always
return ``True`` and defer the check to download time.
"""
if _is_object_storage_path(path_prefix):
return True
return os.path.exists(get_idx_path(path_prefix)) and os.path.exists(get_bin_path(path_prefix))
class IndexedDatasetBuilder:
"""Builder class for the IndexedDataset class
Args:
bin_path (str): The path to the data (.bin) file
dtype (Type[numpy.number], optional): The dtype of the index file. Defaults to numpy.int32.
multimodal (bool, optional): Whether the dataset is multimodal. Defaults to False.
"""
def __init__(self, bin_path: str, dtype: Type[numpy.number] = numpy.int32, multimodal: bool = False) -> None:
"""Initialize the IndexedDatasetBuilder
Args:
bin_path (str): The path to the data (.bin) file.
dtype (Type[numpy.number], optional): The dtype of the index file.
Defaults to numpy.int32.
multimodal (bool, optional): Whether the dataset is multimodal.
Defaults to False.
"""
self.data_file = open(bin_path, "wb")
self.dtype = dtype
self.multimodal = multimodal
self.sequence_lengths = []
self.document_indices = [0]
self.sequence_modes = [] if self.multimodal else None
def add_item(self, tensor: torch.Tensor, mode: int = 0) -> None:
"""Add a single item to the dataset
Args:
tensor (torch.Tensor): The item to add to the data file
mode (int, optional): The mode for the item. Defaults to 0.
"""
np_array = numpy.array(tensor.numpy(), dtype=self.dtype)
self.data_file.write(np_array.tobytes(order="C"))
self.sequence_lengths.append(np_array.size)
if self.multimodal:
self.sequence_modes.append(mode)
def add_document(self, tensor: torch.Tensor, lengths: List[int], modes: Optional[List[int]] = None) -> None:
"""Add an entire document to the dataset
Args:
tensor (torch.Tensor): The document to add
lengths (List[int]): The lengths of each item in the document
modes (Optional[List[int]], optional): The modes for each item in the document.
Defaults to None.
"""
np_array = numpy.array(tensor, dtype=self.dtype)
self.data_file.write(np_array.tobytes(order="C"))
self.sequence_lengths.extend(lengths)
self.document_indices.append(len(self.sequence_lengths))
if self.multimodal:
self.sequence_modes.extend(modes if modes is not None else [0] * lengths)
def end_document(self) -> None:
"""Finalize the document, for use with IndexedDatasetBuilder.add_item"""
self.document_indices.append(len(self.sequence_lengths))
def add_index(self, path_prefix: str) -> None:
"""Add an entire IndexedDataset to the dataset
Args:
path_prefix (str): The index (.idx) and data (.bin) prefix
"""
index = _IndexReader(get_idx_path(path_prefix), multimodal=self.multimodal)
if index.dtype != self.dtype:
raise ValueError(f"Index dtype {index.dtype} does not match builder dtype {self.dtype}")
offset = len(self.sequence_lengths)
self.sequence_lengths.extend(index.sequence_lengths)
self.document_indices.extend((offset + index.document_indices)[1:])
if self.multimodal:
if index.sequence_modes is None:
raise ValueError("sequence_modes cannot be None for a multimodal dataset")
self.sequence_modes.extend(index.sequence_modes)
del index
gc.collect()
with open(get_bin_path(path_prefix), "rb") as f:
shutil.copyfileobj(f, self.data_file)
def finalize(self, idx_path: str) -> None:
"""Clean up and write the index (.idx) file
Args:
idx_path (str): The path to the index file
"""
self.data_file.close()
with _IndexWriter(idx_path, self.dtype) as writer:
writer.write(self.sequence_lengths, self.sequence_modes, self.document_indices)
def get_idx_path(path_prefix: str) -> str:
"""Return the index-file path for an indexed dataset prefix."""
return path_prefix + ".idx"
def get_bin_path(path_prefix: str) -> str:
"""Return the binary-data path for an indexed dataset prefix."""
return path_prefix + ".bin"
def _normalize_prefix(path_prefix: str) -> str:
if path_prefix.endswith(".bin") or path_prefix.endswith(".idx"):
return path_prefix[:-4]
return path_prefix