Diff Coverage

            content_image_mask = [token == image_token_id for token in content_ids]
            images_seq_mask += content_image_mask
            num_image_tokens = sum(content_image_mask)
            if num_image_tokens % 576:
                raise ValueError("Each Janus image must have 576 placeholder tokens")
            images_emb_mask.extend([[True] * 576 for _ in range(num_image_tokens // 576)])

            labels += _janus_labels(content_ids, image_token_id, message["loss_mask"], task_type)

        msg = f"{module} could not be imported"
    if alt is None:
        logger.debug(msg)
        return False, None
    return False, alt


# Optional dependencies for object-storage support. Install via
# Install ``boto3`` for S3 or ``multi-storage-client`` for MSC support.

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``.

            ``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"

        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

        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

        self.idx_writer.write(numpy.array(document_indices, dtype=numpy.int64).tobytes(order="C"))

        # the mode per sequence
        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

        ).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

    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

    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"""

            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`."""

            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,

                ``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,

            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:
            # .idx is already cached locally; determine local path for _IndexReader
            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:
            assert os.path.exists(idx_path) and os.path.exists(bin_path), (
                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

                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

                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]]:

        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.

        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  # existence check deferred to download time
        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

            multimodal (bool, optional): Whether the dataset is multimodal.
                Defaults to False.
        """
        # The data file must stay open until finalize(); it is closed there.
        self.data_file = open(bin_path, "wb")  # pylint: disable=R1732
        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

            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

            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
        """
        # Concatenate index
        index = _IndexReader(get_idx_path(path_prefix), multimodal=self.multimodal)
        assert index.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:
            assert index.sequence_modes is not None, "sequence_modes cannot not be None"
            self.sequence_modes.extend(index.sequence_modes)

        # Free up memory to make space for new indices
        del index
        gc.collect()

        # Concatenate data
        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

    PunktLanguageVars = object
    NLTK_AVAILABLE = False

# Store generated samples in the indexed ``.bin/.idx`` format.
from hyper_parallel.data.tools import io as indexed_dataset


class CustomLanguageVars(PunktLanguageVars):
    """Preserve newline runs when Punkt detects sentence boundaries."""