Diff Coverage

def _compute_vocab_range(input_local: Tensor, vocab_size: int, mesh: DeviceMesh, mesh_dim: int) -> Tuple[int, int]:
    """Compute the global vocabulary range represented by the local tensor."""
    vocab_start = _compute_vocab_start(vocab_size, mesh.size(mesh_dim), mesh.get_local_rank(mesh_dim))
    return vocab_start, vocab_start + input_local.shape[-1]


def distributed_log_softmax(logits_local: Tensor, dim: int, mesh: DeviceMesh, mesh_dim: int = 0) -> Tensor:
    """Stable log-softmax on a class-sharded dimension."""

) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
    """Index target classes and compute the local NLL contribution."""
    target_flat = target.flatten()
    target_mask = (target_flat >= vocab_start) & (target_flat < vocab_end)
    target_mask = target_mask & (target_flat != ignore_index)

    if reduction == "none":
        loss = torch.zeros(target.numel(), dtype=log_probs.dtype, device=log_probs.device)
    else:
        loss = torch.zeros(1, dtype=log_probs.dtype, device=log_probs.device)
    total_weight = torch.zeros(1, dtype=log_probs.dtype, device=log_probs.device)

    total_weight = torch.zeros(1, dtype=log_probs.dtype, device=log_probs.device)

    if target_mask.any():
        local_target = target_flat[target_mask] - vocab_start
        selected_log_probs = log_probs.reshape(-1, log_probs.shape[-1])[torch.where(target_mask)[0], local_target]

        if weight is not None:
            sample_weights = weight[target_flat[target_mask]]
            selected_log_probs = selected_log_probs * sample_weights
            total_weight = sample_weights.sum().reshape(1)
        else:
            total_weight = torch.tensor(

                target_mask.sum().item(), dtype=log_probs.dtype, device=log_probs.device
            ).reshape(1)

        if reduction == "none":
            loss_flat = torch.zeros(target.numel(), dtype=log_probs.dtype, device=log_probs.device)
            loss_flat[target_mask] = -selected_log_probs
            loss = loss_flat.reshape(target.shape)
        elif reduction == "sum":
            loss = -selected_log_probs.sum().unsqueeze(0)
        else:
            loss = -selected_log_probs.sum().unsqueeze(0)
    else:
        if reduction == "none":
            loss = torch.zeros(target.numel(), dtype=log_probs.dtype, device=log_probs.device).reshape(target.shape)
        total_weight = torch.zeros(1, dtype=log_probs.dtype, device=log_probs.device)

    return loss, total_weight, target_mask, torch.tensor(vocab_start, dtype=torch.long, device=log_probs.device)

            mesh: DeviceMesh,
            mesh_dim: int,
    ) -> Tensor:
        """Forward pass."""
        vocab_start, vocab_end = _compute_vocab_range(input_local, vocab_size, mesh, mesh_dim)

        log_probs_local = distributed_log_softmax(input_local, dim=-1, mesh=mesh, mesh_dim=mesh_dim)
        nll_result = distributed_nll_loss_forward(
            log_probs_local, target, weight, ignore_index, reduction, vocab_start, vocab_end
        )

        if reduction == "mean":
            total_loss = _differentiable_all_reduce(nll_result[0], op="sum", group=mesh.get_group(mesh_dim))
            total_weight_sum = _differentiable_all_reduce(nll_result[1], op="sum", group=mesh.get_group(mesh_dim))
            ctx.save_for_backward(
                input_local, log_probs_local, target, weight, total_weight_sum, nll_result[2], nll_result[3]
            )
            ctx.reduction = reduction

            )
            ctx.reduction = reduction
            ctx.ignore_index = ignore_index
            ctx.vocab_size = vocab_size
            ctx.local_vocab_size = input_local.shape[-1]
            ctx.mesh = mesh
            ctx.mesh_dim = mesh_dim
            ctx.vocab_start = vocab_start
            ctx.vocab_end = vocab_end

            )
            ctx.reduction = reduction
            ctx.ignore_index = ignore_index
            ctx.vocab_size = vocab_size
            ctx.local_vocab_size = input_local.shape[-1]
            ctx.mesh = mesh
            ctx.mesh_dim = mesh_dim
            ctx.vocab_start = vocab_start
            ctx.vocab_end = vocab_end
            return _differentiable_all_reduce(nll_result[0], op="sum", group=mesh.get_group(mesh_dim))

        ctx.save_for_backward(
            input_local,
            log_probs_local,

        )
        ctx.reduction = reduction
        ctx.ignore_index = ignore_index
        ctx.vocab_size = vocab_size
        ctx.local_vocab_size = input_local.shape[-1]
        ctx.mesh = mesh
        ctx.mesh_dim = mesh_dim
        ctx.vocab_start = vocab_start
        ctx.vocab_end = vocab_end
        return nll_result[0]

    @staticmethod
    def backward(ctx: Any, grad_output: Tensor) -> Tuple[Optional[Tensor], ...]:
        """Backward pass."""

        """Backward pass."""
        _, log_probs_local, target, weight, total_weight, _, _ = ctx.saved_tensors
        target_flat = target.flatten()
        softmax_local = log_probs_local.exp()
        ignore_mask = target_flat != ctx.ignore_index
        sample_weights = weight[target_flat] if weight is not None else None

        if ctx.reduction == "mean":
            grad_scale = grad_output / total_weight.clamp(min=1e-12)
        elif ctx.reduction == "sum":
            grad_scale = grad_output
        else:
            grad_scale = grad_output.flatten()

        in_vocab_mask = (target_flat >= ctx.vocab_start) & (target_flat < ctx.vocab_end) & ignore_mask
        if ctx.reduction == "none":
            grad_input = softmax_local * grad_scale.unsqueeze(-1)
            if sample_weights is not None:
                grad_input = grad_input * sample_weights.unsqueeze(-1)
        else:
            if sample_weights is not None:
                grad_scale = grad_scale * sample_weights.unsqueeze(-1)
            grad_input = softmax_local * grad_scale.unsqueeze(-1)

                grad_scale = grad_scale * sample_weights.unsqueeze(-1)
            grad_input = softmax_local * grad_scale.unsqueeze(-1)

        if in_vocab_mask.any():
            if ctx.reduction == "none":
                grad_values = -grad_scale * sample_weights if sample_weights is not None else -grad_scale
            else:
                grad_values = -grad_scale.expand_as(target_flat)
            grad_input = grad_input.contiguous()
            grad_input[
                torch.arange(target.numel(), device=target.device, dtype=torch.long)[in_vocab_mask],
                torch.where(in_vocab_mask, target_flat - ctx.vocab_start, torch.zeros_like(target_flat))[in_vocab_mask],
            ] += grad_values[in_vocab_mask]

                torch.where(in_vocab_mask, target_flat - ctx.vocab_start, torch.zeros_like(target_flat))[in_vocab_mask],
            ] += grad_values[in_vocab_mask]

        if not ignore_mask.all():
            if ctx.reduction == "none":
                grad_input[~ignore_mask] = 0.0
            else:
                ignore_indices_expanded = (~ignore_mask).unsqueeze(-1).expand_as(grad_input)
                grad_input[ignore_indices_expanded] = 0.0

    if int(os.environ.get("RANK", "0")) != 0:
        return

    _ = _download_jsonl(config)
    prepare_offline_dataset(config.to_offline_args())


