guidellm.data.builders

def process_dataset(
    data: DataArgs,
    output_path: str | Path,
    tokenizer: DataTokenizerArgs,
    strategy: PreprocessStrategyArgs,
    data_column_mapper: DataPreprocessorArgs,
    push_to_hub: bool,
    hub_dataset_id: str | None,
    random_seed: int,
    data_loader: DataLoaderArgs | None = None,
) -> None:
    """
    Main method to process and save a dataset with sampled prompt/output token counts.

    :param data_loader: Optional loader config. When ``samples`` is greater than 0,
        processing stops after that many successfully processed rows. ``shuffle`` and
        ``num_workers`` are ignored for preprocessing.
    """
    _validate_output_suffix(output_path)
    logger.info(
        "Starting dataset conversion | Input: {} | Output: {}", data, output_path
    )

    # samples > 0 caps successfully processed output rows; -1 means unlimited.
    # shuffle / num_workers on TorchDataLoaderArgs are ignored for preprocess.
    max_samples = data_loader.samples if data_loader is not None else -1

    # Load tokenizer
    tokenizer_factory = TokenizerRegistry.create(tokenizer)
    loaded_tokenizer = tokenizer_factory()

    # Load dataset
    dataset = DatasetDeserializerFactory.deserialize(
        config=data,
        processor_factory=tokenizer_factory,
        random_seed=random_seed,
    )
    # Setup column mapper
    column_mapper: GenerativeColumnMapper = PreprocessorRegistry.create(  # type: ignore[assignment]
        config=data_column_mapper
    )
    column_mapper.setup_data(
        datasets=[dataset],
    )

    # Extract column names from mapper
    prompt_column, prefix_column, output_column = _extract_column_names(column_mapper)

    # Create token samplers
    prompt_token_sampler, output_token_sampler = _create_token_samplers(
        strategy,
        random_seed,
    )

    # Process dataset
    dataset_iterator = iter(dataset)
    processed_prompts = []

    for row in dataset_iterator:
        processed_row = _process_single_row(
            row=row,
            prompt_column=prompt_column,
            prefix_column=prefix_column,
            prompt_token_sampler=prompt_token_sampler,
            output_token_sampler=output_token_sampler,
            tokenizer=loaded_tokenizer,
            strategy=strategy,
            dataset_iterator=dataset_iterator,
            output_column=output_column,
        )
        if processed_row is not None:
            processed_prompts.append(processed_row)
            if max_samples > 0 and len(processed_prompts) >= max_samples:
                break

    _finalize_processed_dataset(
        processed_prompts,
        output_path,
        push_to_hub,
        hub_dataset_id,
    )

def push_dataset_to_hub(
    hub_dataset_id: str | None,
    processed_dataset: Dataset,
) -> None:
    """
    Pushes the processed dataset to Hugging Face Hub using HF_TOKEN.

    :param hub_dataset_id: Identifier on the Hub to push to.
    :param processed_dataset: HuggingFace Dataset object.
    :raises ValueError: If hub_dataset_id or HF_TOKEN is not available.
    """

    hf_token = os.environ.get("HF_TOKEN")
    if not hub_dataset_id or not hf_token:
        raise ValueError(
            "hub_dataset_id and HF_TOKEN env var must be provided when push_to_hub"
            " is True"
        )
    processed_dataset.push_to_hub(hub_dataset_id, token=hf_token)