diff --git a/docs/guides/tool_calling.md b/docs/guides/tool_calling.md index ac5da1d13..2ea2e8c82 100644 --- a/docs/guides/tool_calling.md +++ b/docs/guides/tool_calling.md @@ -144,7 +144,9 @@ guidellm run \ The `tool_response_tokens_stdev`, `tool_response_tokens_min`, and `tool_response_tokens_max` fields work identically to the corresponding `prompt_tokens_*` / `output_tokens_*` variance parameters. -**2. Datasets with a tools column** -- datasets that already contain tool definitions (e.g. `madroid/glaive-function-calling-openai`) work directly. The column mapper auto-detects columns named `tools`, `functions`, or `tool_definitions`: +**2. Datasets with a tools column** -- datasets that already contain tool definitions (e.g. `madroid/glaive-function-calling-openai`) work directly. The column mapper auto-detects columns named `tools`, `functions`, or `tool_definitions`. + +**JSON-wrapped datasets** -- some HuggingFace datasets store all fields inside a single JSON string column (e.g. `madroid/glaive-function-calling-openai` has a `json` column containing `{"messages": [...], "tools": [...]}`). The column mapper automatically detects this pattern and unwraps the JSON to find the inner columns: ```bash guidellm run \ @@ -157,6 +159,19 @@ guidellm run \ --profile kind=constant,rate=1 ``` +**ShareGPT-format datasets** -- datasets using `from`/`value` keys instead of OpenAI's `role`/`content` (e.g. `NousResearch/hermes-function-calling-v1`) are also supported. The `tool_calling_message_extractor` automatically handles both formats and normalizes role aliases (e.g. `"human"` is treated as `"user"`): + +```bash +guidellm run \ + --backend kind=openai_http,target=http://localhost:8000 \ + --data kind=huggingface,source=NousResearch/hermes-function-calling-v1,load_kwargs.split=train \ + --data-column-mapper kind=generative_column_mapper,column_mappings.text_column=conversations,column_mappings.tools_column=tools \ + --data-preprocessor kind=tool_calling_message_extractor \ + --data-preprocessor kind=encode_media \ + --constraint kind=max_requests,count=50 \ + --profile kind=constant,rate=1 +``` + The `tool_calling_message_extractor` preprocessor must be explicitly enabled via `--data-preprocessor` (it is not included by default). It parses each row's `messages` array and extracts prompts, system messages, and tool results into the appropriate columns. If the dataset has no tool result messages, the placeholder (`{"status": "ok"}`) is used as a fallback. ## Tool Choice and Missing Tool-Call Behavior diff --git a/src/guidellm/data/preprocessors/mappers.py b/src/guidellm/data/preprocessors/mappers.py index 65a9801f1..974b8c052 100644 --- a/src/guidellm/data/preprocessors/mappers.py +++ b/src/guidellm/data/preprocessors/mappers.py @@ -1,10 +1,11 @@ from __future__ import annotations +import json import re from collections import defaultdict from typing import Any, ClassVar, Literal, TypeAlias, cast -from datasets import Dataset, IterableDataset +from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict from pydantic import Field from guidellm.data.preprocessors.preprocessor import ( @@ -29,6 +30,86 @@ DatasetColumnValue: TypeAlias = tuple[int, str] +def _unwrap_dataset_dict( + dataset: Dataset | IterableDataset | DatasetDict | IterableDatasetDict, +) -> Dataset | IterableDataset: + """Unwrap a DatasetDict/IterableDatasetDict into a single split. + + Prefers the ``"train"`` split if available, otherwise picks the first split. + Returns the input unchanged if it is already a single Dataset/IterableDataset. + + :param dataset: The dataset or dataset dict to unwrap. + :return: A single Dataset or IterableDataset. + """ + if isinstance(dataset, DatasetDict | IterableDatasetDict): + if "train" in dataset: + return dataset["train"] + return dataset[next(iter(dataset))] + return dataset + + +def _detect_json_wrapper( + dataset: Dataset | IterableDataset, dataset_columns: list[str] +) -> str | None: + """Check if a dataset has a single string column containing JSON dicts. + + :param dataset: The dataset to inspect. + :param dataset_columns: The column names present in the dataset. + :return: The wrapper column name if detected, or None. + """ + if len(dataset_columns) != 1: + return None + + candidate = dataset_columns[0] + sample = next(iter(dataset)) + value = sample[candidate] + if not isinstance(value, str): + return None + + try: + parsed = json.loads(value) + except (json.JSONDecodeError, TypeError): + return None + + if isinstance(parsed, dict) and parsed: + return candidate + return None + + +def _resolve_virtual_columns( + dataset: Dataset | IterableDataset, wrapper_column: str +) -> list[str]: + """Parse the first row's JSON wrapper and return its inner keys. + + :param dataset: The dataset to peek at. + :param wrapper_column: The name of the column containing the JSON string. + :return: List of inner key names from the parsed JSON dict. + """ + sample = next(iter(dataset)) + parsed = json.loads(sample[wrapper_column]) + return list(parsed.keys()) + + +def _extract_json_field( + row_data: dict[str, Any], wrapper_column: str, field: str +) -> Any: + """Parse a JSON wrapper column and extract a specific inner field. + + :param row_data: The raw row dict from the dataset. + :param wrapper_column: The column name containing the JSON string. + :param field: The key to extract from the parsed JSON dict. + :return: The value of the requested field, or None if not present. + """ + raw = row_data.get(wrapper_column) + if not isinstance(raw, str): + return None + try: + parsed = json.loads(raw) + except (json.JSONDecodeError, TypeError): + return None + return parsed.get(field) + + @DataPreprocessorArgs.register( [ "generative_column_mapper", @@ -117,6 +198,9 @@ class GenerativeColumnMapper(DataDependentPreprocessor): "tool_result", "tool_output", ], + "turn_type_column": [ + "turn_type", + ], "relative_timestamp_column": ["relative_timestamp"], "requeue_delay_column": ["requeue_delay"], } @@ -163,7 +247,7 @@ def datasets_mappings( cls, datasets: list[Dataset | IterableDataset], input_mappings: dict[str, str | list[str]] | None = None, - ) -> dict[DatasetColumnKey, list[DatasetColumnValue]]: + ) -> tuple[dict[DatasetColumnKey, list[DatasetColumnValue]], dict[int, str]]: """ Resolve column mappings across one or more datasets. @@ -171,23 +255,30 @@ def datasets_mappings( mapping names (or :attr:`defaults`) using regex patterns that account for pluralisation and turn suffixes (e.g. ``prompt-0``, ``prompt-1``). + When a dataset has no direct column matches but contains a single + JSON-string column, the inner keys of that JSON are used as virtual + column names and matching is retried. + :param datasets: The loaded datasets to inspect for column names. :param input_mappings: Optional explicit column mappings. When ``None``, :attr:`defaults` is used. Values may be a single name or a list of candidate names in priority order. - :return: A dict keyed by ``(column_type, turn_index)`` whose values are - lists of ``(dataset_index, column_name)`` pairs indicating where - each logical column can be found. Categories with no matching - columns are silently omitted from the result. + :return: A tuple of (mappings, json_wrappers) where mappings is a dict + keyed by ``(column_type, turn_index)`` whose values are lists of + ``(dataset_index, column_name)`` pairs, and json_wrappers is a dict + mapping dataset_index to the wrapper column name for datasets that + required JSON unwrapping. """ mappings: dict[DatasetColumnKey, list[DatasetColumnValue]] = defaultdict(list) + json_wrappers: dict[int, str] = {} input_map: dict[str, list[str]] = cls.defaults if input_mappings: input_map = { k: v if isinstance(v, list) else [v] for k, v in input_mappings.items() } - for index, dataset in enumerate(datasets): + for index, raw_dataset in enumerate(datasets): + dataset = _unwrap_dataset_dict(raw_dataset) dataset_name = ( dataset.info.dataset_name if dataset.info and dataset.info.dataset_name @@ -196,34 +287,72 @@ def datasets_mappings( dataset_columns = dataset.column_names or list(next(iter(dataset)).keys()) dataset_columns_str = "\n".join(dataset_columns) - for column_type, names in input_map.items(): - filtered_names = cls._filter_for_dataset( - names, str(index), str(dataset_name) - ) - if not filtered_names: - continue + matched = cls._match_columns( + index, input_map, dataset_name, dataset_columns_str + ) - column_pattern = cls.column_name_pattern.format( - name="|".join(re.escape(n) for n in filtered_names) - ) - # Find the first matching column name - base_match = re.search(column_pattern, dataset_columns_str, re.M | re.I) - if not base_match: - continue + # Fallback: if no matches found, try JSON unwrapping + if not matched: + wrapper = _detect_json_wrapper(dataset, dataset_columns) + if wrapper: + virtual_columns = _resolve_virtual_columns(dataset, wrapper) + virtual_columns_str = "\n".join(virtual_columns) + matched = cls._match_columns( + index, input_map, dataset_name, virtual_columns_str + ) + if matched: + json_wrappers[index] = wrapper - turn_pattern = cls.column_name_pattern.format( - name=base_match.group("match_name"), - ) - turn_columns = cls._extract_turn_columns( - turn_pattern, - dataset_columns_str, - ) + for key, values in matched.items(): + mappings[key].extend(values) - for turn, column_name in sorted(turn_columns): - column_type = cast("GenerativeDatasetColumnType", column_type) - mappings[(column_type, turn)].append((index, column_name)) + return mappings, json_wrappers - return mappings + @classmethod + def _match_columns( + cls, + index: int, + input_map: dict[str, list[str]], + dataset_name: str | int, + dataset_columns_str: str, + ) -> dict[DatasetColumnKey, list[DatasetColumnValue]]: + """Match input_map names against dataset columns using regex patterns. + + :param index: The dataset index in the multi-dataset list. + :param input_map: Mapping of column types to candidate column names. + :param dataset_name: Name or index of the dataset for filtering. + :param dataset_columns_str: Newline-joined string of column names. + :return: Dict of matched (column_type, turn) -> [(index, column_name)]. + """ + matched: dict[DatasetColumnKey, list[DatasetColumnValue]] = defaultdict(list) + + for column_type, names in input_map.items(): + filtered_names = cls._filter_for_dataset( + names, str(index), str(dataset_name) + ) + if not filtered_names: + continue + + column_pattern = cls.column_name_pattern.format( + name="|".join(re.escape(n) for n in filtered_names) + ) + base_match = re.search(column_pattern, dataset_columns_str, re.M | re.I) + if not base_match: + continue + + turn_pattern = cls.column_name_pattern.format( + name=base_match.group("match_name"), + ) + turn_columns = cls._extract_turn_columns( + turn_pattern, + dataset_columns_str, + ) + + for turn, column_name in sorted(turn_columns): + column_type = cast("GenerativeDatasetColumnType", column_type) + matched[(column_type, turn)].append((index, column_name)) + + return matched def __init__( self, @@ -233,6 +362,23 @@ def __init__( self.datasets_column_mappings: ( dict[DatasetColumnKey, list[DatasetColumnValue]] | None ) + self._json_wrappers: dict[int, str] = {} + + def _get_column_value( + self, items: list[dict[str, Any]], dataset_index: int, dataset_column: str + ) -> Any: + """Read a column value, unwrapping JSON if needed for this dataset. + + :param items: The per-dataset row items from the iterator. + :param dataset_index: Index of the dataset to read from. + :param dataset_column: Column name (possibly virtual) to extract. + :return: The column value from the row. + """ + row_data = items[dataset_index]["dataset"] + wrapper = self._json_wrappers.get(dataset_index) + if wrapper is not None: + return _extract_json_field(row_data, wrapper, dataset_column) + return row_data[dataset_column] def __call__(self, items: list[dict[str, Any]]) -> list[dict[str, list[Any]]]: if self.datasets_column_mappings is None: @@ -253,7 +399,7 @@ def __call__(self, items: list[dict[str, Any]]) -> list[dict[str, list[Any]]]: dataset_column, ) in column_mappings: mapped[turn][column_type].append( - items[dataset_index]["dataset"][dataset_column] + self._get_column_value(items, dataset_index, dataset_column) ) return [dict(m) for m in mapped if len(m) > 0] @@ -262,7 +408,7 @@ def setup_data( self, datasets: list[DatasetType], ): - self.datasets_column_mappings = self.datasets_mappings( + self.datasets_column_mappings, self._json_wrappers = self.datasets_mappings( datasets, self.input_mappings ) diff --git a/tests/unit/data/preprocessors/test_mappers.py b/tests/unit/data/preprocessors/test_mappers.py new file mode 100644 index 000000000..24ffbc2f5 --- /dev/null +++ b/tests/unit/data/preprocessors/test_mappers.py @@ -0,0 +1,237 @@ +""" +Unit tests for JSON unwrapping and DatasetDict handling in +guidellm.data.preprocessors.mappers. + +## WRITTEN BY AI ## +""" + +from __future__ import annotations + +import json + +import pytest +from datasets import Dataset, DatasetDict + +from guidellm.data.preprocessors.mappers import ( + GenerativeColumnMapper, + GenerativeColumnMapperArgs, + _detect_json_wrapper, + _extract_json_field, + _resolve_virtual_columns, + _unwrap_dataset_dict, +) + + +class TestDetectJsonWrapper: + """Tests for _detect_json_wrapper helper. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_detects_single_json_string_column(self): + """A dataset with one column containing JSON dicts is detected. + + ## WRITTEN BY AI ## + """ + data = {"json": [json.dumps({"messages": [], "tools": []})]} + ds = Dataset.from_dict(data) + result = _detect_json_wrapper(ds, ["json"]) + assert result == "json" + + @pytest.mark.sanity + def test_returns_none_for_multiple_columns(self): + """Datasets with more than one column are not treated as wrapped. + + ## WRITTEN BY AI ## + """ + data = {"col_a": ["hello"], "col_b": ["world"]} + ds = Dataset.from_dict(data) + result = _detect_json_wrapper(ds, ["col_a", "col_b"]) + assert result is None + + @pytest.mark.sanity + def test_returns_none_for_non_json_string(self): + """A single string column that isn't valid JSON returns None. + + ## WRITTEN BY AI ## + """ + data = {"text": ["this is plain text"]} + ds = Dataset.from_dict(data) + result = _detect_json_wrapper(ds, ["text"]) + assert result is None + + @pytest.mark.sanity + def test_returns_none_for_non_dict_json(self): + """A JSON column containing a list (not dict) returns None. + + ## WRITTEN BY AI ## + """ + data = {"json": [json.dumps([1, 2, 3])]} + ds = Dataset.from_dict(data) + result = _detect_json_wrapper(ds, ["json"]) + assert result is None + + +class TestResolveVirtualColumns: + """Tests for _resolve_virtual_columns helper. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_returns_inner_keys(self): + """Virtual columns are the keys of the parsed JSON dict. + + ## WRITTEN BY AI ## + """ + data = {"json": [json.dumps({"messages": [], "tools": [], "metadata": {}})]} + ds = Dataset.from_dict(data) + result = _resolve_virtual_columns(ds, "json") + assert set(result) == {"messages", "tools", "metadata"} + + +class TestExtractJsonField: + """Tests for _extract_json_field helper. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_extracts_field_from_json_string(self): + """Correctly parses JSON and returns the requested field. + + ## WRITTEN BY AI ## + """ + row = {"json": json.dumps({"messages": [{"role": "user", "content": "hi"}]})} + result = _extract_json_field(row, "json", "messages") + assert result == [{"role": "user", "content": "hi"}] + + @pytest.mark.sanity + def test_returns_none_for_missing_field(self): + """Returns None when the requested field is not in the JSON. + + ## WRITTEN BY AI ## + """ + row = {"json": json.dumps({"messages": []})} + result = _extract_json_field(row, "json", "tools") + assert result is None + + @pytest.mark.sanity + def test_returns_none_for_invalid_json(self): + """Returns None when the wrapper column contains invalid JSON. + + ## WRITTEN BY AI ## + """ + row = {"json": "not valid json{{{"} + result = _extract_json_field(row, "json", "messages") + assert result is None + + +class TestUnwrapDatasetDict: + """Tests for _unwrap_dataset_dict helper. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_returns_train_split_from_dict(self): + """Prefers the 'train' split when available. + + ## WRITTEN BY AI ## + """ + train_ds = Dataset.from_dict({"col": [1, 2, 3]}) + test_ds = Dataset.from_dict({"col": [4, 5]}) + dd = DatasetDict({"train": train_ds, "test": test_ds}) + result = _unwrap_dataset_dict(dd) + assert len(result) == 3 + + @pytest.mark.sanity + def test_returns_first_split_when_no_train(self): + """Falls back to first split if 'train' is not present. + + ## WRITTEN BY AI ## + """ + val_ds = Dataset.from_dict({"col": [10, 20]}) + dd = DatasetDict({"validation": val_ds}) + result = _unwrap_dataset_dict(dd) + assert len(result) == 2 + + @pytest.mark.sanity + def test_returns_dataset_unchanged(self): + """A plain Dataset passes through without modification. + + ## WRITTEN BY AI ## + """ + ds = Dataset.from_dict({"col": [1]}) + result = _unwrap_dataset_dict(ds) + assert result is ds + + +class TestColumnMapperJsonUnwrapping: + """Integration test for JSON unwrapping through the full column mapper flow. + + ## WRITTEN BY AI ## + """ + + @pytest.mark.smoke + def test_json_wrapped_dataset_maps_virtual_columns(self): + """Column mapper resolves virtual columns from a JSON-wrapped dataset. + + ## WRITTEN BY AI ## + """ + messages = [{"role": "user", "content": "hello"}] + tools = [{"type": "function", "function": {"name": "test_fn"}}] + row = json.dumps({"messages": messages, "tools": tools}) + ds = Dataset.from_dict({"json": [row, row]}) + + config = GenerativeColumnMapperArgs( + column_mappings={"text_column": "messages", "tools_column": "tools"} + ) + mapper = GenerativeColumnMapper(config) + mapper.setup_data(datasets=[ds]) + + assert mapper._json_wrappers == {0: "json"} + + # Test __call__ extracts the right values + items = [{"dataset": {"json": row}}] + result = mapper(items) + assert len(result) == 1 + assert result[0]["text_column"] == [messages] + assert result[0]["tools_column"] == [tools] + + @pytest.mark.sanity + def test_normal_dataset_no_json_wrapper(self): + """Datasets with proper columns don't trigger JSON unwrapping. + + ## WRITTEN BY AI ## + """ + ds = Dataset.from_dict( + { + "prompt": ["hello", "world"], + "tools": [None, None], + } + ) + + config = GenerativeColumnMapperArgs( + column_mappings={"text_column": "prompt", "tools_column": "tools"} + ) + mapper = GenerativeColumnMapper(config) + mapper.setup_data(datasets=[ds]) + + assert mapper._json_wrappers == {} + + @pytest.mark.sanity + def test_dataset_dict_unwrapped_before_mapping(self): + """A DatasetDict is unwrapped to a single split before column mapping. + + ## WRITTEN BY AI ## + """ + ds = Dataset.from_dict({"prompt": ["hello"]}) + dd = DatasetDict({"train": ds}) + + config = GenerativeColumnMapperArgs(column_mappings={"text_column": "prompt"}) + mapper = GenerativeColumnMapper(config) + mapper.setup_data(datasets=[dd]) + + assert mapper.datasets_column_mappings is not None