Require datasets>=4.7.0 for Json dtype to prevent insertion of None values

pyproject.toml

@@ -30,7 +30,7 @@ classifiers = [
 requires-python = ">=3.10"
 dependencies = [
     "accelerate>=1.4.0",
-    "datasets>=3.0.0",
+    "datasets>=4.7.0",  # Support Json type and on_mixed_types="use_json"
     "packaging>20.0",
     "transformers>=4.56.2",
 ]

scripts/generate_harmony_dataset.py

@@ -91,7 +91,7 @@ def main(test_size, push_to_hub, repo_id):
             {"reasoning_effort": "high", "model_identity": "You are Tiny ChatGPT, a tiny language model."},
             {"reasoning_effort": "low", "model_identity": "You are Tiny ChatGPT, a tiny language model."},
         ]
-    })
+    }, on_mixed_types="use_json")
     language_modeling_dataset = language_modeling_dataset.train_test_split(test_size=test_size, shuffle=False)
     if push_to_hub:
         language_modeling_dataset.push_to_hub(repo_id, config_name="language_modeling")

scripts/generate_toolcall_dataset.py

@@ -222,7 +222,7 @@ def get_wind_conditions(city: str, unit: str) -> tuple[int, str]:
             json.dumps([start_timer, create_reminder]),
             json.dumps([get_weather_forecast, get_wind_conditions]),
         ]
-    })
+    }, on_mixed_types="use_json")
     language_modeling_dataset = language_modeling_dataset.train_test_split(test_size=test_size, shuffle=False)
     if push_to_hub:
         language_modeling_dataset.push_to_hub(repo_id, config_name="language_modeling")
@@ -328,7 +328,7 @@ def get_wind_conditions(city: str, unit: str) -> tuple[int, str]:
             json.dumps([create_reminder]),
             json.dumps([get_wind_conditions]),
         ],
-    })
+    }, on_mixed_types="use_json")
     preference_dataset = preference_dataset.train_test_split(test_size=test_size, shuffle=False)
     if push_to_hub:
         preference_dataset.push_to_hub(repo_id, config_name="preference")

trl/trainer/dpo_trainer.py

@@ -59,7 +59,6 @@
     get_config_model_id,
     hash_module,
     pad,
-    remove_none_values,
     selective_log_softmax,
     use_adapter,
 )
@@ -863,11 +862,6 @@ def _prepare_dataset(
         args: DPOConfig,
         dataset_name: str,
     ) -> Dataset | IterableDataset:
-        # Tabular backends like Arrow/Parquet insert `None` for mismatched keys in nested structures. Clean them from
-        # sampled data.
-        if isinstance(dataset, Dataset):  # IterableDataset does not support `with_transform`
-            dataset = dataset.with_transform(remove_none_values)
-
         # Build the kwargs for the `map` function
         map_kwargs = {}
         if isinstance(dataset, Dataset):  # IterableDataset does not support num_proc

trl/trainer/reward_trainer.py

@@ -52,7 +52,7 @@
 from ..models import get_act_offloading_ctx_manager
 from .base_trainer import _BaseTrainer
 from .reward_config import RewardConfig
-from .utils import create_model_from_path, disable_dropout_in_model, get_config_model_id, pad, remove_none_values
+from .utils import create_model_from_path, disable_dropout_in_model, get_config_model_id, pad
 
 
 if is_peft_available():
@@ -541,11 +541,6 @@ def _prepare_dataset(
         args: RewardConfig,
         dataset_name: str,
     ) -> Dataset | IterableDataset:
-        # Tabular backends like Arrow/Parquet insert `None` for mismatched keys in nested structures. Clean them from
-        # sampled data.
-        if isinstance(dataset, Dataset):  # IterableDataset does not support `with_transform`
-            dataset = dataset.with_transform(remove_none_values)
-
         # If the dataset is already preprocessed (tokenized), skip the processing steps.
         column_names = get_dataset_column_names(dataset)
         is_processed = "chosen_ids" in column_names and "rejected_ids" in column_names

trl/trainer/sft_trainer.py

@@ -61,7 +61,6 @@
     flush_left,
     get_config_model_id,
     pad,
-    remove_none_values,
     selective_log_softmax,
 )
 
@@ -1054,11 +1053,6 @@ def _prepare_dataset(
         formatting_func: Callable[[dict], str] | None,
         dataset_name: str,
     ) -> Dataset | IterableDataset:
-        # Tabular backends like Arrow/Parquet insert `None` for mismatched keys in nested structures. Clean them from
-        # sampled data.
-        if isinstance(dataset, Dataset):  # IterableDataset does not support `with_transform`
-            dataset = dataset.with_transform(remove_none_values)
-
         # If the dataset is already preprocessed (tokenized), skip the processing steps.
         column_names = get_dataset_column_names(dataset)
         is_processed = "input_ids" in column_names