Enable inherent casting of the datasets.

dushyantbehl · dushyantbehl · commit 2504e0866f05 · 2025-05-12T18:23:50.000+05:30
Handles jumbled columns and mismatching dtypes

Signed-off-by: Dushyant Behl &lt;dushyantbehl@in.ibm.com&gt;
diff --git a/pyproject.toml b/pyproject.toml
@@ -37,7 +37,7 @@ dependencies = [
 "trl>=0.13,<0.18",
 "peft>=0.8.0,<0.14",
 "protobuf>=5.28.0,<6.0.0",
-"datasets>=2.15.0,<4.0",
+"datasets>=3.5.0,<4.0",
 "simpleeval>=0.9.13,<2.0",
 "pillow>=11.0.0,<12.0",
 ]
diff --git a/tests/artifacts/testdata/__init__.py b/tests/artifacts/testdata/__init__.py
@@ -74,6 +74,12 @@
 CHAT_DATA_MULTI_TURN_GRANITE_3_1B = os.path.join(
     JSONL_DATA_DIR, "multi_turn_chat_granite_instruct.jsonl"
 )
+CHAT_DATASET_LARGELIST = os.path.join(
+    PARQUET_DATA_DIR, "chat_dataset_tokenized_largelist.parquet"
+)
+CHAT_DATASET_SEQUENCE = os.path.join(
+    PARQUET_DATA_DIR, "chat_dataset_tokenized_sequence.parquet"
+)
 IMAGE_DATASET = os.path.join(JSONL_DATA_DIR, "image_dataset.jsonl")
 EMPTY_DATA = os.path.join(JSON_DATA_DIR, "empty_data.json")
 MALFORMATTED_DATA = os.path.join(JSON_DATA_DIR, "malformatted_data.json")
diff --git a/tests/artifacts/testdata/parquet/chat_dataset_tokenized_largelist.parquet b/tests/artifacts/testdata/parquet/chat_dataset_tokenized_largelist.parquet
diff --git a/tests/artifacts/testdata/parquet/chat_dataset_tokenized_sequence.parquet b/tests/artifacts/testdata/parquet/chat_dataset_tokenized_sequence.parquet
diff --git a/tests/test_sft_trainer.py b/tests/test_sft_trainer.py
@@ -51,12 +51,15 @@
     DATA_CONFIG_YAML_STREAMING_INPUT_OUTPUT,
     DATA_CONFIG_YAML_STREAMING_PRETOKENIZED,
     GRANITE_3_1_B_CHAT_TEMPLATE,
+    DATA_CONFIG_PRETOKENIZE_JSON_DATA_YAML
 )
 from tests.artifacts.testdata import (
     CHAT_DATA_MULTI_TURN,
     CHAT_DATA_MULTI_TURN_CONVERSATIONS,
     CHAT_DATA_MULTI_TURN_GRANITE_3_1B,
     CHAT_DATA_SINGLE_TURN,
+    CHAT_DATASET_LARGELIST,
+    CHAT_DATASET_SEQUENCE,
     CUSTOM_TOKENIZER_TINYLLAMA,
     EMPTY_DATA,
     MALFORMATTED_DATA,
@@ -820,7 +823,6 @@ def test_run_causallm_ft_pretokenized(dataset_path, packing):
         assert len(output_inference) > 0
         assert "### Text: @NortonSupport Thanks much.\n\n### Label:" in output_inference
 
-
 @pytest.mark.parametrize(
     "datafiles, datasetconfigname",
     [
@@ -836,9 +838,13 @@ def test_run_causallm_ft_pretokenized(dataset_path, packing):
             [TWITTER_COMPLAINTS_TOKENIZED_JSON],
             DATA_CONFIG_YAML_STREAMING_PRETOKENIZED,
         ),
+        (
+            [CHAT_DATASET_LARGELIST, CHAT_DATASET_SEQUENCE],
+            DATA_CONFIG_PRETOKENIZE_JSON_DATA_YAML,
+        ),
     ],
 )
-def test_run_causallm_ft_and_inference_streaming(datasetconfigname, datafiles):
+def test_run_causallm_ft_and_inference(datasetconfigname, datafiles):
     """Check if we can finetune causallm models using multiple datasets with multiple files"""
     with tempfile.TemporaryDirectory() as tempdir:
         data_formatting_args = copy.deepcopy(DATA_ARGS)
diff --git a/tests/utils/test_config_utils.py b/tests/utils/test_config_utils.py
@@ -17,7 +17,6 @@
 
 # Standard
 import base64
-import logging
 import os
 import pickle
 
@@ -31,7 +30,8 @@
 
 # Local
 from tuning.config import peft_config
-from tuning.utils import config_utils, utils
+from tuning.data import utils
+from tuning.utils import config_utils
 
 
 def test_get_hf_peft_config_returns_None_for_tuning_config_None():
