Re-introduce column-mapper arguments

Signed-off-by: Samuel Monson <[email protected]>

Author: sjmonson

src/guidellm/__main__.py

@@ -216,6 +216,12 @@ def benchmark():
         "and dynamically generates more."
     ),
 )
+@click.option(
+    "--data-column-mapper",
+    default=BenchmarkGenerativeTextArgs.get_default("data_column_mapper"),
+    callback=cli_tools.parse_json,
+    help="JSON string of column mappings to apply to the dataset.",
+)
 @click.option(
     "--data-preprocessors",
     default=BenchmarkGenerativeTextArgs.get_default("data_preprocessors"),

src/guidellm/benchmark/entrypoints.py

@@ -19,6 +19,7 @@
 from transformers import PreTrainedTokenizerBase
 from typing_extensions import TypeAliasType
 
+from guidellm import settings
 from guidellm.backends import Backend, BackendType
 from guidellm.benchmark.benchmarker import Benchmarker
 from guidellm.benchmark.outputs import (
@@ -236,6 +237,11 @@ async def resolve_request_loader(
     data_samples: int,
     processor: ProcessorInputT | None,
     processor_args: dict[str, Any] | None,
+    data_column_mapper: (
+        DatasetPreprocessor
+        | dict[str, str | list[str]]
+        | Literal["generative_column_mapper"]
+    ),
     data_preprocessors: list[DatasetPreprocessor | dict[str, str | list[str]] | str],
     data_finalizer: (DatasetFinalizer | dict[str, Any] | str),
     data_collator: Callable | Literal["generative"] | None,
@@ -279,13 +285,20 @@ async def resolve_request_loader(
         else None
     )
 
+    # If no type is specified for the data column mapper, load default
+    if isinstance(data_column_mapper, dict) and "type" not in data_column_mapper:
+        data_column_mapper = {
+            "type": settings.dataset.default_column_mapper,
+            **data_column_mapper,
+        }
+
     preprocessors_list: list[DatasetPreprocessor] = [
         resolve_item_from_registry(
             DatasetPreprocessor,  # type: ignore [type-abstract]
             PreprocessorRegistry,
             preprocessor,
         )
-        for preprocessor in data_preprocessors
+        for preprocessor in ([data_column_mapper] + data_preprocessors)
     ]
 
     finalizer_instance = resolve_item_from_registry(

src/guidellm/benchmark/schemas/generative/entrypoints.py

@@ -202,11 +202,17 @@ def get_default(cls: type[BenchmarkGenerativeTextArgs], field: str) -> Any:
     data_samples: int = Field(
         default=-1, description="Number of samples to use from datasets (-1 for all)"
     )
-    # TODO: Make it easy to cutomize preprocessors without editing the full list
+    data_column_mapper: (
+        DatasetPreprocessor
+        | dict[str, str | list[str]]
+        | Literal["generative_column_mapper"]
+    ) = Field(
+        default_factory=lambda: settings.dataset.default_column_mapper,
+        description="Column mapping preprocessor for dataset fields",
+    )
     data_preprocessors: list[DatasetPreprocessor | dict[str, str | list[str]] | str] = (
         Field(
             default_factory=lambda: [  # type: ignore [arg-type]
-                "generative_column_mapper",
                 "encode_media",
             ],
             description="List of dataset preprocessors to apply in order",
@@ -327,18 +333,23 @@ def serialize_data_collator(
         """Serialize data_collator to string or None."""
         return data_collator if isinstance(data_collator, str) else None
 
+    @field_serializer("data_column_mapper")
+    def serialize_preprocessor(
+        self,
+        data_preprocessor: (DatasetPreprocessor | dict[str, str | list[str]] | str),
+    ) -> dict | str:
+        """Serialize a preprocessor to dict or string."""
+        return data_preprocessor if isinstance(data_preprocessor, dict | str) else {}
+
     @field_serializer("data_preprocessors")
-    def serialize_data_column_mapper(
+    def serialize_preprocessors(
         self,
         data_preprocessors: list[
             DatasetPreprocessor | dict[str, str | list[str]] | str
         ],
     ) -> list[dict | str]:
-        """Serialize data_column_mapper to dict or string."""
-        return [
-            (preprocessor if isinstance(preprocessor, dict | str) else {})
-            for preprocessor in data_preprocessors
-        ]
+        """Serialize each preprocessor to dict or string."""
+        return [self.serialize_preprocessor(p) for p in data_preprocessors]
 
     @field_serializer("data_finalizer")
     def serialize_data_request_formatter(

src/guidellm/settings.py

@@ -73,6 +73,9 @@ class DatasetSettings(BaseModel):
     preferred_data_splits: list[str] = Field(
         default_factory=lambda: ["test", "tst", "validation", "val", "train"]
     )
+    default_column_mapper: Literal["generative_column_mapper"] = (
+        "generative_column_mapper"
+    )
 
 
 class OpenAISettings(BaseModel):