if __name__ == "__main__":

        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)

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

            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

            path_prefix (str): The index (.idx) and data (.bin) prefix
        """
        # Concatenate index
        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:])

        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)

        # Free up memory to make space for new indices
        del index

except ImportError:
    PunktLanguageVars = object
    NLTK_AVAILABLE = False
else:
    NLTK_AVAILABLE = True

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

logger = get_dataset_logger(__name__)


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

            return
        elapsed = time.time() - proc_start
        docs_per_second = count / elapsed
        megabytes_per_second = total_bytes_processed / elapsed / 1024 / 1024
        logger.info(
            "Processed %d documents (%s docs/s, %s MB/s).",
            count,
            docs_per_second,
            megabytes_per_second,

    def split_sentences(self, file_name: tuple[str, str]) -> None:
        """Split every document in one JSONL partition into sentences."""
        input_file_name, output_file_name = file_name
        logger.info("Opening %s", input_file_name)
        encoder = Encoder(self.args)
        with open(input_file_name, "r", encoding="utf-8") as input_file, open(
            output_file_name,
            "w",

            finally:
                pool.close()
                pool.join()

    def _create_builders(self, output_prefix: str, tokenizer: Any) -> tuple[dict[str, str], dict[str, Any]]:
        """Create indexed dataset builders for every configured JSON key."""
        level = "sentence" if self.args.split_sentences else "document"
        output_idx_files = {}
        builders = {}
        for key in self.args.json_keys:
            output_bin_file = f"{output_prefix}_{key}_{level}.bin"
            output_idx_files[key] = f"{output_prefix}_{key}_{level}.idx"
            builders[key] = indexed_dataset.IndexedDatasetBuilder(
                output_bin_file,
                dtype=indexed_dataset.DType.optimal_dtype(len(tokenizer)),
            )
        return output_idx_files, builders

    def _get_chunk_size(self) -> int | None:
        """Return the packed chunk size including the shifted target token."""
        pack_to_seq_len = getattr(self.args, "pack_to_seq_len", None)
        return pack_to_seq_len + 1 if pack_to_seq_len is not None else None

    def _write_encoded_document(
        self,
        builders: dict[str, Any],
        document: dict[str, list[int]],
        sentence_lengths: dict[str, list[int]],

        token_buffers: dict[str, list[int]],
        chunk_size: int | None,
    ) -> None:
        """Write one encoded document, optionally packing fixed-size chunks."""
        for key in self.args.json_keys:
            if chunk_size is None:
                builders[key].add_document(document[key], sentence_lengths[key])
                continue
            token_buffers[key].extend(document[key])
            complete_length = len(token_buffers[key]) // chunk_size * chunk_size
            for offset in range(0, complete_length, chunk_size):
                chunk = token_buffers[key][offset : offset + chunk_size]
                builders[key].add_document(chunk, [chunk_size])
            del token_buffers[key][:complete_length]

    def _process_encoded_documents(
        self,
        pool: Any,
        encoder: Encoder,
        input_file_name: str,

        chunk_size: int | None,
        proc_start: float,
    ) -> None:
        """Encode input records and write them to the indexed dataset builders."""
        total_bytes_processed = 0
        with open(input_file_name, "r", encoding="utf-8") as input_file:
            encoded_docs = pool.imap(encoder.encode, input_file, 32)
            for count, (document, sentence_lengths, bytes_processed) in enumerate(encoded_docs, start=1):
                if self.args.find_optimal_num_workers and count > self.args.max_documents:
                    break
                total_bytes_processed += bytes_processed
                self._write_encoded_document(
                    builders, document, sentence_lengths, token_buffers, chunk_size
                )
                self.print_processing_stats(count, proc_start, total_bytes_processed)

    def process_json_file(self, file_name: tuple[str, str]) -> list[float]:
        """Tokenize one JSONL partition into indexed dataset .bin/.idx files.

        Returns:
            Throughput measurements collected while benchmarking workers.
        """
        input_file_name, output_prefix = file_name
        logger.info("Opening %s", input_file_name)

        startup_start = time.time()
        encoder = Encoder(self.args)
        tokenizer = build_tokenizer(self.args)

        # multiprocessing.Pool must be shut down with close()+join() so
        # in-flight tasks finish; a 'with' block would terminate() them.
        pool = multiprocessing.Pool(self.workers, initializer=encoder.initializer)  # pylint: disable=R1732

        output_idx_files, builders = self._create_builders(output_prefix, tokenizer)
        timing = (time.time() - startup_start, time.time())
        logger.info("Time to startup: %s", timing[0])
        chunk_size = self._get_chunk_size()
        token_buffers = {key: [] for key in self.args.json_keys}
        try:
            self._process_encoded_documents(
                pool, encoder, input_file_name, builders, token_buffers, chunk_size, timing[1]
            )
            for key in self.args.json_keys:
                builders[key].finalize(output_idx_files[key])
        finally:
            for builder in builders.values():
                builder.data_file.close()

            results.append((workers, float(np.mean(measurements))))
    if not results:
        raise ValueError("No worker performance measurements were collected")
    results.sort(key=lambda item: item[1], reverse=True)
    logger.info("Worker performance results:")
    for position, (workers, average_rate) in enumerate(results, start=1):
        logger.info("%d. %d workers: %.4f docs/s", position, workers, average_rate)
    best_workers, best_rate = results[0]
    logger.info(
        "Best configuration: %d total workers (%d per partition), %.4f docs/s.",
        best_workers,
        best_workers // partitions,
        best_rate,

    partition_size: int | None,
) -> None:
    """Distribute input records across partition files."""
    outputs = [open(name["partition"], "w", encoding="utf-8") for name in names]  # pylint: disable=R1732
    partition_index = 0
    line_count = 0
    try:
        for input_file_name in input_files:
            open_file = gzip.open if input_file_name.endswith(".gz") else open
            with open_file(input_file_name, "rt", encoding="utf-8") as input_file:

    performance = {}
    input_files = _resolve_input_files(args.dataset_name_or_path)
    for workers in worker_candidates:
        logger.info("Processing data with %d workers.", workers)
        workers_per_partition = workers // args.partitions

        if args.split_sentences:
            nltk.download("punkt", quiet=True, download_dir=os.environ.get("NLTK_DATA"))