@@ -236,7 +236,7 @@ def test_get_json_config_can_load_from_envvar():
     assert job_config["model_name_or_path"] == "foobar"
 
 
-def test_validate_datasets_logs_warnings_on_mismatch(caplog):
+def test_validate_datasets_throws_error_on_mismatch():
     """Test that `validate_mergeable_datasets` logs warnings when
     datasets have different columns or dtypes."""
     # Create a reference dataset with columns col1:int64 and col2:string
@@ -251,12 +251,5 @@ def test_validate_datasets_logs_warnings_on_mismatch(caplog):
         features=Features({"col1": Value("float64"), "col3": Value("string")}),
     )
 
-    with caplog.at_level(logging.WARNING):
-        utils.validate_mergeable_datasets([ds1, ds2])
-
-    assert (
-        "different columns" in caplog.text
-    ), "Expected a warning about differing columns."
-    assert (
-        "expected int64" in caplog.text
-    ), "Expected a warning about mismatching column dtypes."
+    with pytest.raises(ValueError):
+        utils._validate_mergeable_datasets([ds1, ds2])
diff --git a/tuning/data/collators.py b/tuning/data/collators.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 # Local
-from tuning.utils.utils import try_convert_bytes_dict_to_pil
+from tuning.data.utils import try_convert_bytes_dict_to_pil
 
 
 class VisionDataCollator:
diff --git a/tuning/data/data_config.py b/tuning/data/data_config.py
@@ -20,7 +20,7 @@
 import os
 
 # Local
-from tuning.utils.utils import load_yaml_or_json
+from tuning.data.utils import load_yaml_or_json
 
 logger = logging.getLogger(__name__)
 
diff --git a/tuning/data/data_handlers.py b/tuning/data/data_handlers.py
@@ -34,8 +34,8 @@
 import torch
 
 # Local
+from tuning.data.utils import try_convert_bytes_dict_to_pil, try_convert_image_to_rgb
 from tuning.utils.config_utils import process_jinja_placeholders
-from tuning.utils.utils import try_convert_bytes_dict_to_pil, try_convert_image_to_rgb
 
 logger = logging.getLogger(__name__)
 
diff --git a/tuning/data/data_processors.py b/tuning/data/data_processors.py
@@ -31,10 +31,10 @@
     DataHandler,
     DataHandlerType,
 )
