Reenable and improve preprocess dataset

docs/guides/datasets.md

@@ -212,3 +212,237 @@ benchmark_generative_text(data=data, ...)
 - For lists of dictionaries, all items must have the same keys.
 - For lists of items, all elements must be of the same type.
 - A processor/tokenizer is only required if `GUIDELLM__PREFERRED_PROMPT_TOKENS_SOURCE="local"` or `GUIDELLM__PREFERRED_OUTPUT_TOKENS_SOURCE="local"` is set in the environment. In this case, the processor/tokenizer must be specified using the `--processor` argument. If not set, the processor/tokenizer will be set to the model passed in or retrieved from the server.
+
+## Preprocessing Datasets
+
+GuideLLM provides a preprocessing command that allows you to process datasets to have specific prompt and output token sizes. This is particularly useful when you need to standardize your dataset for benchmarking or when your dataset has prompts that don't match your target token requirements.
+
+The preprocessing command can:
+
+- Resize prompts to target token lengths
+- Handle prompts that are shorter or longer than the target length using various strategies
+- Map columns from your dataset to GuideLLM's expected column names
+- Generate output token counts based on your configuration
+- Save the processed dataset in various formats
+
+### Basic Usage
+
+```bash
+guidellm preprocess dataset \
+    <DATA> \
+    <OUTPUT_PATH> \
+    --processor <PROCESSOR> \
+    --config <CONFIG>
+```
+
+### Required Arguments
+
+| Argument      | Description                                                                                                                                    |
+| ------------- | ---------------------------------------------------------------------------------------------------------------------------------------------- |
+| `DATA`        | Path to the input dataset or Hugging Face dataset ID. Supports all dataset formats documented in the [Dataset Configurations](../datasets.md). |
+| `OUTPUT_PATH` | Path to save the processed dataset, including file suffix (e.g., `processed_dataset.jsonl`, `output.csv`).                                     |
+| `--processor` | **Required.** Processor or tokenizer name/path for calculating token counts. Can be a Hugging Face model ID or local path.                     |
+| `--config`    | **Required.** Configuration specifying target token sizes. Can be a JSON string, key=value pairs, or file path (.json, .yaml, .yml, .config).  |
+
+### Example
+
+```bash
+guidellm preprocess dataset \
+    "path/to/input_dataset.jsonl" \
+    "path/to/processed_dataset.jsonl" \
+    --processor "gpt2" \
+    --config "prompt_tokens=512,output_tokens=256,prefix_tokens_max=100"
+```
+
+### Configuration and Processor Options
+
+The `--config` parameter accepts a `PreprocessDatasetConfig` as a JSON string, key=value pairs, or a configuration file path (.json, .yaml, .yml, .config). This configuration is similar to the synthetic data configuration but includes additional fields specific to preprocessing.
+
+**PreprocessDatasetConfig Options:**
+
+- `prompt_tokens`: Average number of tokens in prompts. If nothing else is specified, all prompts will be resized to this number of tokens.
+- `prompt_tokens_stdev`: Standard deviation for prompt tokens. If not supplied and min/max are not specified, no deviation is applied. If not supplied and min/max are specified, a uniform distribution is used.
+- `prompt_tokens_min`: Minimum number of tokens in prompts. If unset and `prompt_tokens_stdev` is set, the minimum is 1.
+- `prompt_tokens_max`: Maximum number of tokens in prompts. If unset and `prompt_tokens_stdev` is set, the maximum is 5 times the standard deviation.
+- `output_tokens`: Average number of tokens in outputs. If nothing else is specified, all outputs will have this number of tokens.
+- `output_tokens_stdev`: Standard deviation for output tokens. If not supplied and min/max are not specified, no deviation is applied. If not supplied and min/max are specified, a uniform distribution is used.
+- `output_tokens_min`: Minimum number of tokens in outputs. If unset and `output_tokens_stdev` is set, the minimum is 1.
+- `output_tokens_max`: Maximum number of tokens in outputs. If unset and `output_tokens_stdev` is set, the maximum is 5 times the standard deviation.
+- `prefix_tokens_max`: Maximum number of prefix tokens to keep. If set, prefixes will be trimmed to this maximum length. If not set, prefixes are kept as-is (unless `--include-prefix-in-token-count` is used, which disables prefix trimming).
+
+**Example configurations:**
+
+```bash
+# Using key=value pairs
+--config "prompt_tokens=512,output_tokens=256,prefix_tokens_max=100"
+
+# Using JSON string
+--config '{"prompt_tokens": 512, "output_tokens": 256, "prefix_tokens_max": 100}'
+
+# Using a configuration file
+--config "path/to/config.json"
+```
+
+The `--processor` argument specifies the tokenizer to use for calculating token counts. This is required because the preprocessing command needs to tokenize prompts to ensure they match the target token sizes. For information about using processors, including Hugging Face model IDs, local paths, and processor arguments, see the [Data Arguments Overview](../datasets.md#data-arguments-overview) section.
+
+### Column Mapping
+
+When your dataset uses non-standard column names, you can use `--data-column-mapper` to map your columns to GuideLLM's expected column names. This is particularly useful when:
+
+1. **Your dataset uses different column names** (e.g., `question` instead of `prompt`, `instruction` instead of `text_column`)
+2. **You have multiple datasets** and need to specify which dataset's columns to use
+3. **Your dataset has system prompts or prefixes** in a separate column
+
+**Column mapping format:** The `--data-column-mapper` accepts a JSON string mapping column types to column names:
+
+```json
+{
+  "text_column": "question",
+  "prefix_column": "system_prompt",
+  "prompt_tokens_count_column": "input_tokens",
+  "output_tokens_count_column": "completion_tokens"
+}
+```
+
+**Supported column types:**
+
+- `text_column`: The main prompt text (defaults: `prompt`, `instruction`, `question`, `input`, `context`, `content`, `text`)
+- `prefix_column`: System prompt or prefix (defaults: `system_prompt`, `system`, `prefix`)
+- `prompt_tokens_count_column`: Column containing prompt token counts (defaults: `prompt_tokens_count`, `input_tokens_count`)
+- `output_tokens_count_column`: Column containing output token counts (defaults: `output_tokens_count`, `completion_tokens_count`)
+- `image_column`: Image data column
+- `video_column`: Video data column
+- `audio_column`: Audio data column
+
+**Example: Mapping custom column names**
+
+If your dataset has a CSV file with columns `user_query` and `system_message`:
+
+```csv
+user_query,system_message
+"What is AI?","You are a helpful assistant."
+"How does ML work?","You are a technical expert."
+```
+
+You would use:
+
+```bash
+guidellm preprocess dataset \
+    "dataset.csv" \
+    "processed.jsonl" \
+    --processor "gpt2" \
+    --config "prompt_tokens=512,output_tokens=256" \
+    --data-column-mapper '{"text_column": "user_query", "prefix_column": "system_message"}'
+```
+
+**Example: Multiple datasets**
+
+If you're working with multiple datasets and need to specify which dataset's columns to use, you can use the format `<dataset_index>.<column_name>` or `<dataset_name>.<column_name>`:
+
+```bash
+--data-column-mapper '{"text_column": "0.prompt", "prefix_column": "1.system"}'
+```
+
+### Handling Short Prompts
+
+When prompts are shorter than the target token length, you can specify how to handle them using `--short-prompt-strategy`:
+
+| Strategy      | Description                                                                    |
+| ------------- | ------------------------------------------------------------------------------ |
+| `ignore`      | Skip prompts that are shorter than the target length (default)                 |
+| `concatenate` | Concatenate multiple short prompts together until the target length is reached |
+| `pad`         | Pad short prompts with a specified character to reach the target length        |
+| `error`       | Raise an error if a prompt is shorter than the target length                   |
+
+**Example: Concatenating short prompts**
+
+```bash
+guidellm preprocess dataset \
+    "dataset.jsonl" \
+    "processed.jsonl" \
+    --processor "gpt2" \
+    --config "prompt_tokens=512,output_tokens=256" \
+    --short-prompt-strategy "concatenate" \
+    --concat-delimiter "\n\n"
+```
+
+**Example: Padding short prompts**
+
+```bash
+guidellm preprocess dataset \
+    "dataset.jsonl" \
+    "processed.jsonl" \
+    --processor "gpt2" \
+    --config "prompt_tokens=512,output_tokens=256" \
+    --short-prompt-strategy "pad" \
+    --pad-char " "
+```
+
+### Additional Options
+
+| Option                            | Description                                                                                                                             |
+| --------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------- |
+| `--data-args <JSON>`              | JSON string of arguments to pass to dataset loading. See [Data Arguments Overview](../datasets.md#data-arguments-overview) for details. |
+| `--include-prefix-in-token-count` | Include prefix tokens in prompt token count calculation (flag). When enabled, prefix trimming is disabled and the prefix is kept as-is. |
+| `--random-seed <NUMBER>`          | Random seed for reproducible token sampling (default: 42).                                                                              |
+| `--push-to-hub`                   | Push the processed dataset to Hugging Face Hub (flag).                                                                                  |
+| `--hub-dataset-id <ID>`           | Hugging Face Hub dataset ID for upload (required if `--push-to-hub` is set).                                                            |
+
+### Complete Examples
+
+**Example 1: Basic preprocessing with custom column names**
+
+```bash
+guidellm preprocess dataset \
+    "my_dataset.csv" \
+    "processed_dataset.jsonl" \
+    --processor "gpt2" \
+    --config "prompt_tokens=512,output_tokens=256" \
+    --data-column-mapper '{"text_column": "user_question", "prefix_column": "system_instruction"}'
+```
+
+**Example 2: Preprocessing with distribution and short prompt handling**
+
+```bash
+guidellm preprocess dataset \
+    "dataset.jsonl" \
+    "processed.jsonl" \
+    --processor "gpt2" \
+    --config "prompt_tokens=512,prompt_tokens_stdev=50,output_tokens=256,output_tokens_stdev=25" \
+    --short-prompt-strategy "concatenate" \
+    --concat-delimiter "\n\n" \
+    --random-seed 123
+```
+
+**Example 3: Preprocessing with processor arguments and prefix token limits**
+
+```bash
+guidellm preprocess dataset \
+    "dataset.jsonl" \
+    "processed.jsonl" \
+    --processor "gpt2" \
+    --processor-args '{"use_fast": false}' \
+    --config "prompt_tokens=512,output_tokens=256,prefix_tokens_max=100" \
+    --include-prefix-in-token-count
+```
+
+**Example 4: Preprocessing and uploading to Hugging Face Hub**
+
+```bash
+guidellm preprocess dataset \
+    "my_dataset.jsonl" \
+    "processed.jsonl" \
+    --processor "gpt2" \
+    --config "prompt_tokens=512,output_tokens=256" \
+    --push-to-hub \
+    --hub-dataset-id "username/processed-dataset"
+```
+
+### Notes
+
+- The `--config` parameter accepts a `PreprocessDatasetConfig` which includes all token count fields (prompt_tokens, output_tokens, etc.) plus `prefix_tokens_max` for controlling prefix length. See the [Configuration and Processor Options](#configuration-and-processor-options) section above for all available parameters.
+- The processor/tokenizer is required because the preprocessing command needs to tokenize prompts to ensure they match target token sizes. See the [Data Arguments Overview](../datasets.md#data-arguments-overview) for processor usage details.
+- Column mappings are only needed when your dataset uses non-standard column names. GuideLLM will automatically try common column names if no mapping is provided.
+- When using `--short-prompt-strategy concatenate`, ensure your dataset has enough samples to concatenate, or some prompts may be skipped.
+- The output format is determined by the file extension of `OUTPUT_PATH` (e.g., `.jsonl`, `.csv`, `.parquet`).
+- The prefix handling only trims prefixes. It doesn't expand them. Use `prefix_tokens_max` in the config to set a maximum prefix length, which will trim prefixes that exceed this limit.

src/guidellm/__main__.py

@@ -30,6 +30,8 @@
 import click
 from pydantic import ValidationError
 
+from guidellm.data import ShortPromptStrategy, process_dataset
+
 try:
     import uvloop
 except ImportError:
@@ -486,6 +488,142 @@ def preprocess():
     """Dataset preprocessing utilities."""
 
 
+@preprocess.command(
+    "dataset",
+    help=(
+        "Process a dataset to have specific prompt and output token sizes. "
+        "Supports multiple strategies for handling prompts and optional "
+        "Hugging Face Hub upload.\n\n"
+        "DATA: Path to the input dataset or dataset ID.\n\n"
+        "OUTPUT_PATH: Path to save the processed dataset, including file suffix."
+    ),
+    context_settings={"auto_envvar_prefix": "GUIDELLM"},
+)
+@click.argument(
+    "data",
+    type=str,
+    required=True,
+)
+@click.argument(
+    "output_path",
+    type=click.Path(file_okay=True, dir_okay=False, writable=True, resolve_path=True),
+    required=True,
+)
+@click.option(
+    "--processor",
+    type=str,
+    required=True,
+    help="Processor or tokenizer name for calculating token counts.",
+)
+@click.option(
+    "--config",
+    type=str,
+    required=True,
+    help=(
+        "PreprocessDatasetConfig as JSON string, key=value pairs, "
+        "or file path (.json, .yaml, .yml, .config). "
+        "Example: 'prompt_tokens=100,output_tokens=50,prefix_tokens_max=10'"
+        " or '{\"prompt_tokens\": 100, \"output_tokens\": 50, "
+        "\"prefix_tokens_max\": 10}'"
+    ),
+)
+@click.option(
+    "--processor-args",
+    default=None,
+    callback=cli_tools.parse_json,
+    help="JSON string of arguments to pass to the processor constructor.",
+)
+@click.option(
+    "--data-args",
+    callback=cli_tools.parse_json,
+    help="JSON string of arguments to pass to dataset creation.",
+)
+@click.option(
+    "--data-column-mapper",
+    default=None,
+    callback=cli_tools.parse_json,
+    help="JSON string of column mappings to apply to the dataset.",
+)
+@click.option(
+    "--short-prompt-strategy",
+    type=click.Choice([s.value for s in ShortPromptStrategy]),
+    default=ShortPromptStrategy.IGNORE.value,
+    show_default=True,
+    help="Strategy for handling prompts shorter than target length.",
+)
+@click.option(
+    "--pad-char",
+    type=str,
+    default="",
+    callback=decode_escaped_str,
+    help="Character to pad short prompts with when using 'pad' strategy.",
+)
+@click.option(
+    "--concat-delimiter",
+    type=str,
+    default="",
+    help=(
+        "Delimiter for concatenating short prompts (used with 'concatenate' strategy)."
+    ),
+)
+@click.option(
+    "--include-prefix-in-token-count",
+    is_flag=True,
+    default=False,
+    help="Include prefix tokens in prompt token count calculation.",
+)
+@click.option(
+    "--push-to-hub",
+    is_flag=True,
+    help="Push the processed dataset to Hugging Face Hub.",
+)
+@click.option(
+    "--hub-dataset-id",
+    type=str,
+    default=None,
+    help=("Hugging Face Hub dataset ID for upload (required if --push-to-hub is set)."),
+)
+@click.option(
+    "--random-seed",
+    type=int,
+    default=42,
+    show_default=True,
+    help="Random seed for reproducible token sampling.",
+)
+def dataset(
+    data,
+    output_path,
+    processor,
+    config,
+    processor_args,
+    data_args,
+    data_column_mapper,
+    short_prompt_strategy,
+    pad_char,
+    concat_delimiter,
+    include_prefix_in_token_count,
+    push_to_hub,
+    hub_dataset_id,
+    random_seed,
+):
+    process_dataset(
+        data=data,
+        output_path=output_path,
+        processor=processor,
+        config=config,
+        processor_args=processor_args,
+        data_args=data_args,
+        data_column_mapper=data_column_mapper,
+        short_prompt_strategy=short_prompt_strategy,
+        pad_char=pad_char,
+        concat_delimiter=concat_delimiter,
+        include_prefix_in_token_count=include_prefix_in_token_count,
+        push_to_hub=push_to_hub,
+        hub_dataset_id=hub_dataset_id,
+        random_seed=random_seed,
+    )
+
+
 @cli.command(
     "mock-server",
     help=(