-from tuning.utils.utils import (
+from tuning.data.utils import (
     get_loader_for_filepath,
+    maybe_align_datasets,
     resolve_iterable_dataset_features,
-    validate_mergeable_datasets,
 )
 
 logger = logging.getLogger(__name__)
@@ -223,31 +223,30 @@ def _try_load_dataset(dataset_path, dataset_builder, streaming):
 
         for data_path in data_paths:
             dataset = _try_load_dataset(data_path, builder, streaming)
-            if isinstance(dataset, IterableDataset):
-                dataset = resolve_iterable_dataset_features(dataset)
             all_datasets.append(dataset)
 
-        # Logs warning if datasets have different columns
-        validate_mergeable_datasets(all_datasets)
-
         # Concatenate all datasets
         try:
             if len(all_datasets) == 1:
                 return all_datasets[0]
-
+            maybe_align_datasets(all_datasets)
             raw_datasets = datasets.concatenate_datasets(all_datasets)
             logger.info(
-                "Datasets concatenated from %s .Concatenated dataset columns: %s",
+                "Datasets %s concatenated. Final column features: %s",
                 datasetconfig.name,
-                list(raw_datasets.features.keys()),
+                str(list(raw_datasets.features)),
             )
-            return raw_datasets
-
         except Exception as e:
             raise ValueError(
                 f"An error occurred while concatenating datasets from {datasetconfig.name}: {e}"
             ) from e
 
+        # Need to resolve dataset features because data handlers use columns.
+        if isinstance(raw_datasets, IterableDataset):
+            raw_datasets = resolve_iterable_dataset_features(raw_datasets)
+
+        return raw_datasets
+
     def __execute_rename_data_handler(self, raw_datasets, handler, **kwargs):
         """
         Rename columns in the dataset using the provided column mapping.
@@ -456,9 +455,6 @@ def _process_dataset_configs(
             raw_dataset = self.load_dataset(
                 d, self.processor_config.streaming, splitName
             )
-            if isinstance(raw_dataset, IterableDataset):
-                raw_dataset = resolve_iterable_dataset_features(raw_dataset)
-
             logger.info("Loaded raw dataset : %s", str(raw_dataset))
 
             if isinstance(raw_dataset, IterableDataset):
@@ -493,6 +489,9 @@ def _process_dataset_configs(
                 else:
                     final_datasets[k].append(v)
 
+        # Ensure again datasets are aligned before interleaving or concatenating
+        maybe_align_datasets(final_datasets)
+
         if sample_datasets:
             strategy = self.processor_config.sampling_stopping_strategy
             seed = self.processor_config.sampling_seed
@@ -517,6 +516,8 @@ def _process_dataset_configs(
                 )
 
         train_dataset = final_datasets.get("train", None)
+
+        # Just a failsafe in case this is required later.
         if isinstance(train_dataset, IterableDataset):
             train_dataset = resolve_iterable_dataset_features(train_dataset)
 
diff --git a/tuning/data/utils.py b/tuning/data/utils.py
@@ -19,7 +19,7 @@
 import os
 
 # Third Party
-from datasets import IterableDataset
+from datasets import DatasetDict, IterableDataset, IterableDatasetDict
 from PIL import Image
 import yaml
 
@@ -70,39 +70,73 @@ def resolve_iterable_dataset_features(data: IterableDataset):
     return data
 
 
-def validate_mergeable_datasets(datasets):
+def __get_dataset_features(d, default_split="train"):
+    return (
+        d[default_split].features
+        if isinstance(d, (DatasetDict or IterableDatasetDict))
+        else d.features
+    )
+
+
+def _maybe_cast_columns(datasets, default_split="train"):
+    """
+    Given list of datasets, try casting datasets to same features.
+    Assumes that the datasets are aligned in terms of columns which
+    could be ensure by calling validate_mergeable_datasets
+    """
+    if len(datasets) <= 1:
+        return
+
+    # pick the first dataset as the reference
+    features = __get_dataset_features(datasets[0], default_split)
+
+    # Cast remaining datasets according to this
+    for i in range(1, len(datasets)):
+        datasets[i] = datasets[i].cast(features)
+
+
+def _validate_mergeable_datasets(datasets, default_split="train"):
     """Given list of datasets, validate if all datasets have same type and number of columns."""
-    if len(datasets) > 1:
-        ref_columns = datasets[0].features
-        ref_column_names = list(ref_columns.keys())
-        ref_column_types = {col: feat.dtype for col, feat in ref_columns.items()}
-
-        # Check all other datasets
-        for i, ds in enumerate(datasets[1:], start=2):
-            ds_column_names = list(ds.features.keys())
-            ds_column_types = {col: feat.dtype for col, feat in ds.features.items()}
-
-            # Check same set of columns
-            if set(ds_column_names) != set(ref_column_names):
-                logger.warning(
-                    "Dataset %d has different columns: %s. Columns in Dataset 1: %s",
-                    i,
-                    ds_column_names,
-                    ref_column_names,
-                )
+    if len(datasets) <= 1:
+        return
+
+    ref_columns = __get_dataset_features(datasets[0], default_split)
+    ref_column_names = list(ref_columns.keys())
+
+    # Check all other datasets
+    mismatching_ds = []
+    for _, ds in enumerate(datasets[1:], start=1):
+        ds_features = __get_dataset_features(ds, default_split)
+        ds_column_names = list(ds_features.keys())
+
+        # Check same set of columns
+        if set(ds_column_names) != set(ref_column_names):
+            mismatching_ds.append([ds])
 
-            # Check column data types
-            for col in ref_column_names:
-                if (col in ds_column_types) and (
-                    ds_column_types[col] != ref_column_types[col]
-                ):
-                    logger.warning(
-                        "Column '%s' in dataset %d has type %s, expected %s",
-                        col,
-                        i,
-                        ds_column_types[col],
-                        ref_column_types[col],
-                    )
+    if len(mismatching_ds) > 0:
+        raise ValueError(
+            "Datasets passed should have same column names. "
+            + "Found {} datasets with mismatching column names".format(
+                len(mismatching_ds)
+            ),
+        )
+
+
+def maybe_align_datasets(datasets):
+    """
+    Given list of datasets
+     1. validate if all datasets have same type and number of columns.
+     2. try casting dataset columns to same value to ensure mergability
+    """
+    try:
+        for d in datasets:
+            if isinstance(d, IterableDataset):
+                d = resolve_iterable_dataset_features(d)
+
+        _validate_mergeable_datasets(datasets)
+        _maybe_cast_columns(datasets)
+    except Exception as e:  # pylint: disable=broad-exception-raised
+        raise ValueError("Failed to align datasets " + str(datasets)) from e
 
 
 def try_convert_bytes_dict_to_pil(image):

Original file line number	Diff line number	Diff line change
`@@ -37,7 +37,7 @@ dependencies = [`
`37`	`37`	`"trl>=0.13,<0.18",`
`38`	`38`	`"peft>=0.8.0,<0.14",`
`39`	`39`	`"protobuf>=5.28.0,<6.0.0",`
`40`		`-"datasets>=2.15.0,<4.0",`
	`40`	`+"datasets>=3.5.0,<4.0",`
`41`	`41`	`"simpleeval>=0.9.13,<2.0",`
`42`	`42`	`"pillow>=11.0.0,<12.0",`
`43`	`43`	`]`